adding Dependence Manager for cleaner handling of deps (input, outpput and internal)

Author: fschlimb

src/Creator.cpp

@@ -159,7 +159,7 @@ struct DeferredArange : public Deferred
         // set_value(std::move(TypeDispatch<x::Creator>(_dtype, _start, _end, _step)));
     };
 
-    bool generate_mlir(::mlir::OpBuilder & builder, ::mlir::Location loc, jit::IdValueMap & ivm) override
+    bool generate_mlir(::mlir::OpBuilder & builder, ::mlir::Location loc, jit::DepManager & dm) override
     {
         // create start, stop and step
         auto start = jit::createI64(loc, builder, _start);
@@ -170,13 +170,13 @@ struct DeferredArange : public Deferred
         assert(_dtype == INT64 || _dtype == UINT64); // FIXME
         llvm::SmallVector<int64_t> shape(1, -1); //::mlir::ShapedType::kDynamicSize);
         auto artype = ::imex::ptensor::PTensorType::get(builder.getContext(), ::mlir::RankedTensorType::get(shape, dtype), true);
-        auto ar = builder.create<::imex::ptensor::ARangeOp>(loc, artype, start, end, step, true);
-        auto setter = [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides) {
+        dm.addVal(guid(),
+                  builder.create<::imex::ptensor::ARangeOp>(loc, artype, start, end, step, true),
+                  [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides) {
             assert(rank == 1);
             assert(strides[0] == 1);
             this->set_value(std::move(mk_tnsr(_dtype, rank, allocated, aligned, offset, sizes, strides)));
-        };
-        ivm[_guid] = {ar, setter};
+        });
         return false;
     }
 

src/Deferred.cpp

@@ -60,25 +60,6 @@ void Runable::fini()
     _deferred.clear();
 }
 
