[MLIR][Affine] Loop fusion in a block containing Linalg op

[aturetsk]

Handle region-holding operators implementing Linalg interface in MemRefDependenceGraph (RegionBranchOpInterface has been removed from Linalg in https://reviews.llvm.org/D155841). This allows to fuse affine loops if a Linalg operator is present in a block.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-affine

Author: Andrey Turetskiy (aturetsk)

Changes

Handle region-holding operators implementing Linalg interface in MemRefDependenceGraph (RegionBranchOpInterface has been removed from Linalg in https://reviews.llvm.org/D155841). This allows to fuse affine loops if a Linalg operator is present in a block.

---

Full diff: https://github.com/llvm/llvm-project/pull/129136.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/Affine/Analysis/Utils.cpp (+6-3)
• (modified) mlir/test/Dialect/Affine/loop-fusion-4.mlir (+30)

diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index ba6f045cff408..425b58a3a86b1 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -19,6 +19,7 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/Interfaces/CallInterfaces.h"
@@ -252,6 +253,9 @@ bool MemRefDependenceGraph::init() {
   // Create graph nodes.
   DenseMap<Operation *, unsigned> forToNodeMap;
   for (Operation &op : block) {
+    bool hasUnhandleableRegion = op.getNumRegions() != 0 &&
+                                 !isa<RegionBranchOpInterface>(op) &&
+                                 !isa<linalg::LinalgOp>(op);
     if (auto forOp = dyn_cast<AffineForOp>(op)) {
       Node *node = addNodeToMDG(&op, *this, memrefAccesses);
       if (!node)
@@ -277,8 +281,7 @@ bool MemRefDependenceGraph::init() {
       Node *node = addNodeToMDG(&op, *this, memrefAccesses);
       if (!node)
         return false;
-    } else if (!isMemoryEffectFree(&op) &&
-               (op.getNumRegions() == 0 || isa<RegionBranchOpInterface>(op))) {
+    } else if (!isMemoryEffectFree(&op) && !hasUnhandleableRegion) {
       // Create graph node for top-level op unless it is known to be
       // memory-effect free. This covers all unknown/unregistered ops,
       // non-affine ops with memory effects, and region-holding ops with a
@@ -287,7 +290,7 @@ bool MemRefDependenceGraph::init() {
       Node *node = addNodeToMDG(&op, *this, memrefAccesses);
       if (!node)
         return false;
-    } else if (op.getNumRegions() != 0 && !isa<RegionBranchOpInterface>(op)) {
+    } else if (hasUnhandleableRegion) {
       // Return false if non-handled/unknown region-holding ops are found. We
       // won't know what such ops do or what its regions mean; for e.g., it may
       // not be an imperative op.
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index b5b951ad5eb0e..69c7c250404e1 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -548,6 +548,36 @@ func.func @sibling_reduction(%input : memref<10xf32>, %output : memref<10xf32>,
 
 // -----
 
+// Check that presence of a Linalg operator in a block does not prevent
+// fusion from happening in this block.
+
+// ALL-LABEL: func @fusion_in_block_containing_linalg
+func.func @fusion_in_block_containing_linalg(%arg0: memref<5xi8>, %arg1: memref<5xi8>) {
+  %c15_i8 = arith.constant 15 : i8
+  %alloc = memref.alloc() : memref<5xi8>
+  affine.for %arg3 = 0 to 5 {
+    affine.store %c15_i8, %alloc[%arg3] : memref<5xi8>
+  }
+  affine.for %arg3 = 0 to 5 {
+    %0 = affine.load %alloc[%arg3] : memref<5xi8>
+    %1 = affine.load %arg0[%arg3] : memref<5xi8>
+    %2 = arith.muli %0, %1 : i8
+    affine.store %2, %alloc[%arg3] : memref<5xi8>
+  }
+  // ALL:       affine.for
+  // ALL-NEXT:    affine.store
+  // ALL-NEXT:    affine.load
+  // ALL-NEXT:    affine.load
+  // ALL-NEXT:    arith.muli
+  // ALL-NEXT:    affine.store
+  // ALL-NEXT:  }
+  linalg.elemwise_binary ins(%alloc, %alloc: memref<5xi8>, memref<5xi8>) outs(%arg1: memref<5xi8>)
+  // ALL-NEXT:  linalg.elemwise_binary
+  return
+}
+
+// -----
+
 // From  https://github.com/llvm/llvm-project/issues/54541
 
 #map = affine_map<(d0) -> (d0 mod 65536)>

[aturetsk]

@kuhar, fixed, thanks!

[kuhar]

Looks good to me but please wait for someone who has stakes in the affine dialect to approve before merging.