new future type and ReadyCallback for no-tensor values, fixing syncing with insert_slice ops

fschlimb · fschlimb · commit fe1976adad5a · 2023-03-16T07:11:31.000-05:00
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -98,7 +98,6 @@ set(DDPTSrcs
     ${PROJECT_SOURCE_DIR}/src/ManipOp.cpp
     ${PROJECT_SOURCE_DIR}/src/Random.cpp
     ${PROJECT_SOURCE_DIR}/src/ReduceOp.cpp
-    ${PROJECT_SOURCE_DIR}/src/Service.cpp
     ${PROJECT_SOURCE_DIR}/src/SetGetItem.cpp
     ${PROJECT_SOURCE_DIR}/src/Sorting.cpp
 )
@@ -110,6 +109,7 @@ set(RTSrcs
     ${PROJECT_SOURCE_DIR}/src/Mediator.cpp
     ${PROJECT_SOURCE_DIR}/src/MPIMediator.cpp
     ${PROJECT_SOURCE_DIR}/src/Registry.cpp
+    ${PROJECT_SOURCE_DIR}/src/Service.cpp
     ${PROJECT_SOURCE_DIR}/src/jit/mlir.cpp
 )
 set(IDTRSrcs
diff --git a/src/DDPTensorImpl.cpp b/src/DDPTensorImpl.cpp
@@ -8,6 +8,7 @@
 #include <ddptensor/Transceiver.hpp>
 
 #include <algorithm>
+#include <iostream>
 
 DDPTensorImpl::DDPTensorImpl(Transceiver *transceiver, DTypeId dtype,
                              uint64_t ndims, void *allocated, void *aligned,
@@ -161,6 +162,8 @@ void DDPTensorImpl::add_to_args(std::vector<void *> &args, int ndims) {
   buff[2] = static_cast<intptr_t>(_offset);
   memcpy(buff + 3, _sizes, ndims * sizeof(intptr_t));
   memcpy(buff + 3 + ndims, _strides, ndims * sizeof(intptr_t));
+  for (auto i = 0; i < 3 + 2 * ndims; ++i)
+    std::cerr << " " << buff[i];
   args.push_back(buff);
   // second the transceiver
   args.push_back(&_transceiver);
diff --git a/src/Deferred.cpp b/src/Deferred.cpp
@@ -10,6 +10,7 @@
 #include "include/ddptensor/Deferred.hpp"
 #include "include/ddptensor/Mediator.hpp"
 #include "include/ddptensor/Registry.hpp"
+#include "include/ddptensor/Service.hpp"
 #include "include/ddptensor/Transceiver.hpp"
 
 #include <imex/Dialect/Dist/IR/DistOps.h>
@@ -50,7 +51,9 @@ Deferred::future_type defer_tensor(Runable::ptr_type &&_d, bool is_global) {
     throw std::runtime_error("Expected Deferred Tensor promise");
   if (is_global) {
     _dist(d);
-    d->set_guid(Registry::get_guid());
+    if (d->guid() == Registry::NOGUID) {
+      d->set_guid(Registry::get_guid());
+    }
   }
   auto f = d->get_future();
   Registry::put(f);
@@ -154,9 +157,4 @@ void process_promises() {
   } while (!done);
 }
 
-void sync_promises() {
-  // FIXME this does not wait for the last deferred to complete
-  while (!_deferred.empty()) {
-    std::this_thread::sleep_for(std::chrono::milliseconds(1));
-  }
-}
+void sync_promises() { (void)Service::run().get(); }
diff --git a/src/EWBinOp.cpp b/src/EWBinOp.cpp
@@ -459,20 +459,19 @@ struct DeferredEWBinOp : public Deferred {
     auto outTyp =
         ::imex::ptensor::PTensorType::get(shape, aTyp.getElementType());
 
-    dm.addVal(
-        this->guid(),
-        builder.create<::imex::ptensor::EWBinOp>(
-            loc, outTyp, builder.getI32IntegerAttr(ddpt2mlir(_op)), av, bv),
-        [this](Transceiver *transceiver, uint64_t rank, void *allocated,
-               void *aligned, intptr_t offset, const intptr_t *sizes,
-               const intptr_t *strides, uint64_t *gs_allocated,
-               uint64_t *gs_aligned, uint64_t *lo_allocated,
-               uint64_t *lo_aligned, uint64_t balanced) {
-          this->set_value(std::move(
-              mk_tnsr(transceiver, _dtype, rank, allocated, aligned, offset,
-                      sizes, strides, gs_allocated, gs_aligned, lo_allocated,
-                      lo_aligned, balanced)));
-        });
+    auto bop = builder.create<::imex::ptensor::EWBinOp>(
+        loc, outTyp, builder.getI32IntegerAttr(ddpt2mlir(_op)), av, bv);
+    dm.addVal(this->guid(), bop,
+              [this](Transceiver *transceiver, uint64_t rank, void *allocated,
+                     void *aligned, intptr_t offset, const intptr_t *sizes,
+                     const intptr_t *strides, uint64_t *gs_allocated,
+                     uint64_t *gs_aligned, uint64_t *lo_allocated,
+                     uint64_t *lo_aligned, uint64_t balanced) {
+                this->set_value(std::move(
+                    mk_tnsr(transceiver, _dtype, rank, allocated, aligned,
+                            offset, sizes, strides, gs_allocated, gs_aligned,
+                            lo_allocated, lo_aligned, balanced)));
+              });
     return false;
   }
 
diff --git a/src/Registry.cpp b/src/Registry.cpp
@@ -22,7 +22,12 @@ id_type get_guid() { return ++_nguid; }
 
 void put(const tensor_i::future_type &ptr) {
   locker _l(_mutex);
-  _keeper[ptr.id()] = ptr;
+  _keeper.insert({ptr.id(), ptr});
+}
+
+bool has(id_type id) {
+  locker _l(_mutex);
+  return _keeper.find(id) != _keeper.end();
 }
 
 tensor_i::future_type get(id_type id) {
diff --git a/src/Service.cpp b/src/Service.cpp
@@ -39,8 +39,11 @@ namespace x {
 }
 #endif // if 0
 
-struct DeferredService : public Deferred {
-  enum Op : int { REPLICATE, DROP, RUN, SERVICE_LAST };
+// **************************************************************************
+
+struct DeferredService : public DeferredT<Service::service_promise_type,
+                                          Service::service_future_type> {
+  enum Op : int { DROP, RUN, SERVICE_LAST };
 
   id_type _a;
   Op _op;
@@ -51,18 +54,12 @@ struct DeferredService : public Deferred {
 
   void run() {
     switch (_op) {
-    case REPLICATE: {
-      const auto a = std::move(Registry::get(_a).get());
-      auto ddpt = dynamic_cast<DDPTensorImpl *>(a.get());
-      assert(ddpt);
-      ddpt->replicate();
-      set_value(a);
-      break;
-    }
     case RUN:
+      set_value(true);
       break;
     default:
-      throw(std::runtime_error("Unkown Service operation requested."));
+      throw(std::runtime_error(
+          "Execution of unkown service operation requested."));
     }
   }
 
@@ -71,13 +68,14 @@ struct DeferredService : public Deferred {
     switch (_op) {
     case DROP:
       dm.drop(_a);
+      set_value(true);
       // FIXME create delete op and return it
       break;
     case RUN:
-    case REPLICATE:
       return true;
     default:
-      throw(std::runtime_error("Unkown Service operation requested."));
+      throw(std::runtime_error(
+          "MLIR generation for unkown service operation requested."));
     }
 
     return false;
@@ -91,24 +89,52 @@ struct DeferredService : public Deferred {
   }
 };
 
-ddptensor *Service::replicate(const ddptensor &a) {
-  return new ddptensor(
-      defer<DeferredService>(DeferredService::REPLICATE, a.get()));
-}
+// **************************************************************************
 
-void Service::run() {
-  defer<DeferredService>(DeferredService::RUN);
-  // defer_lambda([](){ return true; });
-}
+struct DeferredReplicate : public Deferred {
+  id_type _a;
+
+  DeferredReplicate() : _a() {}
+  DeferredReplicate(const tensor_i::future_type &a) : _a(a.id()) {}
+
+  void run() {
+    const auto a = std::move(Registry::get(_a).get());
+    auto ddpt = dynamic_cast<DDPTensorImpl *>(a.get());
+    assert(ddpt);
+    ddpt->replicate();
+    set_value(a);
+  }
+
+  bool generate_mlir(::mlir::OpBuilder &builder, ::mlir::Location loc,
+                     jit::DepManager &dm) override {
+    return true;
+  }
+
+  FactoryId factory() const { return F_REPLICATE; }
+
+  template <typename S> void serialize(S &ser) {
+    ser.template value<sizeof(_a)>(_a);
+  }
+};
+
+// **************************************************************************
 
 bool inited = false;
 bool finied = false;
 
-void Service::drop(const ddptensor &a) {
+Service::service_future_type Service::drop(const ddptensor &a) {
   if (inited) {
-    // if(getTransceiver()->is_spmd()) getTransceiver()->barrier();
-    defer<DeferredService>(DeferredService::DROP, a.get());
+    return defer<DeferredService>(DeferredService::DROP, a.get());
   }
 }
 
+Service::service_future_type Service::run() {
+  return defer<DeferredService>(DeferredService::RUN);
+}
+
+ddptensor *Service::replicate(const ddptensor &a) {
+  return new ddptensor(defer<DeferredReplicate>(a.get()));
+}
+
 FACTORY_INIT(DeferredService, F_SERVICE);
+FACTORY_INIT(DeferredReplicate, F_REPLICATE);
diff --git a/src/SetGetItem.cpp b/src/SetGetItem.cpp
@@ -273,7 +273,8 @@ struct DeferredSetItem : public Deferred {
   DeferredSetItem(const tensor_i::future_type &a,
                   const tensor_i::future_type &b,
                   const std::vector<py::slice> &v)
-      : _a(a.id()), _b(b.id()), _slc(v) {}
+      : Deferred(a.id(), a.dtype(), a.rank(), a.balanced()), _a(a.id()),
+        _b(b.id()), _slc(v) {}
 
   bool generate_mlir(::mlir::OpBuilder &builder, ::mlir::Location loc,
                      jit::DepManager &dm) override {
@@ -298,19 +299,10 @@ struct DeferredSetItem : public Deferred {
     (void)builder.create<::imex::ptensor::InsertSliceOp>(loc, av, bv, offsV,
                                                          sizesV, stridesV);
     // ... and use av as to later create the ptensor
-    dm.addVal(this->guid(), av,
-              [this](Transceiver *transceiver, uint64_t rank, void *allocated,
-                     void *aligned, intptr_t offset, const intptr_t *sizes,
-                     const intptr_t *strides, uint64_t *gs_allocated,
-                     uint64_t *gs_aligned, uint64_t *lo_allocated,
-                     uint64_t *lo_aligned, uint64_t balanced) {
-                this->set_value(Registry::get(this->_a).get());
-                // this->set_value(std::move(mk_tnsr(dtype, rank, allocated,
-                // aligned, offset, sizes, strides,
-                //                                   gs_allocated, gs_aligned,
-                //                                   lo_allocated,
-                //                                   lo_aligned)));
-              });
+    dm.addReady(this->guid(), [this](id_type guid) {
+      assert(this->guid() == guid);
+      this->set_value(Registry::get(this->_a).get());
+    });
     return false;
   }
 
@@ -368,6 +360,10 @@ struct DeferredGetItem : public Deferred {
     // auto outnd = nd == 0 || _slc.size() == 1 ? 0 : nd;
     auto outTyp =
         ::imex::ptensor::PTensorType::get(shape, oTyp.getElementType());
+    // if(auto dtyp = av.getType().dyn_cast<::imex::dist::DistTensorType>()) {
+    //   av = builder.create<::mlir::UnrealizedConversionCastOp>(loc,
+    //   dtyp.getPTensorType(), av).getResult(0);
+    // }
     // now we can create the PTensor op using the above Values
     auto res = builder.create<::imex::ptensor::SubviewOp>(
         loc, outTyp, av, offsV, sizesV, stridesV);
@@ -413,10 +409,10 @@ 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));
+  a.put(defer<DeferredSetItem>(a.get(), bb.first->get(), v));
   if (bb.second)
     delete bb.first;
-  return res;
+  return &a;
 }
 
 py::object GetItem::get_slice(const ddptensor &a,
diff --git a/src/include/ddptensor/CppTypes.hpp b/src/include/ddptensor/CppTypes.hpp
@@ -210,6 +210,7 @@ enum FactoryId : int {
   F_UNYOP,
   F_GATHER,
   F_GETLOCAL,
+  F_REPLICATE,
   FACTORY_LAST
 };
 
diff --git a/src/include/ddptensor/Registry.hpp b/src/include/ddptensor/Registry.hpp
@@ -26,6 +26,9 @@ tensor_i::future_type get(id_type id);
 /// remove future tensor with guid id from registry
 void del(id_type id);
 
+/// @return true if given guid is registered
+bool has(id_type);
+
 /// finalize registry (before shutdown)
 void fini();
 
diff --git a/src/include/ddptensor/Service.hpp b/src/include/ddptensor/Service.hpp
@@ -6,14 +6,19 @@
 
 #pragma once
 
+#include <future>
+
 class ddptensor;
 
 struct Service {
+  using service_promise_type = std::promise<bool>;
+  using service_future_type = std::shared_future<bool>;
+
   /// replicate the given ddptensor on all ranks
   static ddptensor *replicate(const ddptensor &a);
   /// start running/executing operations, e.g. trigger compile&run
   /// this is not blocking, use futures for synchronization
-  static void run();
+  static service_future_type run();
   /// signal that the given ddptensor is no longer needed and can be deleted
-  static void drop(const ddptensor &a);
+  static service_future_type drop(const ddptensor &a);
 };
diff --git a/src/include/ddptensor/ddptensor.hpp b/src/include/ddptensor/ddptensor.hpp
@@ -15,7 +15,8 @@ class ddptensor {
   tensor_i::future_type _ftx;
 
 public:
-  ddptensor(const tensor_i::future_type &f) : _ftx(f) {}
+  ddptensor(tensor_i::future_type &&f)
+      : _ftx(std::forward<tensor_i::future_type>(f)) {}
   ddptensor(std::shared_future<tensor_i::ptr_type> &&f, id_type id, DTypeId dt,
             int rank, bool balanced)
       : _ftx(std::forward<std::shared_future<tensor_i::ptr_type>>(f), id, dt,
@@ -24,4 +25,7 @@ class ddptensor {
   ~ddptensor() { Service::drop(*this); }
 
   const tensor_i::future_type &get() const { return _ftx; }
+  void const put(tensor_i::future_type &&f) {
+    _ftx = std::forward<tensor_i::future_type>(f);
+  }
 };
diff --git a/src/include/ddptensor/jit/mlir.hpp b/src/include/ddptensor/jit/mlir.hpp
@@ -19,6 +19,7 @@
 #include <mlir/Pass/PassManager.h>
 
 #include <functional>
+#include <map>
 #include <unordered_map>
 #include <utility>
 #include <vector>
@@ -69,21 +70,24 @@ using SetResFunc = std::function<void(
     intptr_t offset, const intptr_t *sizes, const intptr_t *strides,
     uint64_t *gs_allocated, uint64_t *gs_aligned, uint64_t *lo_allocated,
     uint64_t *lo_aligned, uint64_t balanced)>;
+using ReadyFunc = std::function<void(id_type guid)>;
 
 // initialize jit
 void init();
 
 /// Manages iput/output (tensor) dependences
 class DepManager {
 private:
-  using IdValueMap = std::unordered_map<id_type, ::mlir::Value>;
-  using IdCallbackMap = std::unordered_map<id_type, SetResFunc>;
-  using IdRankMap = std::unordered_map<id_type, int>;
+  using IdValueMap = std::map<id_type, ::mlir::Value>;
+  using IdCallbackMap = std::map<id_type, SetResFunc>;
+  using IdReadyMap = std::map<id_type, std::vector<ReadyFunc>>;
+  using IdRankMap = std::map<id_type, int>;
   using ArgList = std::vector<std::pair<id_type, int>>;
 
   ::mlir::func::FuncOp &_func; // MLIR function to which ops are added
   IdValueMap _ivm;             // guid -> mlir::Value
   IdCallbackMap _icm;          // guid -> deliver-callback
+  IdReadyMap _icr;             // guid -> ready-callback
   IdRankMap _irm;              // guid -> rank as computed in MLIR
   ArgList _args;               // input arguments of the generated function
 
@@ -98,6 +102,7 @@ class DepManager {
   /// the promise which generated the value if the tensor is alive when the
   /// function returns it will be added to the list of results
   void addVal(id_type guid, ::mlir::Value val, SetResFunc cb);
+  void addReady(id_type guid, ReadyFunc cb);
 
   /// signals end of lifetime of given tensor: does not need to be returned
   void drop(id_type guid);
diff --git a/src/jit/mlir.cpp b/src/jit/mlir.cpp
diff --git a/test/stencil-2d.py b/test/stencil-2d.py
