adding __setitem__

Author: fschlimb

ddptensor/ddptensor.py

@@ -48,8 +48,8 @@ def __repr__(self):
     def __getitem__(self, *args):
         return dtensor(self._t.__getitem__(*args))
 
-    # def __setitem__(self, key, value):
-    #     x = self._t.__setitem__(key, value._t if isinstance(value, dtensor) else value)
+    def __setitem__(self, key, value):
+        x = self._t.__setitem__(key, value._t) # if isinstance(value, dtensor) else value)
 
     # def get_slice(self, *args):
     #     return self._t.get_slice(*args)

src/MPIMediator.cpp

@@ -53,20 +53,20 @@ uint64_t MPIMediator::register_array(tensor_i::ptr_type ary)
     return s_last_id;
 }
 
-void MPIMediator::pull(rank_type from, const tensor_i * ary, const NDSlice & slice, void * rbuff)
+void MPIMediator::pull(rank_type from, const tensor_i & ary, const NDSlice & slice, void * rbuff)
 {
     MPI_Comm comm = MPI_COMM_WORLD;
     MPI_Request request[2];
     MPI_Status status[2];
     Buffer buff;
 
     bitsery::Serializer<OutputAdapter> ser{buff};
-    uint64_t id = ary->id();
+    uint64_t id = ary.id();
     ser.value8b(id);
     ser.object(slice);
     ser.adapter().flush();
 
-    auto sz = slice.size() * ary->item_size();
+    auto sz = slice.size() * ary.item_size();
     std::cerr << "alsdkjf " << sz << " " << buff.size() << " " << rbuff << std::endl;
     MPI_Irecv(rbuff, sz, MPI_CHAR, from, PUSH_TAG, comm, &request[1]);
     MPI_Isend(buff.data(), buff.size(), MPI_CHAR, from, PULL_TAG, comm, &request[0]);

src/ddptensor.cpp

@@ -67,7 +67,7 @@ struct Creator
         return TypeDispatch<x::Creator>(dtype, op, std::forward<shape_type>(shape), std::forward<py::object>(val));
     }
 };
-
+#if 0
 struct IEWBinOp
 {
     static auto op(IEWBinOpId op, x::DPTensorBaseX::ptr_type a, x::DPTensorBaseX::ptr_type b)
@@ -117,6 +117,7 @@ struct ReduceOp
         return TypeDispatch<x::ReduceOp>(a->dtype(), op, a, dim);
     }
 };
+#endif
 
 struct GetItem
 {
@@ -125,6 +126,13 @@ struct GetItem
         return TypeDispatch<x::GetItem>(a->dtype(), a, NDSlice(v));
     }
 };
+struct SetItem
+{
+    static auto op(x::DPTensorBaseX::ptr_type a, const std::vector<py::slice> & v, x::DPTensorBaseX::ptr_type b)
+    {
+        return TypeDispatch<x::SetItem>(a->dtype(), a, NDSlice(v), b);
+    }
+};
 
 rank_type myrank()
 {
@@ -162,7 +170,7 @@ PYBIND11_MODULE(_ddptensor, m) {
     py::class_<Creator>(m, "Creator")
         .def("create_from_shape", &Creator::create_from_shape)
         .def("full", &Creator::full);
-
+#if 0
     py::class_<EWUnyOp>(m, "EWUnyOp")
         .def("op", &EWUnyOp::op);
 
@@ -179,10 +187,12 @@ PYBIND11_MODULE(_ddptensor, m) {
 
     py::class_<ReduceOp>(m, "ReduceOp")
         .def("op", &ReduceOp::op);
+#endif
 
     py::class_<x::DPTensorBaseX, x::DPTensorBaseX::ptr_type>(m, "DPTensorX")
         .def("__repr__", &x::DPTensorBaseX::__repr__)
-        .def("__getitem__", &GetItem::op);
+        .def("__getitem__", &GetItem::op)
+        .def("__setitem__", &SetItem::op);
 
 #if 0
     py::class_<dtensor>(m, "dtensor")

src/include/ddptensor/MPIMediator.hpp

@@ -13,7 +13,7 @@ class MPIMediator : public Mediator
     MPIMediator();
     virtual ~MPIMediator();
     virtual uint64_t register_array(tensor_i::ptr_type ary);
-    virtual void pull(rank_type from, const tensor_i * ary, const NDSlice & slice, void * buffer);
+    virtual void pull(rank_type from, const tensor_i & ary, const NDSlice & slice, void * buffer);
 
 protected:
     void listen();

src/include/ddptensor/Mediator.hpp

@@ -3,7 +3,6 @@
 #pragma once
 
 #include <vector>
-#include "UtilsAndTypes.hpp"
 #include "tensor_i.hpp"
 
 class NDSlice;
@@ -13,7 +12,7 @@ class Mediator
 public:
     virtual ~Mediator() {}
     virtual uint64_t register_array(tensor_i::ptr_type ary) = 0;
-    virtual void pull(rank_type from, const tensor_i * ary, const NDSlice & slice, void * buffer) = 0;
+    virtual void pull(rank_type from, const tensor_i & ary, const NDSlice & slice, void * buffer) = 0;
 };
 
 extern Mediator * theMediator;

src/include/ddptensor/PVSlice.hpp

@@ -172,7 +172,7 @@ class PVSlice
         return _slice.map(slc);
     }
 
-    NDSlice slice_of_rank(rank_type rank) const
+    NDSlice slice_of_rank(rank_type rank = theTransceiver->rank()) const
     {
         if(_base->split_dim() == NOSPLIT) {
             return rank == theTransceiver->rank() ? slice() : NDSlice();
@@ -279,4 +279,11 @@ class PVSlice
     {
         return iterator();
     }
+
+    friend std::ostream &operator<<(std::ostream &output, const PVSlice & slc) {
+        output << "{slice=" << slc.slice()
+               << "base=" << to_string(slc._base->shape())
+               << "offset=" << slc._base->offset()<< "}";
+        return output;
+    }
 };

src/include/ddptensor/tensor_i.hpp

@@ -21,6 +21,16 @@ class tensor_i
 public:
     typedef std::shared_ptr<tensor_i> ptr_type;
 
+    virtual ~tensor_i() {};
+    virtual std::string __repr__() const = 0;
+    virtual DType dtype() const = 0;
+    virtual shape_type shape() const = 0;
+    virtual void bufferize(const NDSlice & slice, Buffer & buff) = 0;
+    virtual int item_size() const = 0;
+    virtual uint64_t id() const = 0;
+};
+ #if 0
+
     virtual ~tensor_i(){}
 
     virtual const PVSlice & pvslice() = 0;
@@ -47,3 +57,4 @@ class tensor_i
 
     virtual py::object get_slice(const NDSlice & slice) const = 0;
 };
+#endif