Dist runtime (#49)

* updating to new DistRuntimeDilaect in imex

Author: fschlimb

CMakeLists.txt

@@ -177,6 +177,8 @@ target_link_libraries(_ddpt_rt PRIVATE
     IMEXDistDialect
     IMEXDistTransforms
     IMEXDistToStandard
+    IMEXDistRuntimeDialect
+    IMEXDistRuntimeTransforms
     IMEXUtil
     IMEXTransforms
 	  MLIROptLib

imex_version.txt

@@ -1 +1 @@
-a3e33e0756c29d004736f7c655f98e9ee085eebb
+2e1500739abe2e6195171723c3e85c3c2f2266de

src/Deferred.cpp

@@ -14,6 +14,7 @@
 #include "include/ddptensor/Transceiver.hpp"
 
 #include <imex/Dialect/Dist/IR/DistOps.h>
+#include <imex/Dialect/DistRuntime/IR/DistRuntimeOps.h>
 #include <imex/Dialect/PTensor/IR/PTensorOps.h>
 #include <mlir/Dialect/Func/IR/FuncOps.h>
 #include <mlir/Dialect/LLVMIR/LLVMDialect.h>
@@ -94,7 +95,7 @@ void process_promises() {
 
     // Create a MLIR module
     auto module = builder.create<::mlir::ModuleOp>(loc);
-    auto protos = builder.create<::imex::dist::RuntimePrototypesOp>(loc);
+    auto protos = builder.create<::imex::distruntime::RuntimePrototypesOp>(loc);
     module.push_back(protos);
     // Create the jit func
     // create dummy type, we'll replace it with the actual type later

src/idtr.cpp

@@ -464,7 +464,7 @@ void _idtr_update_halo(DTypeId ddpttype, int64_t ndims, int64_t *ownedOff,
   tc = getTransceiver();
 #endif
   auto nworkers = tc->nranks();
-  if (nworkers <= 1)
+  if (nworkers <= 1 || getenv("DDPT_SKIP_COMM"))
     return;
   auto myWorkerIndex = tc->rank();
 

src/jit/mlir.cpp

@@ -435,6 +435,7 @@ static const char *pass_pipeline =
                             "func.func(dist-infer-elementwise-cores),"
                             "convert-dist-to-standard,"
                             "canonicalize,"
+                            "lower-distruntime-to-idtr,"
                             "convert-ptensor-to-linalg,"
                             "canonicalize,"
                             "func.func(tosa-to-linalg),"
@@ -483,6 +484,7 @@ JIT::JIT()
   _context.getOrLoadDialect<::mlir::linalg::LinalgDialect>();
   _context.getOrLoadDialect<::imex::ptensor::PTensorDialect>();
   _context.getOrLoadDialect<::imex::dist::DistDialect>();
+  _context.getOrLoadDialect<::imex::distruntime::DistRuntimeDialect>();
   // create the pass pipeline from string
   if (::mlir::failed(::mlir::parsePassPipeline(pass_pipeline, _pm)))
     throw std::runtime_error("failed to parse pass pipeline");
@@ -582,6 +584,7 @@ void init() {
 
   ::imex::registerPTensorPasses();
   ::imex::registerDistPasses();
+  ::imex::registerDistRuntimePasses();
   ::imex::registerConvertDistToStandard();
   ::imex::registerConvertPTensorToLinalg();