adding reshape (partial) and arange

Author: fschlimb

ddptensor/__init__.py

@@ -56,9 +56,20 @@
         exec(
             f"{func} = lambda shape, val, dtype: dtensor(_cdt.Creator.full(_cdt.{FUNC}, shape, val, dtype))"
         )
+    elif func == "arange":
+        exec(
+            f"{func} = lambda start, end, step, dtype: dtensor(_cdt.Creator.arange(start, end, step, dtype))"
+        )
 
 for func in api.api_categories["ReduceOp"]:
     FUNC = func.upper()
     exec(
         f"{func} = lambda this, dim: dtensor(_cdt.ReduceOp.op(_cdt.{FUNC}, this._t, dim))"
     )
+
+for func in api.api_categories["ManipOp"]:
+    FUNC = func.upper()
+    if func == "reshape":
+        exec(
+            f"{func} = lambda this, shape: dtensor(_cdt.ManipOp.reshape(this._t, shape))"
+        )

setup.py

@@ -29,7 +29,7 @@ def build_cmake(self, ext):
         extdir.parent.mkdir(parents=True, exist_ok=True)
 
         # example of cmake args
-        config = 'Debug' if self.debug else 'Release'
+        config = 'Debug'# if self.debug else 'Release'
         cmake_args = [
             '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=' + str(extdir.parent.absolute()),
             '-DCMAKE_BUILD_TYPE=' + config

src/Creator.cpp

@@ -38,6 +38,15 @@ namespace x {
             }
             throw std::runtime_error("Unknown creator");
         }
+
+        static ptr_type op(uint64_t start, uint64_t end, uint64_t step)
+        {
+            PVSlice pvslice({Slice(start, end, step).size()});
+            auto lslc = pvslice.slice_of_rank();
+            const auto & l1dslc = lslc.dim(0);
+            auto a = xt::arange<T>(start + l1dslc._start*step, start + l1dslc._end * step, l1dslc._step);
+            return operatorx<T>::mk_tx(std::move(pvslice), std::move(a));
+        }
     }; // class creatorx
 } // namespace x
 
@@ -51,3 +60,8 @@ tensor_i::ptr_type Creator::full(const shape_type & shape, const py::object & va
     auto op = FULL;
     return TypeDispatch<x::Creator>(dtype, op, shape, val);
 }
+
+tensor_i::ptr_type Creator::arange(uint64_t start, uint64_t end, uint64_t step, DTypeId dtype)
+{
+    return TypeDispatch<x::Creator>(dtype, start, end, step);
+}

src/MPITransceiver.cpp

@@ -72,3 +72,23 @@ void MPITransceiver::reduce_all(void * inout, DTypeId T, size_t N, RedOpType op)
 {
     MPI_Allreduce(MPI_IN_PLACE, inout, N, to_mpi(T), to_mpi(op), MPI_COMM_WORLD);
 }
+
+void MPITransceiver::alltoall(const void* buffer_send,
+                              const int* counts_send,
+                              const int* displacements_send,
+                              DTypeId datatype_send,
+                              void* buffer_recv,
+                              const int* counts_recv,
+                              const int* displacements_recv,
+                              DTypeId datatype_recv)
+{
+    MPI_Alltoallv(buffer_send,
+                  counts_send,
+                  displacements_send,
+                  to_mpi(datatype_send),
+                  buffer_recv,
+                  counts_recv,
+                  displacements_recv,
+                  to_mpi(datatype_recv),
+                  MPI_COMM_WORLD);
+}

src/ddptensor.cpp

