Unitsize (#37)

* updates for latest imex
* generating correct GShape in DistTensorType args and for subview/getitem

---------

Co-authored-by: Tuomas Karna <tuomas.karna@intel.com>

Author: fschlimb

CMakeLists.txt

@@ -186,7 +186,6 @@ target_link_libraries(_ddpt_rt PRIVATE
     MLIRFuncToLLVM
     MLIRFuncTransforms
     MLIRLinalgDialect
-    MLIRLinalgToLLVM
     MLIRLinalgTransforms
     MLIRLLVMDialect
     MLIRMathDialect

examples/stencil-2d.py

@@ -181,7 +181,7 @@ def main():
                 for s in range(-r, r + 1):
                     for t in range(-r, r + 1):
                         B[r:b, r:b] += W[r + t, r + s] * A[r + t : b + t, r + s : b + s]
-        A = A + 1.0
+        A[:, :] = A + 1.0
 
     np.sync()
     t1 = timer()

imex_version.txt

@@ -1 +1 @@
-cfe20696de8126dee8fcd596ee8e4cc8522977f7
+e2ecbb3c08e4a87256cdf11d8057fe765d3a04b8

setup.py

@@ -40,7 +40,8 @@ def build_cmake(self, ext):
         build_args = [
             "--config",
             config,
-            # '--', '-j8'
+            "-j8"
+            # '--', '-j4'
         ]
 
         os.chdir(str(build_temp))

src/SetGetItem.cpp

@@ -280,9 +280,10 @@ struct DeferredGetItem : public Deferred {
     // get params and extract offsets/sizes/strides
     const auto dtype = this->dtype();
     auto av = dm.getDependent(builder, _a);
-    auto &offs = _slc.offsets();
-    auto &sizes = _slc.sizes();
-    auto &strides = _slc.strides();
+    const auto &offs = _slc.offsets();
+    const auto &sizes =
+        shape(); // we already converted ALL_SIZE as much as posible
+    const auto &strides = _slc.strides();
     auto nd = offs.size();
     // convert C++ slices into vectors of MLIR Values
     std::vector<::mlir::OpFoldResult> offsV(nd);
@@ -291,7 +292,7 @@ struct DeferredGetItem : public Deferred {
     for (auto i = 0; i < nd; ++i) {
       offsV[i] = ::imex::createIndex(loc, builder, offs[i]);
       stridesV[i] = ::imex::createIndex(loc, builder, strides[i]);
-      if (sizes[i] == ALL_SIZE) {
+      if (sizes[i] < 0) {
         sizesV[i] =
             builder.create<::imex::ptensor::DimOp>(loc, av, i).getResult();
       } else {

src/jit/mlir.cpp

@@ -117,9 +117,10 @@ static ::mlir::Type makeSignlessType(::mlir::Type type) {
 
 // convert ddpt's DTYpeId into MLIR type
 static ::mlir::Type getTType(::mlir::OpBuilder &builder, DTypeId dtype,
-                             ::mlir::SmallVector<int64_t> &lhShape,
-                             ::mlir::SmallVector<int64_t> &ownShape,
-                             ::mlir::SmallVector<int64_t> &rhShape,
+                             const ::mlir::SmallVector<int64_t> &gShape,
+                             const ::mlir::SmallVector<int64_t> &lhShape,
+                             const ::mlir::SmallVector<int64_t> &ownShape,
+                             const ::mlir::SmallVector<int64_t> &rhShape,
                              uint64_t team, bool balanced) {
   ::mlir::Type etyp;
 
@@ -154,9 +155,7 @@ static ::mlir::Type getTType(::mlir::OpBuilder &builder, DTypeId dtype,
   };
 
   if (team) {
-    if (ownShape.size()) {
-      auto gShape = ownShape;
-      gShape[0] += lhShape[0] + rhShape[0];
+    if (gShape.size()) {
       return ::imex::dist::DistTensorType::get(gShape, etyp,
                                                {lhShape, ownShape, rhShape});
     } else {
@@ -183,8 +182,10 @@ ::mlir::Value DepManager::getDependent(::mlir::OpBuilder &builder,
       ownShape[i] = impl->local_shape()[i];
       rhShape[i] = impl->rh_shape()[i];
     }
-    auto typ = getTType(builder, fut.dtype(), lhShape, ownShape, rhShape,
-                        fut.team(), fut.balanced());
+    auto typ = getTType(
+        builder, fut.dtype(),
+        ::mlir::SmallVector<int64_t>(impl->shape(), impl->shape() + rank),
+        lhShape, ownShape, rhShape, fut.team(), fut.balanced());
     _func.insertArgument(idx, typ, {}, loc);
     auto val = _func.getArgument(idx);
     _args.push_back({guid, std::move(fut)});
@@ -516,7 +517,8 @@ JIT::JIT()
   crunner = crunner ? crunner : "libmlir_c_runner_utils.so";
   const char *idtr = getenv("DDPT_IDTR_SO");
   idtr = idtr ? idtr : "libidtr.so";
-  _sharedLibPaths = {idtr, crunner};
+  _sharedLibPaths = {idtr, crunner,
+                     "/home/fschlimb/llvm/lib/libmlir_runner_utils.so"};
 
   // detect target architecture
   auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
@@ -562,14 +564,12 @@ void init() {
   ::mlir::registerConvertFuncToLLVMPass();
   ::mlir::bufferization::registerBufferizationPasses();
   ::mlir::arith::registerArithPasses();
-  ::mlir::registerAffinePasses();
   ::mlir::registerCanonicalizerPass();
   ::mlir::registerConvertAffineToStandardPass();
   ::mlir::registerFinalizeMemRefToLLVMConversionPass();
   ::mlir::registerArithToLLVMConversionPass();
   ::mlir::registerConvertMathToLLVMPass();
   ::mlir::registerConvertControlFlowToLLVMPass();
-  ::mlir::registerConvertLinalgToLLVMPass();
   ::mlir::registerConvertOpenMPToLLVMPass();
   ::mlir::memref::registerMemRefPasses();
   ::mlir::registerReconcileUnrealizedCastsPass();