start cleaning repo

usainzg · Oct 28, 2024 · 9653e27 · 9653e27
1 parent e0a2cee
commit 9653e27
Show file tree

Hide file tree

Showing 4 changed files with 110 additions and 50 deletions.
diff --git a/lib/Transform/Affine/AffineDistributeToMPI.cpp b/lib/Transform/Affine/AffineDistributeToMPI.cpp
@@ -1,7 +1,9 @@
 #include "lib/Transform/Affine/AffineDistributeToMPI.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/MPI/IR/MPI.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/include/mlir/Pass/Pass.h"
 
@@ -70,6 +72,50 @@ struct AffineDistributeToMPI
 
     // create affine loop for the second half
     // receive processed first half
+
+    // send first half of data to rank 1
+    // get all memref operands from the loop body
+    SmallVector<Value, 4> memrefOperands;
+    forOp.walk([&](Operation *op) {
+      if (auto loadOp = dyn_cast<memref::LoadOp>(op)) {
+        if (!llvm::is_contained(memrefOperands, loadOp.getMemref()))
+          memrefOperands.push_back(loadOp.getMemref());
+      }
+      if (auto storeOp = dyn_cast<memref::StoreOp>(op)) {
+        if (!llvm::is_contained(memrefOperands, storeOp.getMemref()))
+          memrefOperands.push_back(storeOp.getMemref());
+      }
+    });
+
+    // send each memref to rank 1
+    for (auto memref : memrefOperands) {
+      builder.create<mpi::SendOp>(loc, retvalType, memref,
+                                  /*dest=*/builder.getI32IntegerAttr(1),
+                                  /*tag=*/builder.getI32IntegerAttr(0));
+    }
+
+    // create affine loop for the second half
+    auto upperBound = forOp.getUpperBound();
+    auto lowerBound = getHalfPoint(builder, forOp);
+
+    auto newLoop = builder.create<AffineForOp>(
+        loc, lowerBound, upperBound,
+        [&](OpBuilder &nestedBuilder, Location loc, ValueRange ivs) {
+          // clone the original loop body
+        BlockAndValueMapping mapping;
+          mapping.map(forOp.getInductionVar(), ivs.front());
+          for (auto &op : forOp.getBody()->without_terminator())
+            nestedBuilder.clone(op, mapping);
+        });
+
+    // receive processed first half
+    // only receive the result memref (assumed to be the last operand)
+    if (!memrefOperands.empty()) {
+      auto resultMemref = memrefOperands.back();
+      builder.create<mpi::RecvOp>(loc, retvalType, resultMemref,
+                                  /*source=*/builder.getI32IntegerAttr(1),
+                                  /*tag=*/builder.getI32IntegerAttr(0));
+    }
   }
 
   void processRankOne(OpBuilder &builder, affine::AffineForOp forOp) {
@@ -79,6 +125,18 @@ struct AffineDistributeToMPI
     // send result back to rank 0
     // cleanup local buffers (memrefs dealloc)
   }
+
+  // helper function to get the midpoint of the loop range
+  AffineMap getHalfPoint(OpBuilder &builder, AffineForOp forOp) {
+    auto context = builder.getContext();
+    auto upperMap = forOp.getUpperBoundMap();
+    auto upperBound = upperMap.getResult(0);
+
+    // create an affine map that divides the upper bound by 2
+    auto halfExpr = upperBound.floorDiv(2);
+    return AffineMap::get(upperMap.getNumDims(), upperMap.getNumSymbols(),
+                          halfExpr, context);
+  }
 };
 
 } // namespace tutorial

diff --git a/lib/Transform/Affine/BUILD b/lib/Transform/Affine/BUILD
@@ -42,52 +42,53 @@ cc_library(
       "@llvm-project//mlir:MPIDialect",
       "@llvm-project//mlir:SCFDialect",
       "@llvm-project//mlir:FuncDialect",
+      "@llvm-project//mlir:MemRefDialect",
       "@llvm-project//mlir:Pass",
       "@llvm-project//mlir:Transforms",
   ],
 )
 
-cc_library(
-    name = "AffineFullUnroll",
-    srcs = ["AffineFullUnroll.cpp"],
-    hdrs = [
-        "AffineFullUnroll.h",
-        "Passes.h",
-    ],
-    deps = [
-        ":pass_inc_gen",
-        "@llvm-project//mlir:AffineDialect",
-        "@llvm-project//mlir:AffineUtils",
-        "@llvm-project//mlir:FuncDialect",
-        "@llvm-project//mlir:Pass",
-        "@llvm-project//mlir:Transforms",
-    ],
-)
+# cc_library(
+#     name = "AffineFullUnroll",
+#     srcs = ["AffineFullUnroll.cpp"],
+#     hdrs = [
+#         "AffineFullUnroll.h",
+#         "Passes.h",
+#     ],
+#     deps = [
+#         ":pass_inc_gen",
+#         "@llvm-project//mlir:AffineDialect",
+#         "@llvm-project//mlir:AffineUtils",
+#         "@llvm-project//mlir:FuncDialect",
+#         "@llvm-project//mlir:Pass",
+#         "@llvm-project//mlir:Transforms",
+#     ],
+# )
 