@@ -60,7 +60,8 @@ PYBIND11_MODULE(_ddptensor, m) {
 
     py::class_<Creator>(m, "Creator")
         .def("create_from_shape", &Creator::create_from_shape)
-        .def("full", &Creator::full);
+        .def("full", &Creator::full)
+        .def("arange", &Creator::arange);
 
     py::class_<EWUnyOp>(m, "EWUnyOp")
         .def("op", &EWUnyOp::op);
@@ -74,6 +75,9 @@ PYBIND11_MODULE(_ddptensor, m) {
     py::class_<ReduceOp>(m, "ReduceOp")
         .def("op", &ReduceOp::op);
 
+    py::class_<ManipOp>(m, "ManipOp")
+        .def("reshape", &ManipOp::reshape);
+
     py::class_<tensor_i, tensor_i::ptr_type>(m, "DPTensorX")
         .def_property_readonly("dtype", &tensor_i::dtype)
         .def_property_readonly("shape", &tensor_i::shape)

src/include/ddptensor/MPITransceiver.hpp

@@ -23,6 +23,14 @@ class MPITransceiver : public Transceiver
     virtual void barrier();
     virtual void bcast(void * ptr, size_t N, rank_type root);
     virtual void reduce_all(void * inout, DTypeId T, size_t N, RedOpType op);
+    virtual void alltoall(const void* buffer_send,
+                          const int* counts_send,
+                          const int* displacements_send,
+                          DTypeId datatype_send,
+                          void* buffer_recv,
+                          const int* counts_recv,
+                          const int* displacements_recv,
+                          DTypeId datatype_recv);
 
 private:
     rank_type _nranks, _rank;

src/include/ddptensor/NDSlice.hpp

@@ -201,6 +201,26 @@ class NDSlice {
             } );
     }
 
+    ///
+    /// @return Copy of NDSlice which was trimmed by given slice
+    ///
+    NDSlice trim(const NDSlice & slc) const
+    {
+        return _convert([&](uint64_t i) {
+                return _slice_vec[i].trim(slc.dim(i)._start, slc.dim(i)._end);
+            } );
+    }
+
+    ///
+    /// @return Copy of NDSlice which was trimmed by given slice
+    ///
+    NDSlice overlap(const NDSlice & slc) const
+    {
+        return _convert([&](uint64_t i) {
+                return _slice_vec[i].overlap(slc.dim(i));
+            } );
+    }
+    
     ///
     /// @return Copy of NDSlice which was trimmed by t_slc and shifted by s_slc._start's
     ///

src/include/ddptensor/Operations.hpp

@@ -11,6 +11,7 @@ struct Creator
 {
     static tensor_i::ptr_type create_from_shape(CreatorId op, const shape_type & shape, DTypeId dtype=FLOAT64);
     static tensor_i::ptr_type full(const shape_type & shape, const py::object & val, DTypeId dtype=FLOAT64);
+    static tensor_i::ptr_type arange(uint64_t start, uint64_t end, uint64_t step, DTypeId dtype=INT64);
 };
 
 struct IEWBinOp
@@ -44,6 +45,11 @@ struct SetItem
     static void __setitem__(x::DPTensorBaseX::ptr_type a, const std::vector<py::slice> & v, x::DPTensorBaseX::ptr_type b);
 };
 
+struct ManipOp
+{
+    static tensor_i::ptr_type reshape(x::DPTensorBaseX::ptr_type a, const shape_type & shape);
+};
+
 
 // Dependent on dt, dispatch arguments to a operation class.
 // The operation must

src/include/ddptensor/PVSlice.hpp

@@ -12,6 +12,7 @@ constexpr static int NOSPLIT = -1;
 class BasePVSlice
 {
     uint64_t   _offset;
+    uint64_t   _tile_size;
     shape_type _shape;
     int        _split_dim;
 
@@ -24,15 +25,18 @@ class BasePVSlice
           _shape(shape),
           _split_dim(split)
     {
+        _tile_size = VPROD(_shape) / shape[_split_dim] * _offset;
     }
     BasePVSlice(shape_type && shape, int split=0)
         : _offset(split == NOSPLIT ? 0 : (shape[split] + theTransceiver->nranks() - 1) / theTransceiver->nranks()),
           _shape(std::move(shape)),
           _split_dim(split)
     {
+        _tile_size = VPROD(_shape) / _shape[_split_dim] * _offset;
     }
 
     uint64_t offset() const { return _offset; }
+    uint64_t tile_size() const { return _tile_size; }
     int split_dim() const { return _split_dim; }
     const shape_type & shape() const { return _shape; }
     shape_type shape(rank_type rank) const
@@ -125,6 +129,11 @@ class PVSlice
         return _base->split_dim();
     }
 
+    const uint64_t tile_size() const
+    {
+        return _base->tile_size();
+    }
+
     const shape_type & shape() const
     {
         if(_shape.size() != _slice.ndims()) {