[Mlir-commits] [mlir] [MLIR][Affine] Generalize/fix mdg init for region-holding ops with known control flow (PR #71754)

[Uday Bondhugula]

https://github.com/bondhugula created https://github.com/llvm/llvm-project/pull/71754

Generalize/fix mdg init for region-holding ops with well-defined control
flow. Use the memory effect interface to determine when to create a
node. While on this, remove the special treatment there for call ops.

This allows fusion of affine nests even in the presence of scf region
ops elsewhere in the Block. Previously, even a single scf.for/if/while
op in a block would have made fusion bail out on all affine fusion that
was possible. Addressed.


>From e49e2b8b2be0f36ab66f000ac9e6116bb48a2c61 Mon Sep 17 00:00:00 2001
From: Uday Bondhugula <uday at polymagelabs.com>
Date: Tue, 7 Nov 2023 16:51:36 +0530
Subject: [PATCH] [MLIR][Affine] Generalize/fix mdg init for region-holding ops
 with known control flow

Generalize/fix mdg init for region-holding ops with well-defined control
flow. Use the memory effect interface to determine when to create a
node. While on this, remove the special treatment there for call ops.

This allows fusion of affine nests even in the presence of scf region
ops elsewhere in the Block. Previously, even a single scf.for/if/while
op in a block would have made fusion bail out on all affine fusion that
was possible. Addressed.
---
 mlir/lib/Dialect/Affine/Analysis/Utils.cpp    | 30 +++++------
 .../Dialect/Affine/loop-fusion-scf-mixed.mlir | 51 +++++++++++++++++++
 2 files changed, 66 insertions(+), 15 deletions(-)
 create mode 100644 mlir/test/Dialect/Affine/loop-fusion-scf-mixed.mlir

diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 23921b700b6669b..9d2998456a4b679 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -152,28 +152,28 @@ bool MemRefDependenceGraph::init() {
       auto memref = cast<AffineWriteOpInterface>(op).getMemRef();
       memrefAccesses[memref].insert(node.id);
       nodes.insert({node.id, node});
-    } else if (op.getNumRegions() != 0) {
-      // Return false if another region is found (not currently supported).
-      return false;
     } else if (op.getNumResults() > 0 && !op.use_empty()) {
       // Create graph node for top-level producer of SSA values, which
       // could be used by loop nest nodes.
       Node node(nextNodeId++, &op);
       nodes.insert({node.id, node});
-    } else if (isa<CallOpInterface>(op)) {
-      // Create graph node for top-level Call Op that takes any argument of
-      // memref type. Call Op that returns one or more memref type results
-      // is already taken care of, by the previous conditions.
-      if (llvm::any_of(op.getOperandTypes(),
-                       [&](Type t) { return isa<MemRefType>(t); })) {
-        Node node(nextNodeId++, &op);
-        nodes.insert({node.id, node});
-      }
-    } else if (hasEffect<MemoryEffects::Write, MemoryEffects::Free>(&op)) {
-      // Create graph node for top-level op, which could have a memory write
-      // side effect.
+    } else if (!isMemoryEffectFree(&op) &&
+               (op.getNumRegions() == 0 || isa<RegionBranchOpInterface>(op))) {
+      // 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
+      // well-defined control flow. During the fusion validity checks, we look
+      // for non-affine ops on the path from source to destination, at which
+      // point we check which memrefs if any are used in the region.
       Node node(nextNodeId++, &op);
       nodes.insert({node.id, node});
+    } else if (op.getNumRegions() != 0) {
+      // 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.
+      LLVM_DEBUG(llvm::dbgs()
+                 << "MDG init failed; unknown region-holding op found!\n");
+      return false;
     }
   }
 
diff --git a/mlir/test/Dialect/Affine/loop-fusion-scf-mixed.mlir b/mlir/test/Dialect/Affine/loop-fusion-scf-mixed.mlir
new file mode 100644
index 000000000000000..cf67b71ff4a5490
--- /dev/null
+++ b/mlir/test/Dialect/Affine/loop-fusion-scf-mixed.mlir
@@ -0,0 +1,51 @@
+// RUN: mlir-opt -pass-pipeline='builtin.module(func.func(affine-loop-fusion))' %s | FileCheck %s
+
+// Test fusion of affine nests in the presence of other region-holding ops
+// (scf.for in the test case below) in the block.
+
+// CHECK-LABEL: func @scf_and_affine
+func.func @scf_and_affine(%A : memref<10xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c10 = arith.constant 10 : index
+  %cst = arith.constant 0.0 : f32
+
+  %B = memref.alloc() : memref<10xf32>
+  %C = memref.alloc() : memref<10xf32>
+
+  affine.for %j = 0 to 10 {
+    %v = affine.load %A[%j] : memref<10xf32>
+    affine.store %v, %B[%j] : memref<10xf32>
+  }
+
+  affine.for %j = 0 to 10 {
+    %v = affine.load %B[%j] : memref<10xf32>
+    affine.store %v, %C[%j] : memref<10xf32>
+  }
+  // Nests are fused.
+  // CHECK:     affine.for %{{.*}} = 0 to 10
+  // CHECK-NOT: affine.for
+  // CHECK:     scf.for
+
+  scf.for %i = %c0 to %c10 step %c1 {
+    memref.store %cst, %B[%i] : memref<10xf32>
+  }
+
+  // The nests below shouldn't be fused.
+  affine.for %j = 0 to 10 {
+    %v = affine.load %A[%j] : memref<10xf32>
+    affine.store %v, %B[%j] : memref<10xf32>
+  }
+  scf.for %i = %c0 to %c10 step %c1 {
+    memref.store %cst, %B[%i] : memref<10xf32>
+  }
+  affine.for %j = 0 to 10 {
+    %v = affine.load %B[%j] : memref<10xf32>
+    affine.store %v, %C[%j] : memref<10xf32>
+  }
+  // CHECK: affine.for
+  // CHECK: scf.for
+  // CHECK: affine.for
+
+  return
+}