-#if 0
-class DepManager
-{
-private:
-    IdValueMap _ivm;
-    std::unordered_set<id_type> _args;
-public:
-    ::mlir::Value getDependent(i::mlir::OpBuilder & builder, d_type guid)
-    {
-        if(auto d = _ivm.find(guid); d == _ivm.end()) {
-            _func.insertArg
-            _ivm[guid] = {val, {}}
-        } else {
-            return d->second.first;
-        }
-    }
-};
-#endif
-
 void process_promises()
 {
     bool done = false;
@@ -87,7 +68,6 @@ void process_promises()
     do {
         ::mlir::OpBuilder builder(&jit._context);
         auto loc = builder.getUnknownLoc();
-        jit::IdValueMap ivp;
 
         // Create a MLIR module
         auto module = builder.create<::mlir::ModuleOp>(loc);
@@ -104,11 +84,13 @@ void process_promises()
         // we need to keep runables/deferred/futures alive until we set their values below
         std::vector<Runable::ptr_type> runables;
 
+        jit::DepManager dm(function);
+
         while(true) {
             Runable::ptr_type d;
             _deferred.pop(d);
             if(d) {
-                if(d->generate_mlir(builder, loc, ivp)) {
+                if(d->generate_mlir(builder, loc, dm)) {
                     d.reset();
                     break;
                 };
@@ -123,39 +105,8 @@ void process_promises()
 
         if(runables.empty()) continue;
 
-        // Now we have to define the return type as a ValueRange of all arrays which we have created
-        // (runnables have put them into ivp when generating mlir)
-        // We also compute the total size of the struct llvm created for this return type
-        // llvm will basically return a struct with all the arrays as members, each of type JIT::MemRefDescriptor
-
-        // Need a container to put all return values, will be used to construct TypeRange
-        std::vector<::mlir::Type> ret_types;
-        ret_types.reserve(ivp.size());
-        std::vector<::mlir::Value> ret_values;
-        ret_types.reserve(ivp.size());
-        std::unordered_map<id_type, uint64_t> rank_map;
-        // here we store the total size of the llvm struct
-        uint64_t sz = 0;
-        for(auto & v : ivp) {
-            auto value = v.second.first;
-            // append the type and array/value
-            ret_types.push_back(value.getType());
-            ret_values.push_back(value);
-            auto ptt = value.getType().dyn_cast<::imex::ptensor::PTensorType>();
-            assert(ptt);
-            auto rank = ptt.getRtensor().getShape().size();
-            rank_map[v.first] = rank;
-            // add sizeof(MemRefDescriptor<elementtype, rank>) to sz
-            sz += 3 + 2 * rank;
-        }
-        ::mlir::TypeRange ret_tr(ret_types);
-        ::mlir::ValueRange ret_vr(ret_values);
-
-        // add return statement
-        auto ret_value = builder.create<::mlir::func::ReturnOp>(loc, ret_vr);
-        // Define and assign correct function type
-        auto funcTypeAttr = ::mlir::TypeAttr::get(builder.getFunctionType({}, ret_tr));
-        function.setFunctionTypeAttr(funcTypeAttr);
+        // create return statement and adjust function type
+        uint64_t sz = dm.handleResult(builder);
         // also request generation of c-wrapper function
         function->setAttr(::mlir::LLVM::LLVMDialect::getEmitCWrapperAttrName(), ::mlir::UnitAttr::get(&jit._context));
         // add the function to the module
@@ -165,22 +116,10 @@ void process_promises()
         // compile and run the module
         assert(sizeof(intptr_t) == sizeof(void*));
         intptr_t * output = new intptr_t[sz];
-        std::cout << ivp.size() << " sz: " << sz << std::endl;
         if(jit.run(module, fname, output)) throw std::runtime_error("failed running jit");
 
-        // push results to fulfill promises
-        size_t pos = 0;
-        for(auto & v : ivp) {
-            auto value = v.second.first;
-            auto rank = rank_map[v.first];
-            void * allocated = (void*)output[pos];
-            void * aligned = (void*)output[pos+1];
-            intptr_t offset = output[pos+2];
-            intptr_t * sizes = output + pos + 3;
-            intptr_t * stride = output + pos + 3 + rank;
-            pos += 3 + 2 * rank;
-            v.second.second(rank, allocated, aligned, offset, sizes, stride);
-        }
+        // push results to deliver promises
+        dm.deliver(output, sz);
     } while(!done);
 }
 

src/EWBinOp.cpp

@@ -456,15 +456,16 @@ struct DeferredEWBinOp : public Deferred
 #endif
     }
 
-    bool generate_mlir(::mlir::OpBuilder & builder, ::mlir::Location loc, jit::IdValueMap & ivm) override
+    bool generate_mlir(::mlir::OpBuilder & builder, ::mlir::Location loc, jit::DepManager & dm) override
     {
         // FIXME the type of the result is based on a only
-        auto rtyp = ivm[_a].first.getType();
-        auto ewbo = builder.create<::imex::ptensor::EWBinOp>(loc, rtyp, builder.getI32IntegerAttr(ddpt2mlir(_op)), ivm[_a].first, ivm[_b].first);
-        auto setter = [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides) {
+        auto a = dm.getDependent(builder, _a);
+        auto b = dm.getDependent(builder, _b);
+        dm.addVal(guid(),
+                  builder.create<::imex::ptensor::EWBinOp>(loc, a.getType(), builder.getI32IntegerAttr(ddpt2mlir(_op)), a, b),
+                  [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides) {
             this->set_value(std::move(mk_tnsr(_dtype, rank, allocated, aligned, offset, sizes, strides)));
-        };
-        ivm[_guid] = {ewbo, setter};
+        });
         return false;
     }
 

src/ReduceOp.cpp

@@ -116,10 +116,10 @@ struct DeferredReduceOp : public Deferred
 #endif
     }
 
-    bool generate_mlir(::mlir::OpBuilder & builder, ::mlir::Location loc, jit::IdValueMap & ivm) override
+    bool generate_mlir(::mlir::OpBuilder & builder, ::mlir::Location loc, jit::DepManager & dm) override
     {
         // FIXME reduction over individual dimensions is not supported
-        auto a = ivm[_a].first;
+        auto a = dm.getDependent(builder, _a);
         auto a_ptt = a.getType().dyn_cast<::imex::ptensor::PTensorType>();
         assert(a_ptt);
 
@@ -128,11 +128,11 @@ struct DeferredReduceOp : public Deferred
             ::mlir::RankedTensorType::get(llvm::SmallVector<int64_t>(), a_ptt.getRtensor().getElementType()),
             true
         );
-        auto rop = builder.create<::imex::ptensor::ReductionOp>(loc, rtyp, builder.getI32IntegerAttr(ddpt2mlir(_op)), a);
-        auto setter = [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides) {
+        dm.addVal(guid(),
+                  builder.create<::imex::ptensor::ReductionOp>(loc, rtyp, builder.getI32IntegerAttr(ddpt2mlir(_op)), a),
+                  [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides) {
             this->set_value(std::move(mk_tnsr(_dtype, rank, allocated, aligned, offset, sizes, strides)));
-        };
-        ivm[_guid] = {rop, setter};
+        });
         return false;
     }
 

