[Mlir-commits] [mlir] [mlir][affine] Treat failed dependence checks conservatively (PR #211014)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Takayuki Todokoro (takatodo)

<details>
<summary>Changes</summary>

Affine dependence analysis can return `DependenceResult::Failure` when it
cannot construct an access relation. Some loop transformations currently
treat this result like `NoDependence`, which can allow an invalid
transformation and cause a miscompilation.

This change handles analysis failures conservatively:

- Reject loop tiling when a dependence check fails.
- Reject validity-checked loop permutation when the complete dependence
  graph cannot be constructed.
- Preserve the original loop nest when sequential-loop sinking encounters
  the same failure.
- Add regression tests using a semi-affine same-memref dependence.

Fixes #<!-- -->210585.

Assisted-by: OpenAI Codex

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


5 Files Affected:

- (modified) mlir/include/mlir/Dialect/Affine/Analysis/AffineAnalysis.h (+2-2) 
- (modified) mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp (+4-1) 
- (modified) mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp (+4-1) 
- (modified) mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp (+4-2) 
- (added) mlir/test/Dialect/Affine/loop-transformation-validity.mlir (+25) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/AffineAnalysis.h b/mlir/include/mlir/Dialect/Affine/Analysis/AffineAnalysis.h
index 3e4b8648061ff..6bfa887e95fd1 100644
--- a/mlir/include/mlir/Dialect/Affine/Analysis/AffineAnalysis.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/AffineAnalysis.h
@@ -190,8 +190,8 @@ inline bool noDependence(DependenceResult result) {
 
 /// Returns in 'depCompsVec', dependence components for dependences between all
 /// load and store ops in loop nest rooted at 'forOp', at loop depths in range
-/// [1, maxLoopDepth].
-void getDependenceComponents(
+/// [1, maxLoopDepth]. Returns failure if any dependence cannot be analyzed.
+LogicalResult getDependenceComponents(
     AffineForOp forOp, unsigned maxLoopDepth,
     std::vector<SmallVector<DependenceComponent, 2>> *depCompsVec);
 
diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
index 3d1a73417d1ea..aef5ec85858a7 100644
--- a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
@@ -694,7 +694,7 @@ DependenceResult mlir::affine::checkMemrefAccessDependence(
 
 /// Gathers dependence components for dependences between all ops in loop nest
 /// rooted at 'forOp' at loop depths in range [1, maxLoopDepth].
-void mlir::affine::getDependenceComponents(
+LogicalResult mlir::affine::getDependenceComponents(
     AffineForOp forOp, unsigned maxLoopDepth,
     std::vector<SmallVector<DependenceComponent, 2>> *depCompsVec) {
   // Collect all load and store ops in loop nest rooted at 'forOp'.
@@ -719,9 +719,12 @@ void mlir::affine::getDependenceComponents(
         DependenceResult result = checkMemrefAccessDependence(
             srcAccess, dstAccess, d, /*dependenceConstraints=*/nullptr,
             &depComps);
+        if (result.value == DependenceResult::Failure)
+          return failure();
         if (hasDependence(result))
           depCompsVec->push_back(depComps);
       }
     }
   }
+  return success();
 }
diff --git a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
index 40802cc6e85e5..2dbc1baa02236 100644
--- a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
@@ -524,8 +524,11 @@ bool mlir::affine::isTilingValid(ArrayRef<AffineForOp> loops) {
             srcAccess, dstAccess, d, /*dependenceConstraints=*/nullptr,
             &depComps);
 
+        if (result.value == DependenceResult::Failure)
+          return false;
+
         // Skip if there is no dependence in this case.
-        if (!hasDependence(result))
+        if (noDependence(result))
           continue;
 
         // Check whether there is any negative direction vector in the
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index 90bc57e950cf1..bb002229d7b50 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -1357,7 +1357,8 @@ bool mlir::affine::isValidLoopInterchangePermutation(
   // Gather dependence components for dependences between all ops in loop nest
   // rooted at 'loops[0]', at loop depths in range [1, maxLoopDepth].
   std::vector<SmallVector<DependenceComponent, 2>> depCompsVec;
-  getDependenceComponents(loops[0], maxLoopDepth, &depCompsVec);
+  if (failed(getDependenceComponents(loops[0], maxLoopDepth, &depCompsVec)))
+    return false;
   return checkLoopInterchangeDependences(depCompsVec, loops, loopPermMap);
 }
 
@@ -1466,7 +1467,8 @@ AffineForOp mlir::affine::sinkSequentialLoops(AffineForOp forOp) {
   // rooted at 'loops[0]', at loop depths in range [1, maxLoopDepth].
   unsigned maxLoopDepth = loops.size();
   std::vector<SmallVector<DependenceComponent, 2>> depCompsVec;
-  getDependenceComponents(loops[0], maxLoopDepth, &depCompsVec);
+  if (failed(getDependenceComponents(loops[0], maxLoopDepth, &depCompsVec)))
+    return forOp;
 
   // Mark loops as either parallel or sequential.
   SmallVector<bool, 8> isParallelLoop(maxLoopDepth, true);
diff --git a/mlir/test/Dialect/Affine/loop-transformation-validity.mlir b/mlir/test/Dialect/Affine/loop-transformation-validity.mlir
new file mode 100644
index 0000000000000..b57e9ddcc357b
--- /dev/null
+++ b/mlir/test/Dialect/Affine/loop-transformation-validity.mlir
@@ -0,0 +1,25 @@
+// RUN: mlir-opt %s -test-loop-permutation="permutation-map=1,0 check-validity=1" | FileCheck %s
+// RUN: mlir-opt %s -affine-loop-tile="tile-size=4" | FileCheck %s
+
+#dynamic_index = affine_map<()[s0, s1] -> (s0 * s1)>
+
+// Dependence analysis cannot represent the common semi-affine index. The
+// remaining indices carry a (1, -1) dependence, so both transforms must fail.
+// CHECK-LABEL: func.func @unknown_dependence
+func.func @unknown_dependence(
+    %A: memref<?x9x9xi32>, %B: memref<9x9xi32>,
+    %p: index, %q: index, %value: i32) {
+  // CHECK:      affine.for %[[I:.*]] = 1 to 8 {
+  // CHECK-NEXT:   affine.for %[[J:.*]] = 1 to 8 {
+  affine.for %i = 1 to 8 {
+    affine.for %j = 1 to 8 {
+      %z = affine.apply #dynamic_index()[%p, %q]
+      // CHECK: affine.store %{{.*}}, %{{.*}}[%{{.*}}, %[[I]], %[[J]]]
+      affine.store %value, %A[%z, %i, %j] : memref<?x9x9xi32>
+      %loaded = affine.load %A[%z, %i - 1, %j + 1]
+          : memref<?x9x9xi32>
+      affine.store %loaded, %B[%i, %j] : memref<9x9xi32>
+    }
+  }
+  return
+}

``````````

</details>


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