load_store.hpp

 
#pragma once
#include <cstddef>
#include <cstdint>
#include <iterator>
#include <ranges>
#include <span>
#include "argon/argon_full.hpp"
#include "argon/helpers/argon_for.hpp"
#include "argon/helpers/size.hpp"
#include "arm_simd/helpers/vec128.hpp"
 
#ifdef __ARM_FEATURE_MVE
#define simd mve
#else
#define simd neon
#endif
 
namespace argon::vectorize {

 
template <typename ScalarType>
struct load_store : public std::ranges::view_interface<load_store<ScalarType>> {
  using intrinsic_type = simd::Vec128_t<ScalarType>;

  static constexpr size_t lanes = sizeof(intrinsic_type) / sizeof(ScalarType);
 
 public:
  struct Iterator {
    using iterator_category = std::forward_iterator_tag;     
    using value_type = helpers::ArgonFor_t<intrinsic_type>;  
    using difference_type = std::ptrdiff_t;                  
 
    Iterator() = default;

    Iterator(ScalarType* ptr) : ptr{ptr}, vec{value_type::Load(ptr)} {}
 
    const value_type& operator*() const { return vec; }

    value_type& operator*() {
      dirty_ = true;
      return vec;
    }
 
    const value_type* operator->() const { return &vec; }

    value_type* operator->() {
      dirty_ = true;
      return &vec;
    }

    Iterator& operator++() {
      if (dirty_) {
        vec.StoreTo(ptr);  // store before increment
      }
      ptr += lanes;
      vec = value_type::Load(ptr);
      return *this;
    }

    Iterator operator++(int) {
      Iterator tmp = *this;
      ++(*this);
      return tmp;
    }

    friend bool operator==(const Iterator& a, const Iterator& b) { return a.ptr == b.ptr; }

    friend bool operator==(const Iterator& a, const ScalarType* ptr) { return a.ptr == ptr; }

    friend bool operator!=(const Iterator& a, const Iterator& b) { return a.ptr != b.ptr; }

    friend bool operator!=(const Iterator& a, const ScalarType* ptr) { return a.ptr != ptr; }

    void reload() { vec = value_type::Load(ptr); }
 
   private:
    bool dirty_;
    ScalarType* ptr = nullptr;
    value_type vec;
  };
  static_assert(std::input_or_output_iterator<Iterator>);

  struct ConstIterator {
    using iterator_category = std::forward_iterator_tag;     
    using value_type = helpers::ArgonFor_t<intrinsic_type>;  
    using difference_type = std::ptrdiff_t;                  
 
    ConstIterator() = default;

    ConstIterator(const ScalarType* ptr) : ptr{ptr}, vec{value_type::Load(ptr)} {}

    const value_type operator*() const { return vec; }

    ConstIterator& operator++() {
      ptr += lanes;
      vec = value_type::Load(ptr);
      return *this;
    }

    ConstIterator operator++(int) {
      ConstIterator tmp = *this;
      ++(*this);
      return tmp;
    }

    friend bool operator==(const ConstIterator& a, const ConstIterator& b) { return a.ptr == b.ptr; }

    friend bool operator!=(const ConstIterator& a, const ConstIterator& b) { return a.ptr != b.ptr; }
 
   private:
    const ScalarType* ptr = nullptr;
    value_type vec;
  };
  static_assert(std::input_iterator<ConstIterator>);
 
  using iterator = Iterator;             
  using const_iterator = ConstIterator;  

  load_store(ScalarType* start, ScalarType* end)
      : start_{start}, size_{helpers::vectorizeable_size<ScalarType>(end - start) / lanes} {};

  load_store(ScalarType* start, const size_t size)
      : start_{start}, size_{helpers::vectorizeable_size<ScalarType>(size) / lanes} {};

  load_store(std::span<ScalarType> span)
      : start_{span.data()}, size_{helpers::vectorizeable_size<ScalarType>(span.size()) / lanes} {};

  template <std::ranges::contiguous_range R>
  load_store(R&& r)
      : start_{&*std::ranges::begin(r)}, size_{helpers::vectorizeable_size<ScalarType>(std::ranges::size(r)) / lanes} {}

  iterator begin() const { return Iterator(start_); }

  ScalarType* end() const { return start_ + (size_ * lanes); }

  const_iterator cbegin() const { return ConstIterator(start_); }

  const ScalarType* cend() const { return start_ + (size_ * lanes); }

  size_t size() const { return size_; }
 
 private:
  ScalarType* start_;
  size_t size_;
};

template <std::ranges::contiguous_range R>
load_store(R&& r) -> load_store<std::ranges::range_value_t<R>>;
 
static_assert(std::ranges::range<load_store<int32_t>>);
static_assert(std::ranges::view<load_store<int32_t>>);
static_assert(std::movable<load_store<int32_t>>);
static_assert(std::ranges::viewable_range<load_store<int32_t>>);
 
}  // namespace argon::vectorize
#undef simd