expose NDCuVec.strides()

AMYPAD · Mar 13, 2024 · b6ed960 · b6ed960
1 parent e631d2a
commit b6ed960
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/cuvec/include/cuvec.cuh b/cuvec/include/cuvec.cuh
@@ -112,6 +112,13 @@ template <class T> struct NDCuVec {
     if (size != vec.size()) throw std::length_error("reshape: size mismatch");
     this->shape = shape;
   }
+  std::vector<size_t> strides() const {
+    const size_t ndim = this->shape.size();
+    std::vector<size_t> s(ndim);
+    s[ndim - 1] = sizeof(T);
+    for (int i = ndim - 2; i >= 0; i--) s[i] = this->shape[i + 1] * s[i + 1];
+    return s;
+  }
 };
 
 #endif // _CUVEC_H_
diff --git a/cuvec/include/cuvec_pybind11.cuh b/cuvec/include/cuvec_pybind11.cuh
@@ -32,13 +32,9 @@ PYBIND11_MAKE_OPAQUE(NDCuVec<double>);
   pybind11::class_<NDCuVec<T>>(m, PYBIND11_TOSTRING(NDCuVec_##typechar),                          \
                                pybind11::buffer_protocol())                                       \
       .def_buffer([](NDCuVec<T> &v) -> pybind11::buffer_info {                                    \
-        size_t ndim = v.shape.size();                                                             \
-        std::vector<size_t> strides(ndim);                                                        \
-        strides[ndim - 1] = sizeof(T);                                                            \
-        for (int i = ndim - 2; i >= 0; i--) strides[i] = v.shape[i + 1] * strides[i + 1];         \
         return pybind11::buffer_info(v.vec.data(), sizeof(T),                                     \
                                      pybind11::format_descriptor<T>::format(), v.shape.size(),    \
-                                     v.shape, strides);                                           \
+                                     v.shape, v.strides());                                       \
       })                                                                                          \
       .def(pybind11::init<>())                                                                    \
       .def(pybind11::init<std::vector<size_t>>())                                                 \