fix buffer overflow in packed iterator in non-padded arrays by padding the end of the array to the next multiple

Author: Fikret Ardal

c++/nda/mem/handle.hpp

@@ -263,7 +263,11 @@ namespace nda::mem {
      */
     handle_heap(long size, do_not_initialize_t) {
       if (size == 0) return;
-      auto b = allocator.allocate(size * sizeof(T), type_alignment_info<T>::required_alignment);
+      long new_size = size;
+      if constexpr(Vectorizable<T>) {
+        new_size = next_multiple(size, native_simd<T>::size());
+      }
+      auto b = allocator.allocate(new_size * sizeof(T), type_alignment_info<T>::required_alignment);
       if (not b.ptr) throw std::bad_alloc{};
       _data = (T *)b.ptr;
       _size = size;
@@ -275,7 +279,11 @@ namespace nda::mem {
      */
     handle_heap(long size, init_zero_t) {
       if (size == 0) return;
-      auto b = allocator.allocate_zero(size * sizeof(T), type_alignment_info<T>::required_alignment);
+      long new_size = size;
+      if constexpr(Vectorizable<T>) {
+        new_size = next_multiple(size, native_simd<T>::size());
+      }
+      auto b = allocator.allocate_zero(new_size * sizeof(T), type_alignment_info<T>::required_alignment);
       if (not b.ptr) throw std::bad_alloc{};
       _data = (T *)b.ptr;
       _size = size;
@@ -295,10 +303,14 @@ namespace nda::mem {
     handle_heap(long size) {
       if (size == 0) return;
       blk_t b;
+      long new_size = size;
+      if constexpr(Vectorizable<T>) {
+        new_size = next_multiple(size, native_simd<T>::size());
+      }
       if constexpr (is_complex_v<T> && init_dcmplx)
-        b = allocator.allocate_zero(size * sizeof(T), type_alignment_info<T>::required_alignment);
+        b = allocator.allocate_zero(new_size * sizeof(T), type_alignment_info<T>::required_alignment);
       else
-        b = allocator.allocate(size * sizeof(T), type_alignment_info<T>::required_alignment);
+        b = allocator.allocate(new_size * sizeof(T), type_alignment_info<T>::required_alignment);
       if (not b.ptr) throw std::bad_alloc{};
       _data = (T *)b.ptr;
       _size = size;

c++/nda/packed.hpp

@@ -1,8 +1,11 @@
 #pragma once
 #include "./concepts.hpp"
+#include "./mem/alignment.hpp"
 
 #include <iterator>
 #include <utility>
+#include <tuple>
+#include <algorithm>
 
 namespace nda {
   template <MemoryArray A>
@@ -24,9 +27,11 @@ namespace nda {
 
       size_t index;
 
-      packed_iterator(const pointer data_ptr, const std::size_t idx) : data(data_ptr), index(idx) {}
+      size_t last_index;
 
-      std::pair<simd_t, pointer> operator*() const {
+      packed_iterator(const pointer data_ptr, const size_t idx, const size_t last_idx) : data(data_ptr), index(idx), last_index(last_idx) {}
+
+      std::tuple<simd_t, pointer, size_t> operator*() const noexcept {
         // Valid field counts= let j be the fastest dimension: let p be the padding. simd_size
         // (index + simd_size) % lengths[j] (index = 0 lengths[j] = 10,  then valid_fields = 8;, index 8 lengtsah[j] = 10
         // ((lengths[j] - (index + simd_size + (lengths[j] / (index))^-1 * padding) % lengths[j]) % lengths[j])
@@ -35,64 +40,68 @@ namespace nda {
         // 8 + 8 - 10
         // 16 + 8 - 10 * 2 - 6
         if constexpr (aligned_and_padded) {
-          return {simd_t(data + index), data + index};
+          return {simd_t(data + index), data + index, std::min(simd_t::size(), last_index - index)};
         } else {
           simd_t tmp;
           tmp.load_unaligned(data + index);
-          return {tmp, data + index};
+          return {tmp, data + index, std::min(simd_t::size(), last_index - index)};
         }
       }
 
-      pointer operator->() const { return data + index; }
+      pointer operator->() const noexcept { return data + index; }
 
-      packed_iterator &operator++() {
+      packed_iterator &operator++() noexcept {
         index += simd_t::size();
         return *this;
       }
 
-      packed_iterator operator++(int) {
+      packed_iterator operator++(int) noexcept {
         packed_iterator tmp = *this;
         index += simd_t::size();
         return tmp;
       }
 
-      packed_iterator &operator--() {
+      packed_iterator &operator--() noexcept {
         index -= simd_t::size();
         return *this;
       }
 
-      packed_iterator operator--(int) {
+      packed_iterator operator--(int) noexcept {
         packed_iterator tmp = *this;
         index -= simd_t::size();
         return tmp;
       }
 
-      packed_iterator &operator+=(std::ptrdiff_t n) {
+      packed_iterator &operator+=(std::ptrdiff_t n) noexcept {
         index += n * simd_t::size();
         return *this;
       }
 
-      packed_iterator &operator-=(std::ptrdiff_t n) {
+      packed_iterator &operator-=(std::ptrdiff_t n) noexcept {
         index -= n * simd_t::size();
         return *this;
       }
 
-      packed_iterator operator+(std::ptrdiff_t n) const { return packed_iterator(data, index + n * simd_t::size()); }
-      packed_iterator operator-(std::ptrdiff_t n) const { return packed_iterator(data, index - n * simd_t::size()); }
-      std::ptrdiff_t operator-(const packed_iterator &other) const { return (index - other.index) / simd_t::size(); }
-      std::pair<simd_t, pointer> operator[](std::ptrdiff_t n) const {
-        return {simd_t(data + index + n * simd_t::size()), data + index + n * simd_t::size()};
+      packed_iterator operator+(std::ptrdiff_t n) const noexcept { return packed_iterator(data, index + n * simd_t::size()); }
+      packed_iterator operator-(std::ptrdiff_t n) const noexcept { return packed_iterator(data, index - n * simd_t::size()); }
+      std::ptrdiff_t operator-(const packed_iterator &other) const noexcept { return (index - other.index) / simd_t::size(); }
+      std::tuple<simd_t, pointer, size_t> operator[](std::ptrdiff_t n) const noexcept {
+        return {simd_t(data + index + n * simd_t::size()), data + index + n * simd_t::size(),
+                std::min(simd_t::size(), last_index - index + n * simd_t::size())};
       }
 
-      bool operator==(const packed_iterator &other) const { return index == other.index; }
-      bool operator!=(const packed_iterator &other) const { return !(*this == other); }
-      bool operator<(const packed_iterator &other) const { return index < other.index; }
-      bool operator>(const packed_iterator &other) const { return index > other.index; }
-      bool operator<=(const packed_iterator &other) const { return index <= other.index; }
-      bool operator>=(const packed_iterator &other) const { return index >= other.index; }
+      bool operator==(const packed_iterator &other) const noexcept { return index == other.index; }
+      bool operator!=(const packed_iterator &other) const noexcept { return !(*this == other); }
+      bool operator<(const packed_iterator &other) const noexcept { return index < other.index; }
+      bool operator>(const packed_iterator &other) const noexcept { return index > other.index; }
+      bool operator<=(const packed_iterator &other) const noexcept { return index <= other.index; }
+      bool operator>=(const packed_iterator &other) const noexcept { return index >= other.index; }
     };
 
-    packed_iterator begin() { return packed_iterator(array.data(), 0); }
-    packed_iterator end() { return packed_iterator(array.data(), array.indexmap().capacity()); }
+    packed_iterator begin() const noexcept { return packed_iterator(array.data(), 0, array.indexmap().capacity()); }
+    packed_iterator end() const noexcept {
+      return packed_iterator(array.data(), mem::next_multiple(array.indexmap().capacity(), native_simd<std::remove_cvref_t<get_value_t<A>>>::size()),
+                             array.indexmap().capacity());
+    }
   };
 } // namespace nda

c++/nda/simd/arch/simd.hpp

@@ -11,8 +11,35 @@
 #include "./AVX512/functions.hpp"
 #endif
 
+#include <algorithm>
+
 namespace nda {
   template <typename T>
   using native_simd = simd_type<std::conditional_t<std::is_same_v<T, long long>, int64_t, T>, abi::get_native_width<T>(), abi::get_native_abi_tag()>;
 
+  template <typename T, bool Aligned>
+  struct simd_block {
+    using simd_t = native_simd<T>;
+
+    simd_t value;
+    T *data;
+    size_t valid_field;
+
+    simd_block(const simd_t &v, const T *d, const size_t valid) : value(v), data(d), valid_field(valid) {}
+
+    void store() {
+      if (valid_field == simd_t::size()) [[likely]] {
+        if constexpr (Aligned) {
+          value.store(data);
+        } else {
+          value.store_unaligned(data);
+        }
+      } else {
+        alignas(simd_t::alignment()) std::array<T, simd_t::size()> tmp;
+        value.store(tmp);
+        std::copy_n(data, valid_field, tmp);
+      }
+    }
+  };
+
 } // namespace nda

test/c++/nda_simd.cpp

@@ -1569,29 +1569,33 @@ TEST(NDA, OurSIMD) {
   const long size1 = 2;
   const long size2 = 10;
   float k          = 1;
-  array_aligned<float,2> s({size1, size2});
-  array_aligned<float,2> x({size1, size2});
+  array<float, 2> s({size1, size2});
+  array<float, 2> x({size1, size2});
   for (int i = 0; i < size1; ++i) {
     for (int j = 0; j < size2; ++j) { s(i, j) = k++; }
   }
   for (int i = 0; i < size1; ++i) {
     for (int j = 0; j < size2; ++j) { x(i, j) = k++; }
   }
+  for (auto [simd, pointer, valid_field] : packed(s)) {
+    using simd_t = native_simd<float>;
+    alignas(simd_t::alignment()) std::array<float, simd_t::size()> tmp;
+    simd.store(tmp.data());
+    std::cout << tmp << std::endl;
+    std::cout << valid_field << std::endl;
+  }
   // s = array_aligned<float,2>::rand({size1, size2});
   auto hop = make_array_view(s);
   std::cout << hop << std::endl;
   std::cout << hop.is_aligned << std::endl;
   std::cout << typeid(hop).name() << std::endl;
   std::cout << typeid(make_regular(hop)).name() << std::endl;
   std::cout << hop.indexmap().capacity() << std::endl;
-  for (int i= 0; i < hop.indexmap().capacity() ; ++i) {
-    std::cout << *(hop.data() + i) << " ";
-  }
+  for (int i = 0; i < hop.indexmap().capacity(); ++i) { std::cout << *(hop.data() + i) << " "; }
   std::cout << std::endl;
   auto aa = make_regular(hop);
   std::cout << aa.size() << std::endl;
   std::cout << aa.is_aligned << std::endl;
   bool a = s == hop;
   std::cout << (hop == s) << std::endl;
-
 }