adding ddpt.gather/to_numpy

Author: fschlimb

src/CollComm.cpp

@@ -2,6 +2,80 @@
 
 #include "ddptensor/CollComm.hpp"
 
+void bufferize(DDPTensorImpl::ptr_type a_ptr, void *outPtr) {
+  dispatch(a_ptr->dtype(), a_ptr->data(), [&a_ptr, outPtr](auto *ptr) {
+    auto buff = static_cast<decltype(ptr)>(outPtr);
+
+    forall(0, ptr, a_ptr->local_shape(), a_ptr->local_strides(), a_ptr->ndims(),
+           [&buff](const auto *in) {
+             *buff = *in;
+             ++buff;
+           });
+  });
+}
+
+// @param outPtr touched only if on root and/or root==REPLICATED
+void gather_tensor(DDPTensorImpl::ptr_type a_ptr, rank_type root,
+                   void *outPtr) {
+  auto trscvr = getTransceiver();
+  auto nranks = trscvr->nranks();
+  auto myrank = trscvr->rank();
+  bool sendonly = root != REPLICATED && root != myrank;
+
+  auto dtype = a_ptr->dtype();
+  auto mysizes = a_ptr->local_shape();
+  auto mysz = mysizes[0];
+  auto myoff = a_ptr->local_offsets()[0];
+  auto nd = a_ptr->ndims();
+  auto gshape = a_ptr->shape();
+  auto myTileSz =
+      std::accumulate(&gshape[1], &gshape[nd], 1, std::multiplies<int64_t>());
+
+  // allgather process local offset and sizes
+  std::vector<int> displacements(nranks);
+  std::vector<int> counts(nranks, 2);
+  std::vector<int> szsAndOffs(2 * nranks);
+  for (size_t i = 0; i < nranks; ++i) {
+    displacements[i] = i * 2;
+  }
+  szsAndOffs[2 * myrank + 0] = myoff; // FIXME split dim
+  szsAndOffs[2 * myrank + 1] = mysz;
+  trscvr->gather(szsAndOffs.data(), counts.data(), displacements.data(), INT32,
+                 REPLICATED);
+
+  // compute each pranks local contribution
+  int64_t curr = 0;
+  for (auto i = 0; i < nranks; ++i) {
+    assert(szsAndOffs[i * 2] * myTileSz == curr);
+    displacements[i] = curr;
+    counts[i] = szsAndOffs[i * 2 + 1] * myTileSz;
+    curr += counts[i];
+  }
+
+  // create buffer/numpy array and copy
+  void *ptr = nullptr;
+  bool need_del = false;
+  if (sendonly) {
+    if (mysz > 0 && a_ptr->is_sliced()) {
+      ptr = new char[mysz * sizeof_dtype(dtype) * myTileSz];
+      bufferize(a_ptr, ptr);
+      need_del = true;
+    } else if (mysz > 0) {
+      ptr = a_ptr->data();
+    }
+  } else {
+    ptr = outPtr;
+    bufferize(a_ptr, &(static_cast<char *>(
+                         ptr)[displacements[myrank] * sizeof_dtype(dtype)]));
+  }
+
+  // final gather
+  trscvr->gather(ptr, counts.data(), displacements.data(), dtype, root);
+
+  if (need_del)
+    delete[] static_cast<char *>(ptr);
+}
+
 // Compute offset and displacements when mapping n_slc to o_slc. This is
 // necessary when slices are not equally partitioned.
 //

src/IO.cpp

@@ -10,33 +10,7 @@
 #include "ddptensor/Transceiver.hpp"
 #include "ddptensor/TypeDispatch.hpp"
 