src/Service.cpp

@@ -68,14 +68,12 @@ struct DeferredService : public Deferred
 #endif
     }
 
-    bool generate_mlir(::mlir::OpBuilder & builder, ::mlir::Location loc, jit::IdValueMap & ivm) override
+    bool generate_mlir(::mlir::OpBuilder & builder, ::mlir::Location loc, jit::DepManager & dm) override
     {
         switch(_op) {
         case DROP:
-            if(auto e = ivm.find(_a); e != ivm.end()) {
-                ivm.erase(e);
-                // FIXME create delete op and return it
-            }
+            dm.drop(_a);
+            // FIXME create delete op and return it
             break;
         case RUN:
             return true;

src/ddptensor.cpp

@@ -137,8 +137,7 @@ PYBIND11_MODULE(_ddptensor, m) {
         .def("_get_slice", &GetItem::get_slice)
         .def("_get_local", &GetItem::get_local)
         .def("_gather", &GetItem::gather)
-        .def("to_numpy", &IO::to_numpy)
-        .def("ttt", &jit::ttt);
+        .def("to_numpy", &IO::to_numpy);
 
     py::class_<Creator>(m, "Creator")
         .def("create_from_shape", &Creator::create_from_shape)

src/include/ddptensor/Deferred.hpp

@@ -17,7 +17,7 @@ struct Runable
     virtual void run() = 0;
     /// generate MLIR code for jit
     /// @return true if last operation in to-be-compiled region, false otherwise
-    virtual bool generate_mlir(::mlir::OpBuilder &, ::mlir::Location, jit::IdValueMap &)
+    virtual bool generate_mlir(::mlir::OpBuilder &, ::mlir::Location, jit::DepManager &)
     {
         throw(std::runtime_error("No MLIR support for this operation."));
         return false;
@@ -145,7 +145,7 @@ struct DeferredLambda : public Runable
         _l();
     }
 
-    bool generate_mlir(::mlir::OpBuilder &, ::mlir::Location, jit::IdValueMap &)
+    bool generate_mlir(::mlir::OpBuilder &, ::mlir::Location, jit::DepManager &)
     {
         return _l();
     }

src/include/ddptensor/jit/mlir.hpp

@@ -8,26 +8,60 @@
 #include <mlir/IR/MLIRContext.h>
 #include <mlir/Pass/PassManager.h>
 #include <mlir/IR/Builders.h>
+#include <mlir/Dialect/Func/IR/FuncOps.h>
 
 #include <unordered_map>
 #include <functional>
 #include  <utility>
+#include <vector>
 
 namespace jit {
 
 // function type for building body for linalg::generic
 using SetResFunc = std::function<void(
     uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides)>;
-using IdValueMap = std::unordered_map<id_type, std::pair<::mlir::Value, SetResFunc>>;
 
 // initialize jit
 void init();
 
-void ttt();
-
 // create a constant integer with given value
 extern ::mlir::Value createI64(const ::mlir::Location & loc, ::mlir::OpBuilder & builder, int64_t val);
 
+/// Manages iput/output (tensor) dependences
+class DepManager
+{
+private:
+    using IdValueMap = std::unordered_map<id_type, std::pair<::mlir::Value, SetResFunc>>;
+    using IdRankMap = std::unordered_map<id_type, int>;
+    ::mlir::func::FuncOp & _func; // MLIR function to which ops are added
+    IdValueMap _ivm;              // guid -> {mlir::Value, deliver-callback}
+    IdRankMap _irm;               // guid -> rank as computed in MLIR
+    std::vector<id_type> _args;   // input args to generated function
+
+public:
+    DepManager(::mlir::func::FuncOp & f)
+    : _func(f)
+    {}
+    /// @return the ::mlir::Value representing the tensor with guid guid
+    /// If the tensor is not created wtihin the current function, it will
+    /// be added as a function argument.
+    ::mlir::Value getDependent(::mlir::OpBuilder & builder, id_type guid);
+
+    /// for the tensor guid register the ::mlir::value and a callback to deliver 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);
+
+    /// signals end of lifetime of given tensor: does not need to be returned
+    void drop(id_type guid);
+
+    /// create return statement and add results to function
+    /// @return size of output in number of intptr_t's 
+    uint64_t handleResult(::mlir::OpBuilder & builder);
+
+    /// devlier promise after execution
+    void deliver(intptr_t *, uint64_t);
+};
+
 // A class to manage the MLIR business (compilation and execution).
 // Just a stub for now, will need to be extended with paramters and maybe more.
 class JIT {