now releasing memory of tensor data; protecting all runtime output by _verbose flag (#11)

Author: fschlimb

src/DDPTensorImpl.cpp

@@ -56,6 +56,16 @@ DDPTensorImpl::DDPTensorImpl(const int64_t *shape, uint64_t N, rank_type owner)
   assert(!_transceiver || _transceiver == getTransceiver());
 }
 
+DDPTensorImpl::~DDPTensorImpl() {
+  if (!_base) {
+    delete[] reinterpret_cast<char *>(_allocated);
+  }
+  delete[] _gs_allocated;
+  delete[] _lo_allocated;
+  delete[] _sizes;
+  delete[] _strides;
+}
+
 DDPTensorImpl::ptr_type DDPTensorImpl::clone(bool copy) {
   // FIXME memory leak
   auto nd = ndims();
@@ -163,8 +173,8 @@ int64_t DDPTensorImpl::__int__() const {
   return res;
 }
 
-void DDPTensorImpl::add_to_args(std::vector<void *> &args, int ndims) {
-  assert(ndims == this->ndims());
+void DDPTensorImpl::add_to_args(std::vector<void *> &args) {
+  int ndims = this->ndims();
   // local tensor first
   intptr_t *buff = new intptr_t[dtensor_sz(ndims)];
   buff[0] = reinterpret_cast<intptr_t>(_allocated);

src/Deferred.cpp

@@ -89,7 +89,7 @@ void process_promises() {
   bool done = false;
   jit::JIT jit;
   do {
-    ::mlir::OpBuilder builder(&jit._context);
+    ::mlir::OpBuilder builder(&jit.context());
     auto loc = builder.getUnknownLoc();
 
     // Create a MLIR module
@@ -146,7 +146,8 @@ void process_promises() {
       // also request generation of c-wrapper function
       function->setAttr(::mlir::LLVM::LLVMDialect::getEmitCWrapperAttrName(),
                         builder.getUnitAttr());
-      function.getFunctionType().dump();
+      if (jit.verbose())
+        function.getFunctionType().dump();
       // add the function to the module
       module.push_back(function);
 
@@ -159,7 +160,8 @@ void process_promises() {
         // push results to deliver promises
         dm.deliver(output, osz);
       } else {
-        std::cerr << "\tskipping\n";
+        if (jit.verbose())
+          std::cerr << "\tskipping\n";
       }
     } // no else needed
 

src/SetGetItem.cpp

@@ -292,18 +292,22 @@ struct DeferredGetItem : public Deferred {
     // now we can create the PTensor op using the above Values
     auto res = builder.create<::imex::ptensor::SubviewOp>(
         loc, outTyp, av, offsV, sizesV, stridesV);
-    dm.addVal(this->guid(), res,
-              [this, dtype](Transceiver *transceiver, uint64_t rank,
-                            void *allocated, void *aligned, intptr_t offset,
-                            const intptr_t *sizes, const intptr_t *strides,
-                            int64_t *gs_allocated, int64_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 future_a = Registry::get(_a);
+
+    dm.addVal(
+        this->guid(), res,
+        [this, dtype, future_a](
+            Transceiver *transceiver, uint64_t rank, void *allocated,
+            void *aligned, intptr_t offset, const intptr_t *sizes,
+            const intptr_t *strides, int64_t *gs_allocated, int64_t *gs_aligned,
+            uint64_t *lo_allocated, uint64_t *lo_aligned, uint64_t balanced) {
+          auto t = mk_tnsr(transceiver, dtype, rank, allocated, aligned, offset,
+                           sizes, strides, gs_allocated, gs_aligned,
+                           lo_allocated, lo_aligned, balanced);
+          t->set_base(future_a.get());
+          this->set_value(std::move(t));
+        });
     return false;
   }
 

src/include/ddptensor/DDPTensorImpl.hpp

@@ -35,6 +35,7 @@ class DDPTensorImpl : public tensor_i {
   uint64_t _ndims = 0;
   uint64_t _balanced = 1;
   DTypeId _dtype = DTYPE_LAST;
+  tensor_i::ptr_type _base;
 
 public:
   using ptr_type = std::shared_ptr<DDPTensorImpl>;
@@ -66,11 +67,10 @@ class DDPTensorImpl : public tensor_i {
   // helper
   void alloc(bool all = true);
 
-  // FIXME proper lifetime management for tensor data
-  ~DDPTensorImpl() {
-    delete[] _sizes;
-    delete[] _strides;
-  }
+  // set the base tensor
+  void set_base(const tensor_i::ptr_type &base) { _base = base; }
+
+  virtual ~DDPTensorImpl();
 
   // @return pointer to raw data
   void *data();
@@ -158,7 +158,7 @@ class DDPTensorImpl : public tensor_i {
 
   /// add tensor to list of args in the format expected by MLIR
   /// assuming tensor has ndims dims.
-  virtual void add_to_args(std::vector<void *> &args, int ndims) override;
+  virtual void add_to_args(std::vector<void *> &args) override;
 
   /// @return local offsets into global tensor
   const uint64_t *local_offsets() const { return _lo_aligned; }

src/include/ddptensor/jit/mlir.hpp

@@ -3,6 +3,7 @@
 #pragma once
 
 #include <ddptensor/CppTypes.hpp>
+#include <ddptensor/tensor_i.hpp>
 
 #include <imex/Dialect/PTensor/IR/PTensorOps.h>
 #include <mlir/Conversion/LLVMCommon/MemRefBuilder.h>
@@ -104,14 +105,15 @@ using ReadyFunc = std::function<void(id_type guid)>;
 // initialize jit
 void init();
 
-/// Manages iput/output (tensor) dependences
+/// Manages iput/output (tensor) dependencies
 class DepManager {
 private:
   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>>;
+  // here we store the future (and not the guid) to keep it alive long enough
+  using ArgList = std::vector<std::pair<id_type, tensor_i::future_type>>;
 
   ::mlir::func::FuncOp &_func; // MLIR function to which ops are added
   IdValueMap _ivm;             // guid -> mlir::Value
@@ -166,6 +168,10 @@ class JIT {
   std::vector<intptr_t> run(::mlir::ModuleOp &, const std::string &,
                             std::vector<void *> &, size_t);
 
+  int verbose() { return _verbose; };
+  ::mlir::MLIRContext &context() { return _context; };
+
+private:
   ::mlir::MLIRContext _context;
   ::mlir::PassManager _pm;
   int _verbose;

src/include/ddptensor/tensor_i.hpp

@@ -106,7 +106,7 @@ class tensor_i {
   // store tensor information in form of corresponding
   // jit::JIT::DistMemRefDescriptor
   // @return stored size in number of intptr_t
-  virtual void add_to_args(std::vector<void *> &args, int ndims) = 0;
+  virtual void add_to_args(std::vector<void *> &args) = 0;
 };
 
 #if 0

src/jit/mlir.cpp

@@ -169,7 +169,7 @@ ::mlir::Value DepManager::getDependent(::mlir::OpBuilder &builder,
         getTType(builder, fut.dtype(), fut.rank(), fut.team(), fut.balanced());
     _func.insertArgument(idx, typ, {}, loc);
     auto val = _func.getArgument(idx);
-    _args.push_back({guid, fut.rank()});
+    _args.push_back({guid, std::move(fut)});
     _ivm[guid] = val;
     return val;
   } else {
@@ -180,8 +180,7 @@ ::mlir::Value DepManager::getDependent(::mlir::OpBuilder &builder,
 std::vector<void *> DepManager::store_inputs() {
   std::vector<void *> res;
   for (auto a : _args) {
-    auto f = Registry::get(a.first);
-    f.get().get()->add_to_args(res, a.second);
+    a.second.get().get()->add_to_args(res);
     _ivm.erase(a.first); // inputs need no delivery
     _icm.erase(a.first);
   }
@@ -202,7 +201,7 @@ void DepManager::drop(id_type guid) {
   _ivm.erase(guid);
   _icm.erase(guid);
   _icr.erase(guid);
-  // Registry::del(guid);
+  Registry::del(guid);
   // FIXME create delete op
 }
 
@@ -338,9 +337,11 @@ std::vector<intptr_t> JIT::run(::mlir::ModuleOp &module,
     }
     if (cached) {
       module = cached;
-      // std::cerr << "using cached module" << std::endl;
+      if (_verbose)
+        std::cerr << "cached..." << std::endl;
     } else {
-      // std::cerr << "compiling..." << std::endl;
+      if (_verbose)
+        std::cerr << "compiling..." << std::endl;
       cache.push_back(std::make_pair(cksm, module));
       if (_verbose > 1)
         module.dump();
@@ -432,6 +433,7 @@ JIT::JIT()
   }
   // some verbosity
   if (_verbose) {
+    std::cerr << "pass pipeline: " << pass_pipeline << std::endl;
     // _pm.enableStatistics();
     _pm.enableTiming();
     // if(_verbose > 1)