-using promise_type = std::promise<py::object>;
-using future_type = std::shared_future<py::object>;
-
-struct DeferredToNumpy : public DeferredT<promise_type, future_type> {
-  id_type _a;
-
-  DeferredToNumpy() = default;
-  DeferredToNumpy(const tensor_i::future_type &a) : _a(a.id()) {}
-
-  void run() {
-    const auto a = std::move(Registry::get(_a).get());
-    set_value(
-        GetItem::do_gather(a, getTransceiver()->is_cw() ? 0 : REPLICATED));
-  }
-
-  FactoryId factory() const { return F_TONUMPY; }
-
-  template <typename S> void serialize(S &ser) {
-    ser.template value<sizeof(_a)>(_a);
-  }
-};
-
-py::object IO::to_numpy(const ddptensor &a) {
+GetItem::py_future_type IO::to_numpy(const ddptensor &a) {
   assert(!getTransceiver()->is_cw() || getTransceiver()->rank() == 0);
-  auto f = defer<DeferredToNumpy>(a.get());
-  auto x = f.get();
-  return x;
+  return GetItem::gather(a, 0);
 }
-
-FACTORY_INIT(DeferredToNumpy, F_TONUMPY);

src/SetGetItem.cpp

@@ -6,11 +6,13 @@
 */
 
 #include "ddptensor/SetGetItem.hpp"
+#include "ddptensor/CollComm.hpp"
 #include "ddptensor/Creator.hpp"
 #include "ddptensor/DDPTensorImpl.hpp"
 #include "ddptensor/Factory.hpp"
 #include "ddptensor/Mediator.hpp"
 #include "ddptensor/NDSlice.hpp"
+#include "ddptensor/Transceiver.hpp"
 #include "ddptensor/TypeDispatch.hpp"
 #include "ddptensor/UtilsAndTypes.hpp"
 
@@ -19,6 +21,10 @@
 #include <imex/Utils/PassUtils.h>
 #include <mlir/IR/Builders.h>
 
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+namespace py = pybind11;
+
 #if 0
 namespace x {
 
@@ -140,59 +146,75 @@ namespace x {
             T * data = a_ptr->xarray().data();
             return py::array(std::move(slc.shape()), std::move(strides), data + off, handle);
         }
-
-        // gather
-        // We simply create a local buffer, copy our local data to the right place
-        // and then call AllGatherV via inplace operation.
-        template<typename T>
-        static py::object op(rank_type root, const std::shared_ptr<DPTensorX<T>> & a_ptr)
-        {
-            auto nranks = getTransceiver()->nranks();
-            auto rank = getTransceiver()->rank();
-            bool sendonly = root != REPLICATED && root != rank;
-            const auto & slc = a_ptr->slice();
-            auto mysz = slc.local_slice().size();
-
-            // create buffer/numpy array
-            T * ptr = nullptr;
-            py::array res;
-            if(sendonly) {
-                if(mysz > 0 && a_ptr->is_sliced()) ptr = new T[mysz];
-            } else {
-                res = py::array_t<T>(slc.shape());
-                ptr = reinterpret_cast<T*>(res.mutable_data());
-            }
-            int displacements[nranks];
-            int counts[nranks];
-            int off = 0;
-            // for each rank compute counts and displacements
-            for(auto i=0; i<nranks; ++i) {
-                uint64_t szi = i == rank ? mysz : slc.local_slice(i).size();
-                counts[i] = szi;
-                displacements[i] = off;
-                // copy our local data
-                if(i == rank) {
-                    if(a_ptr->is_sliced()) {
-                        // if non-contiguous copy element by element
-                        const auto & av = xt::strided_view(a_ptr->xarray(), a_ptr->lslice());
-                        uint64_t j = sendonly ? -1 : off - 1;
-                        for(auto v : av) ptr[++j] = v;
-                    } else {
-                        if(sendonly && mysz > 0) ptr = a_ptr->xarray().data();
-                        else memcpy(&ptr[off], a_ptr->xarray().data(), szi*sizeof(T));
-                    }
-                }
-                off += szi;
-            }
-            getTransceiver()->gather(ptr, counts, displacements, DTYPE<T>::value, root);
-            if(sendonly && mysz > 0 && a_ptr->is_sliced()) delete [] ptr;
-            return res;
-        }
     };
 
 } // namespace x
 #endif // if 0
 
