[PATCH] D79558: [MLIR] Make ParallelLoopFusion pass scan through all nested regions.

[Alexander Belyaev via Phabricator via llvm-commits]

pifon2a created this revision.
pifon2a added a reviewer: herhut.
Herald added subscribers: llvm-commits, Kayjukh, frgossen, grosul1, Joonsoo, stephenneuendorffer, liufengdb, lucyrfox, mgester, arpith-jacob, nicolasvasilache, antiagainst, shauheen, jpienaar, rriddle, mehdi_amini.
Herald added a reviewer: nicolasvasilache.
Herald added a project: LLVM.

Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D79558

Files:
  mlir/lib/Dialect/LoopOps/Transforms/ParallelLoopFusion.cpp
  mlir/test/Dialect/Loops/parallel-loop-fusion.mlir


Index: mlir/test/Dialect/Loops/parallel-loop-fusion.mlir
===================================================================
--- mlir/test/Dialect/Loops/parallel-loop-fusion.mlir
+++ mlir/test/Dialect/Loops/parallel-loop-fusion.mlir
@@ -307,3 +307,53 @@
 // CHECK-LABEL: func @do_not_fuse_loops_with_memref_defined_in_loop_bodies
 // CHECK:        loop.parallel
 // CHECK:        loop.parallel
+
+// -----
+
+func @nested_fuse(%A: memref<2x2xf32>, %B: memref<2x2xf32>,
+                    %C: memref<2x2xf32>, %result: memref<2x2xf32>) {
+  %c2 = constant 2 : index
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %sum = alloc()  : memref<2x2xf32>
+  loop.parallel (%k) = (%c0) to (%c2) step (%c1) {
+    loop.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
+      %B_elem = load %B[%i, %j] : memref<2x2xf32>
+      %C_elem = load %C[%i, %j] : memref<2x2xf32>
+      %sum_elem = addf %B_elem, %C_elem : f32
+      store %sum_elem, %sum[%i, %j] : memref<2x2xf32>
+      loop.yield
+    }
+    loop.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
+      %sum_elem = load %sum[%i, %j] : memref<2x2xf32>
+      %A_elem = load %A[%i, %j] : memref<2x2xf32>
+      %product_elem = mulf %sum_elem, %A_elem : f32
+      store %product_elem, %result[%i, %j] : memref<2x2xf32>
+      loop.yield
+    }
+  }
+  dealloc %sum : memref<2x2xf32>
+  return
+}
+// CHECK-LABEL: func @nested_fuse
+// CHECK-SAME:   ([[A:%.*]]: {{.*}}, [[B:%.*]]: {{.*}}, [[C:%.*]]: {{.*}},
+// CHECK-SAME:    [[RESULT:%.*]]: {{.*}}) {
+// CHECK:      [[C2:%.*]] = constant 2 : index
+// CHECK:      [[C0:%.*]] = constant 0 : index
+// CHECK:      [[C1:%.*]] = constant 1 : index
+// CHECK:      [[SUM:%.*]] = alloc()
+// CHECK:      loop.parallel
+// CHECK:        loop.parallel ([[I:%.*]], [[J:%.*]]) = ([[C0]], [[C0]])
+// CHECK-SAME:       to ([[C2]], [[C2]]) step ([[C1]], [[C1]]) {
+// CHECK:          [[B_ELEM:%.*]] = load [[B]]{{\[}}[[I]], [[J]]]
+// CHECK:          [[C_ELEM:%.*]] = load [[C]]{{\[}}[[I]], [[J]]]
+// CHECK:          [[SUM_ELEM:%.*]] = addf [[B_ELEM]], [[C_ELEM]]
+// CHECK:          store [[SUM_ELEM]], [[SUM]]{{\[}}[[I]], [[J]]]
+// CHECK:          [[SUM_ELEM_:%.*]] = load [[SUM]]{{\[}}[[I]], [[J]]]
+// CHECK:          [[A_ELEM:%.*]] = load [[A]]{{\[}}[[I]], [[J]]]
+// CHECK:          [[PRODUCT_ELEM:%.*]] = mulf [[SUM_ELEM_]], [[A_ELEM]]
+// CHECK:          store [[PRODUCT_ELEM]], [[RESULT]]{{\[}}[[I]], [[J]]]
+// CHECK:          loop.yield
+// CHECK:        }
+// CHECK:      }
+// CHECK:      dealloc [[SUM]]
Index: mlir/lib/Dialect/LoopOps/Transforms/ParallelLoopFusion.cpp
===================================================================
--- mlir/lib/Dialect/LoopOps/Transforms/ParallelLoopFusion.cpp
+++ mlir/lib/Dialect/LoopOps/Transforms/ParallelLoopFusion.cpp
@@ -162,11 +162,12 @@
 struct ParallelLoopFusion
     : public LoopParallelLoopFusionBase<ParallelLoopFusion> {
   void runOnOperation() override {
-    for (Region &region : getOperation()->getRegions())
-      naivelyFuseParallelOps(region);
+    getOperation()->walk([&](Operation *child) {
+      for (Region &region : child->getRegions())
+        naivelyFuseParallelOps(region);
+    });
   }
 };
-
 } // namespace
 
 std::unique_ptr<Pass> mlir::createParallelLoopFusionPass() {


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D79558.262567.patch
Type: text/x-patch
Size: 3294 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200507/07b20316/attachment.bin>