-cc_library(
-    name = "AffineFullUnrollPatternRewrite",
-    srcs = ["AffineFullUnrollPatternRewrite.cpp"],
-    hdrs = [
-        "AffineFullUnrollPatternRewrite.h",
-        "Passes.h",
-    ],
-    deps = [
-        ":pass_inc_gen",
-        "@llvm-project//mlir:AffineDialect",
-        "@llvm-project//mlir:AffineUtils",
-        "@llvm-project//mlir:FuncDialect",
-        "@llvm-project//mlir:Pass",
-        "@llvm-project//mlir:Transforms",
-    ],
-)
+# cc_library(
+#     name = "AffineFullUnrollPatternRewrite",
+#     srcs = ["AffineFullUnrollPatternRewrite.cpp"],
+#     hdrs = [
+#         "AffineFullUnrollPatternRewrite.h",
+#         "Passes.h",
+#     ],
+#     deps = [
+#         ":pass_inc_gen",
+#         "@llvm-project//mlir:AffineDialect",
+#         "@llvm-project//mlir:AffineUtils",
+#         "@llvm-project//mlir:FuncDialect",
+#         "@llvm-project//mlir:Pass",
+#         "@llvm-project//mlir:Transforms",
+#     ],
+# )
 
 cc_library(
     name = "Passes",
     hdrs = ["Passes.h"],
     deps = [
         "AffineDistributeToMPI",
-        ":AffineFullUnroll",
-        ":AffineFullUnrollPatternRewrite",
+        # ":AffineFullUnroll",
+        # ":AffineFullUnrollPatternRewrite",
         ":pass_inc_gen",
     ],
 )
diff --git a/lib/Transform/Affine/Passes.h b/lib/Transform/Affine/Passes.h
@@ -2,8 +2,8 @@
 #define LIB_TRANSFORM_AFFINE_PASSES_H_
 
 #include "lib/Transform/Affine/AffineDistributeToMPI.h"
-#include "lib/Transform/Affine/AffineFullUnroll.h"
-#include "lib/Transform/Affine/AffineFullUnrollPatternRewrite.h"
+// #include "lib/Transform/Affine/AffineFullUnroll.h"
+// #include "lib/Transform/Affine/AffineFullUnrollPatternRewrite.h"
 
 namespace mlir {
 namespace tutorial {

diff --git a/lib/Transform/Affine/Passes.td b/lib/Transform/Affine/Passes.td
@@ -3,21 +3,21 @@
 
 include "mlir/Pass/PassBase.td"
 
-def AffineFullUnroll : Pass<"affine-full-unroll"> {
-  let summary = "Fully unroll all affine loops";
-  let description = [{
-    Fully unroll all affine loops.
-  }];
-  let dependentDialects = ["mlir::affine::AffineDialect"];
-}
+// def AffineFullUnroll : Pass<"affine-full-unroll"> {
+//   let summary = "Fully unroll all affine loops";
+//   let description = [{
+//     Fully unroll all affine loops.
+//   }];
+//   let dependentDialects = ["mlir::affine::AffineDialect"];
+// }
 
-def AffineFullUnrollPatternRewrite : Pass<"affine-full-unroll-rewrite"> {
-  let summary = "Fully unroll all affine loops using the pattern rewrite engine";
-  let description = [{
-    Fully unroll all affine loops using the pattern rewrite engine.
-  }];
-  let dependentDialects = ["mlir::affine::AffineDialect"];
-}
+// def AffineFullUnrollPatternRewrite : Pass<"affine-full-unroll-rewrite"> {
+//   let summary = "Fully unroll all affine loops using the pattern rewrite engine";
+//   let description = [{
+//     Fully unroll all affine loops using the pattern rewrite engine.
+//   }];
+//   let dependentDialects = ["mlir::affine::AffineDialect"];
+// }
 
 def AffineDistributeToMPI : Pass<"affine-distribute-to-mpi"> {
   let summary = "Distribute affine loops to MPI processes";
@@ -28,6 +28,7 @@ def AffineDistributeToMPI : Pass<"affine-distribute-to-mpi"> {
     "mlir::affine::AffineDialect",
     "mlir::scf::SCFDialect",
     "mlir::mpi::MPIDialect",
+    "mlir::memref::MemRefDialect",
   ];
 }