+// ***************************************************************************
+
+struct DeferredGather
+    : public DeferredT<GetItem::py_promise_type, GetItem::py_future_type> {
+  id_type _a;
+  rank_type _root;
+
+  DeferredGather() = default;
+  DeferredGather(const tensor_i::future_type &a, rank_type root)
+      : _a(a.id()), _root(root) {}
+
+  template <typename T> struct mk_array {
+    template <typename C> static py::object op(C &&shp, void *&outPtr) {
+      auto ary = py::array_t<T>(std::forward<C>(shp));
+      outPtr = ary.mutable_data();
+      return ary;
+    }
+  };
+
+  void run() override {
+    // gather
+    // We simply create a local buffer, copy our local data to the right place
+    // and then call AllGatherV via inplace operation.
+    auto trscvr = getTransceiver();
+    auto myrank = trscvr->rank();
+    auto aa = std::move(Registry::get(_a).get());
+    auto a_ptr = std::dynamic_pointer_cast<DDPTensorImpl>(aa);
+    assert(a_ptr);
+    bool sendonly = _root != REPLICATED && _root != myrank;
+
+    void *outPtr = nullptr;
+    py::object res;
+    if (!sendonly) {
+      auto tmp = a_ptr->shape();
+      // std::vector<ssize_t> shp(tmp, &tmp[a_ptr->ndims()]);
+      res = dispatch<mk_array>(a_ptr->dtype(),
+                               std::vector<ssize_t>(tmp, &tmp[a_ptr->ndims()]),
+                               outPtr);
+      // (void*)nullptr, [&shp, &res, &outPtr](auto * ptr) {
+      //     auto ary = py::array_t<double>({4,4});
+      //     res = ary;
+      //     outPtr = ary.mutable_data();
+      // });
+    }
+
+    gather_tensor(a_ptr, _root, outPtr);
+
+    set_value(res);
+  }
+
+  bool generate_mlir(::mlir::OpBuilder &builder, ::mlir::Location loc,
+                     jit::DepManager &dm) override {
+    return true;
+  }
+
+  FactoryId factory() const { return F_GATHER; }
+
+  template <typename S> void serialize(S &ser) {
+    ser.template value<sizeof(_a)>(_a);
+  }
+};
+
+// ***************************************************************************
+
 struct DeferredSetItem : public Deferred {
   id_type _a;
   id_type _b;
@@ -252,15 +274,7 @@ struct DeferredSetItem : public Deferred {
   }
 };
 
-ddptensor *SetItem::__setitem__(ddptensor &a, const std::vector<py::slice> &v,
-                                const py::object &b) {
-
-  auto bb = Creator::mk_future(b);
-  auto res = new ddptensor(defer<DeferredSetItem>(a.get(), bb.first->get(), v));
-  if (bb.second)
-    delete bb.first;
-  return res;
-}
+// ***************************************************************************
 
 struct DeferredGetItem : public Deferred {
   id_type _a;
@@ -331,11 +345,27 @@ struct DeferredGetItem : public Deferred {
   }
 };
 
+// ***************************************************************************
+
 ddptensor *GetItem::__getitem__(const ddptensor &a,
                                 const std::vector<py::slice> &v) {
   return new ddptensor(defer<DeferredGetItem>(a.get(), v));
 }
 
+GetItem::py_future_type GetItem::gather(const ddptensor &a, rank_type root) {
+  return defer<DeferredGather>(a.get(), root);
+}
+
+ddptensor *SetItem::__setitem__(ddptensor &a, const std::vector<py::slice> &v,
+                                const py::object &b) {
+
+  auto bb = Creator::mk_future(b);
+  auto res = new ddptensor(defer<DeferredSetItem>(a.get(), bb.first->get(), v));
+  if (bb.second)
+    delete bb.first;
+  return res;
+}
+
 py::object GetItem::get_slice(const ddptensor &a,
                               const std::vector<py::slice> &v) {
   const auto aa = std::move(a.get());
@@ -347,14 +377,6 @@ py::object GetItem::get_local(const ddptensor &a, py::handle h) {
   return {}; // FIXME TypeDispatch<x::SPMD>(aa, h);
 }
 
-py::object GetItem::do_gather(const tensor_i::ptr_type &a, rank_type root) {
-  return {}; // FIXME TypeDispatch<x::SPMD>(a, root);
-}
-
-py::object GetItem::gather(const ddptensor &a, rank_type root) {
-  const auto aa = std::move(a.get().get());
-  return do_gather(aa, root);
-}
-
 FACTORY_INIT(DeferredGetItem, F_GETITEM);
 FACTORY_INIT(DeferredSetItem, F_SETITEM);
+FACTORY_INIT(DeferredGather, F_GATHER);