adding get_local()

Author: fschlimb

src/Deferred.cpp

@@ -127,22 +127,25 @@ void process_promises() {
       function->setAttr(::mlir::LLVM::LLVMDialect::getEmitCWrapperAttrName(),
                         builder.getUnitAttr());
       function.getFunctionType().dump();
-      std::cout << std::endl;
       // add the function to the module
       module.push_back(function);
 
       // get input buffers (before results!)
       auto input = std::move(dm.store_inputs());
 
-      // compile and run the module
-      intptr_t *output = new intptr_t[osz];
-      if (jit.run(module, fname, input, output))
-        throw std::runtime_error("failed running jit");
+      if (osz > 0 || !input.empty()) {
+        // compile and run the module
+        intptr_t *output = new intptr_t[osz];
+        if (jit.run(module, fname, input, output))
+          throw std::runtime_error("failed running jit");
 
-      // push results to deliver promises
-      dm.deliver(output, osz);
+        // push results to deliver promises
+        dm.deliver(output, osz);
 
-      delete[] output;
+        delete[] output;
+      } else {
+        std::cerr << "\tskipping\n";
+      }
     } // no else needed
 
     // now we execute the deferred action which could not be compiled

src/SetGetItem.cpp

@@ -151,6 +151,63 @@ namespace x {
 } // namespace x
 #endif // if 0
 
+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;
+  }
+};
+
+template <typename T> struct wrap_array {
+  template <typename C, typename S>
+  static py::object op(C &&shp, S &&str, void *data, const py::handle &handle) {
+    return py::array(std::forward<C>(shp), std::forward<S>(str),
+                     reinterpret_cast<T *>(data), handle);
+  }
+};
+
+// ***************************************************************************
+
+struct DeferredGetLocal
+    : public DeferredT<GetItem::py_promise_type, GetItem::py_future_type> {
+  id_type _a;
+  py::handle _handle;
+
+  DeferredGetLocal() = default;
+  DeferredGetLocal(const tensor_i::future_type &a, py::handle &handle)
+      : _a(a.id()), _handle(handle) {}
+
+  void run() override {
+    auto aa = std::move(Registry::get(_a).get());
+    auto a_ptr = std::dynamic_pointer_cast<DDPTensorImpl>(aa);
+    assert(a_ptr);
+    auto tmp_shp = a_ptr->local_shape();
+    auto tmp_str = a_ptr->local_strides();
+    auto nd = a_ptr->ndims();
+    auto eSz = sizeof_dtype(a_ptr->dtype());
+    std::vector<ssize_t> strides(nd);
+    for (auto i = 0; i < nd; ++i) {
+      strides[i] = eSz * tmp_str[i];
+    }
+    auto res = dispatch<wrap_array>(a_ptr->dtype(),
+                                    std::vector<ssize_t>(tmp_shp, &tmp_shp[nd]),
+                                    strides, a_ptr->data(), _handle);
+    set_value(res);
+  }
+
+  bool generate_mlir(::mlir::OpBuilder &builder, ::mlir::Location loc,
+                     jit::DepManager &dm) override {
+    return true;
+  }
+
+  FactoryId factory() const { return F_GETLOCAL; }
+
+  template <typename S> void serialize(S &ser) {
+    ser.template value<sizeof(_a)>(_a);
+  }
+};
+
 // ***************************************************************************
 
 struct DeferredGather
@@ -162,14 +219,6 @@ struct DeferredGather
   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
@@ -352,6 +401,10 @@ ddptensor *GetItem::__getitem__(const ddptensor &a,
   return new ddptensor(defer<DeferredGetItem>(a.get(), v));
 }
 
+GetItem::py_future_type GetItem::get_local(const ddptensor &a, py::handle h) {
+  return defer<DeferredGetLocal>(a.get(), h);
+}
+
 GetItem::py_future_type GetItem::gather(const ddptensor &a, rank_type root) {
   return defer<DeferredGather>(a.get(), root);
 }
@@ -372,11 +425,7 @@ py::object GetItem::get_slice(const ddptensor &a,
   return {}; // FIXME TypeDispatch<x::SPMD>(aa.get(), NDSlice(v), aa.id());
 }
 
-py::object GetItem::get_local(const ddptensor &a, py::handle h) {
-  const auto aa = std::move(a.get().get());
-  return {}; // FIXME TypeDispatch<x::SPMD>(aa, h);
-}
-
 FACTORY_INIT(DeferredGetItem, F_GETITEM);
 FACTORY_INIT(DeferredSetItem, F_SETITEM);
 FACTORY_INIT(DeferredGather, F_GATHER);
+FACTORY_INIT(DeferredGather, F_GETLOCAL);

src/ddptensor.cpp

@@ -127,6 +127,7 @@ PYBIND11_MODULE(_ddptensor, m) {
   Factory::init<F_SETITEM>();
   Factory::init<F_SORTOP>();
   Factory::init<F_GATHER>();
+  Factory::init<F_GETITEM>();
 
   jit::init();
 
@@ -140,7 +141,10 @@ PYBIND11_MODULE(_ddptensor, m) {
       .def("sync", &sync_promises)
       .def("myrank", &myrank)
       .def("_get_slice", &GetItem::get_slice)
-      .def("_get_local", &GetItem::get_local)
+      .def("_get_local",
+           [](const ddptensor &f, py::handle h) {
+             PY_SYNC_RETURN(GetItem::get_local(f, h));
+           })
       .def("_gather",
            [](const ddptensor &f, rank_type root = REPLICATED) {
              PY_SYNC_RETURN(GetItem::gather(f, root));

src/include/ddptensor/CppTypes.hpp

@@ -209,6 +209,7 @@ enum FactoryId : int {
   F_SORTOP,
   F_UNYOP,
   F_GATHER,
+  F_GETLOCAL,
   FACTORY_LAST
 };
 

src/include/ddptensor/Deferred.hpp

@@ -91,7 +91,8 @@ class Deferred
 
 extern void _dist(const Runable *p);
 
-// defer operations which do not return a tensor, e.g. which are not a Deferred
+// defer operations which do *not* return a tensor, e.g. which are not a
+// Deferred
 template <typename T, typename... Ts,
           std::enable_if_t<!std::is_base_of_v<Deferred, T>, bool> = true>
 typename T::future_type defer(Ts &&...args) {

src/include/ddptensor/SetGetItem.hpp

@@ -19,7 +19,7 @@ struct GetItem {
                                 const std::vector<py::slice> &v);
   static py::object get_slice(const ddptensor &a,
                               const std::vector<py::slice> &v);
-  static py::object get_local(const ddptensor &a, py::handle h);
+  static py_future_type get_local(const ddptensor &a, py::handle h);
   static py_future_type gather(const ddptensor &a, rank_type root);
 };
 

test/stencil-2d.py

@@ -149,7 +149,8 @@ def main():
 
     for k in range(iterations + 1):
         # start timer after a warmup iteration
-        if k < 1:
+        if k <= 1:
+            np.sync()
             t0 = timer()
 
         if pattern == "star":
@@ -191,8 +192,8 @@ def main():
     # * Analyze and output results.
     # ******************************************************************************
 
-    BB = np.spmd.gather(B)
-    norm = numpy.linalg.norm(numpy.reshape(BB, n * n), ord=1)
+    B = np.spmd.gather(B)
+    norm = numpy.linalg.norm(numpy.reshape(B, n * n), ord=1)
     active_points = (n - 2 * r) ** 2
     norm /= active_points