adding setitem and adjustsing to new imex

Author: fschlimb

ddptensor/ddptensor.py

@@ -54,4 +54,4 @@ def __getitem__(self, key):
         return dtensor(self._t.__getitem__(key if isinstance(key, list) else [key,]))
 
     def __setitem__(self, key, value):
-         self._t = self._t.__setitem__(key, value._t) # if isinstance(value, dtensor) else value)
+         self._t.__setitem__(key if isinstance(key, list) else [key,], value._t) # if isinstance(value, dtensor) else value)

src/Creator.cpp

@@ -5,7 +5,7 @@
 #include "ddptensor/DDPTensorImpl.hpp"
 
 #include <imex/Dialect/PTensor/IR/PTensorOps.h>
-#include <imex/internal/PassUtils.h>
+#include <imex/Utils/PassUtils.h>
 
 #include <mlir/IR/Builders.h>
 #include <mlir/Dialect/Arith/IR/Arith.h>

src/SetGetItem.cpp

@@ -7,7 +7,7 @@
 
 #include <imex/Dialect/PTensor/IR/PTensorOps.h>
 #include <imex/Dialect/Dist/IR/DistOps.h>
-#include <imex/internal/PassUtils.h>
+#include <imex/Utils/PassUtils.h>
 #include <mlir/IR/Builders.h>
 
 #if 0
@@ -202,6 +202,38 @@ struct DeferredSetItem : public Deferred
         //set_value(std::move(TypeDispatch<x::SetItem>(a, b, _slc, _b)));
     }
 
+    bool generate_mlir(::mlir::OpBuilder & builder, ::mlir::Location loc, jit::DepManager & dm) override
+    {
+        // get params and extract offsets/sizes/strides
+        const auto dtype = this->dtype();
+        auto av = dm.getDependent(builder, _a);
+        auto bv = dm.getDependent(builder, _b);
+        auto & offs = _slc.offsets();
+        auto & sizes = _slc.sizes();
+        auto & strides = _slc.strides();
+        auto nd = offs.size();
+        // convert C++ slices into vectors of MLIR Values
+        std::vector<::mlir::Value> offsV(nd);
+        std::vector<::mlir::Value> sizesV(nd);
+        std::vector<::mlir::Value> stridesV(nd);
+        for(auto i = 0; i<nd; ++i) {
+            offsV[i] = ::imex::createIndex(loc, builder, offs[i]);
+            sizesV[i] = ::imex::createIndex(loc, builder, sizes[i]);
+            stridesV[i] = ::imex::createIndex(loc, builder, strides[i]);
+        }
+        // insertsliceop has no return value, so we just craete the op...
+        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](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) {
+            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)));
+        });
+        return false;
+    }
+
     FactoryId factory() const
     {
         return F_SETITEM;
@@ -264,7 +296,7 @@ struct DeferredGetItem : public Deferred
                       sizesV,
                       stridesV),
                   [this, dtype](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 * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned) {
             this->set_value(std::move(mk_tnsr(dtype, rank, allocated, aligned, offset, sizes, strides,
                                               gs_allocated, gs_aligned, lo_allocated, lo_aligned)));
         });

src/idtr.cpp

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 #include <ddptensor/idtr.hpp>
+#include <ddptensor/jit/mlir.hpp>
 #include <ddptensor/DDPTensorImpl.hpp>
 #include <ddptensor/MPITransceiver.hpp>
 
@@ -120,14 +121,15 @@ static ReduceOpId mlir2ddpt(const ::imex::ptensor::ReduceOpId rop)
 // Elementwise inplace allreduce
 void idtr_reduce_all(void * inout, DTypeId dtype, uint64_t N, int op)
 {
-
     getTransceiver()->reduce_all(inout, dtype, N, mlir2ddpt(static_cast<imex::ptensor::ReduceOpId>(op)));
 }
 
 // FIXME hard-coded 0d tensor
-void _idtr_reduce_all(uint64_t * allocated, uint64_t * aligned, uint64_t offset, DTypeId dtype, int op)
+void _idtr_reduce_all(uint64_t rank, uint64_t * mrd, DTypeId dtype, int op)
 {
-    idtr_reduce_all(aligned + offset, dtype, 1, op);
+    assert(rank==0);
+    auto descr = reinterpret_cast<jit::JIT::MemRefDescriptor<uint64_t, 0>*>(mrd);
+    idtr_reduce_all(descr->aligned + descr->offset, dtype, 1, op);
 }
 
 } // extern "C"

src/jit/mlir.cpp

@@ -296,7 +296,7 @@ int JIT::run(::mlir::ModuleOp & module, const std::string & fname, std::vector<v
 static const char * pass_pipeline =
    getenv("DDPT_PASSES")
    ? getenv("DDPT_PASSES")
-   : "func.func(ptensor-dist),convert-dist-to-standard,convert-ptensor-to-linalg,arith-expand,canonicalize,arith-bufferize,func.func(empty-tensor-to-alloc-tensor,scf-bufferize,linalg-bufferize,tensor-bufferize),func-bufferize,canonicalize,func.func(finalizing-bufferize,convert-linalg-to-parallel-loops),canonicalize,func.func(lower-affine),fold-memref-alias-ops,convert-scf-to-cf,convert-memref-to-llvm,convert-func-to-llvm,convert-scf-to-cf,reconcile-unrealized-casts";
+   : "func.func(ptensor-dist),convert-dist-to-standard,convert-ptensor-to-linalg,arith-expand,canonicalize,arith-bufferize,func.func(empty-tensor-to-alloc-tensor,scf-bufferize,linalg-bufferize,tensor-bufferize),func-bufferize,canonicalize,func.func(finalizing-bufferize,convert-linalg-to-parallel-loops),canonicalize,fold-memref-alias-ops,lower-affine,convert-scf-to-cf,convert-memref-to-llvm,convert-func-to-llvm,reconcile-unrealized-casts";
 
 JIT::JIT()
     : _context(::mlir::MLIRContext::Threading::DISABLED),