[Mlir-commits] [mlir] [MLIR] control-flow-sink: don't sink non-pure ops with affine-valid results (PR #218216)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Peruere1828 (Peruere1828)

<details>
<summary>Changes</summary>

A non-pure operation whose result is currently a valid affine dimension or symbol must stay at the top level of its affine scope to remain valid. Sinking such an operation into a conditionally executed region will break that invariant and produce invalid IR (e.g. `index.remu` used by `affine.load`, see issue #<!-- -->216542).

The previous `isMemoryEffectFree` check was too permissive because `index.remu` is memory-effect-free but not pure. Switching the legality predicate to `isPure` was too aggressive and broke sinking of ops such as `test.region_if` whose results are not affine-relevant. Only reject sinking when the operation's result is actually a valid affine dim/symbol. Also adds MLIRAffineDialect to MLIRTransforms LINK_LIBS (required for BUILD_SHARED_LIBS builds).

Fixes: #<!-- -->216542

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


3 Files Affected:

- (modified) mlir/lib/Transforms/CMakeLists.txt (+1) 
- (modified) mlir/lib/Transforms/ControlFlowSink.cpp (+18-2) 
- (modified) mlir/test/Transforms/control-flow-sink.mlir (+32) 


``````````diff
diff --git a/mlir/lib/Transforms/CMakeLists.txt b/mlir/lib/Transforms/CMakeLists.txt
index 66b39f53c91df..bf31554526730 100644
--- a/mlir/lib/Transforms/CMakeLists.txt
+++ b/mlir/lib/Transforms/CMakeLists.txt
@@ -31,6 +31,7 @@ add_mlir_library(MLIRTransforms
   MLIRTransformsDialectInterfaceIncGen
 
   LINK_LIBS PUBLIC
+  MLIRAffineDialect
   MLIRAnalysis
   MLIRFunctionInterfaces
   MLIRLoopLikeInterface
diff --git a/mlir/lib/Transforms/ControlFlowSink.cpp b/mlir/lib/Transforms/ControlFlowSink.cpp
index eed9c1cdab36d..a093ee9fa4ec4 100644
--- a/mlir/lib/Transforms/ControlFlowSink.cpp
+++ b/mlir/lib/Transforms/ControlFlowSink.cpp
@@ -15,6 +15,7 @@
 
 #include "mlir/Transforms/Passes.h"
 
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
@@ -40,10 +41,25 @@ void ControlFlowSink::runOnOperation() {
     SmallVector<Region *> regionsToSink;
     // Get the regions are that known to be executed at most once.
     getSinglyExecutedRegionsToSink(branch, regionsToSink);
-    // Sink side-effect free operations.
+    // Sink memory-effect-free operations.
     numSunk = controlFlowSink(
         regionsToSink, domInfo,
-        [](Operation *op, Region *) { return isMemoryEffectFree(op); },
+        [](Operation *op, Region *) {
+          if (!isMemoryEffectFree(op))
+            return false;
+          // A non-pure operation whose result is currently a valid affine
+          // dimension or symbol must remain at the top level of its affine
+          // scope to stay valid. Sinking such an operation into a
+          // conditionally executed region moves its definition into a nested
+          // region, breaking that invariant and producing invalid IR (e.g.
+          // `index.remu` used by `affine.load`).
+          if (isPure(op))
+            return true;
+          for (Value result : op->getResults())
+            if (affine::isValidDim(result) || affine::isValidSymbol(result))
+              return false;
+          return true;
+        },
         [](Operation *op, Region *region) {
           // Move the operation to the beginning of the region's entry block.
           // This guarantees the preservation of SSA dominance of all of the
diff --git a/mlir/test/Transforms/control-flow-sink.mlir b/mlir/test/Transforms/control-flow-sink.mlir
index 4efcac6cdbc29..eafbe2f05277d 100644
--- a/mlir/test/Transforms/control-flow-sink.mlir
+++ b/mlir/test/Transforms/control-flow-sink.mlir
@@ -235,3 +235,35 @@ func.func @test_not_sunk_deeply(%arg0: i32) -> i32 {
   }) : () -> i32
   return %1 : i32
 }
+
+// Test that a non-pure operation whose result is used as an affine dimension
+// is not sunk into a conditionally executed region. `index.remu` is
+// memory-effect-free but not pure, so moving it into the `scf.if` would make
+// its result an invalid affine dimension operand (`affine.load` requires its
+// index to be defined at the top level of the affine scope or produced by a
+// pure op).
+
+// CHECK-LABEL: func.func @test_not_sunk_non_pure(
+// CHECK-SAME:    %[[ARG0:.*]]: memref<8xi32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: i1) -> i32 {
+// CHECK-NEXT:  %[[V0:.*]] = index.constant 8
+// CHECK-NEXT:  %[[V1:.*]] = index.remu %[[ARG1]], %[[V0]]
+// CHECK-NEXT:  %[[V2:.*]] = scf.if %[[ARG2]] -> (i32) {
+// CHECK-NEXT:    %[[V3:.*]] = affine.load %[[ARG0]][%[[V1]]] : memref<8xi32>
+// CHECK-NEXT:    scf.yield %[[V3]] : i32
+// CHECK-NEXT:  } else {
+// CHECK-NEXT:    %[[V4:.*]] = index.castu %[[ARG1]] : index to i32
+// CHECK-NEXT:    scf.yield %[[V4]] : i32
+// CHECK-NEXT:  }
+// CHECK-NEXT:  return %[[V2]] : i32
+func.func @test_not_sunk_non_pure(%arg0: memref<8xi32>, %arg1: index, %arg2: i1) -> i32 {
+  %0 = index.constant 8
+  %1 = index.remu %arg1, %0
+  %2 = scf.if %arg2 -> (i32) {
+    %3 = affine.load %arg0[%1] : memref<8xi32>
+    scf.yield %3 : i32
+  } else {
+    %4 = index.castu %arg1 : index to i32
+    scf.yield %4 : i32
+  }
+  return %2 : i32
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/218216