[MLIR][Affine] Fix getSliceBounds for missing handling of no lower/upper bound in certain cases

[bondhugula]

Fix FlatLinearValueConstraints::getSliceBounds for missing checks on no
lower/upper bound bound. Obvious bug.

Fixes: #119525
Fixes: #108374

[llvmbot]

@llvm/pr-subscribers-mlir-affine

@llvm/pr-subscribers-mlir

Author: Uday Bondhugula (bondhugula)

Changes

Fix FlatLinearValueConstraints::getSliceBounds for missing checks on no
lower/upper bound bound. Obvious bug.

Fixes: #119525
Fixes: #108374

---

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

2 Files Affected:

• (modified) mlir/lib/Analysis/FlatLinearValueConstraints.cpp (+2-2)
• (modified) mlir/test/Dialect/Affine/loop-fusion-4.mlir (+70)

diff --git a/mlir/lib/Analysis/FlatLinearValueConstraints.cpp b/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
index 0d6ff2fd908db..e9ada0a386723 100644
--- a/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
+++ b/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
@@ -679,7 +679,7 @@ void FlatLinearConstraints::getSliceBounds(unsigned offset, unsigned num,
       // TODO: being conservative for the moment in cases that
       // lead to multiple bounds - until getConstDifference in LoopFusion.cpp is
       // fixed (b/126426796).
-      if (!lbMap || lbMap.getNumResults() > 1) {
+      if (!lbMap || lbMap.getNumResults() != 1) {
         LLVM_DEBUG(llvm::dbgs()
                    << "WARNING: Potentially over-approximating slice lb\n");
         auto lbConst = getConstantBound64(BoundType::LB, pos + offset);
@@ -688,7 +688,7 @@ void FlatLinearConstraints::getSliceBounds(unsigned offset, unsigned num,
                                  getAffineConstantExpr(*lbConst, context));
         }
       }
-      if (!ubMap || ubMap.getNumResults() > 1) {
+      if (!ubMap || ubMap.getNumResults() != 1) {
         LLVM_DEBUG(llvm::dbgs()
                    << "WARNING: Potentially over-approximating slice ub\n");
         auto ubConst = getConstantBound64(BoundType::UB, pos + offset);
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index 788d7f9470530..3e4135898d46c 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -391,3 +391,73 @@ func.func @memref_index_type() {
   // PRODUCER-CONSUMER-MAXIMAL: return
   return
 }
+
+// -----
+
+#map = affine_map<(d0) -> (d0)>
+#map1 = affine_map<(d0) -> (d0 + 1)>
+
+// Exercises fix for crash reported at https://github.com/llvm/llvm-project/issues/119525
+
+// No fusion of  producer into consumer happens here as the slice is determined
+// to be invalid. This is a limitation and it is possible to compute a slice
+// (reduction along %arg4) and fuse.
+
+// PRODUCER-CONSUMER-LABEL: func @slice_compute_check
+func.func @slice_compute_check(%arg0: memref<1x8x26xi32, strided<[?, ?, ?], offset: ?>>, %arg1: memref<1x8x26xi32, strided<[?, ?, ?], offset: ?>>, %arg2: memref<1x8x26xi32, strided<[?, ?, ?], offset: ?>>) {
+  %alloc_14 = memref.alloc() : memref<1x8x26xi32>
+  %alloc_15 = memref.alloc() : memref<1x26xi32>
+  affine.for %arg3 = 0 to 1 {
+    affine.for %arg4 = 0 to 8 {
+      affine.for %arg5 = 0 to 26 {
+        affine.for %arg6 = #map(%arg3) to #map1(%arg3) {
+          affine.for %arg7 = #map(%arg4) to #map1(%arg4) {
+            affine.for %arg8 = #map(%arg5) to #map1(%arg5) {
+              %61 = affine.load %alloc_14[%arg6, %arg7, %arg8] : memref<1x8x26xi32>
+              %62 = affine.load %alloc_15[%arg6, %arg8] : memref<1x26xi32>
+              %63 = llvm.intr.smin(%61, %62) : (i32, i32) -> i32
+              affine.store %63, %alloc_15[%arg6, %arg8] : memref<1x26xi32>
+            }
+          }
+        }
+      }
+    }
+  }
+  affine.for %arg3 = 0 to 26 {
+    %61 = affine.load %alloc_15[0, %arg3] : memref<1x26xi32>
+  }
+  memref.dealloc %alloc_15 : memref<1x26xi32>
+  memref.dealloc %alloc_14 : memref<1x8x26xi32>
+  return
+}
+
+// -----
+
+// Exercises fix for crash reported at https://github.com/llvm/llvm-project/issues/108374
+
+// No fusion of  producer into consumer happens here. The slice will not be
+// valid as the producer doesn't supply to all of the consumer.
+
+#map = affine_map<(d0) -> (d0)>
+#map1 = affine_map<(d0) -> (d0 + 1)>
+// PRODUCER-CONSUMER-LABEL: func @test_add_slice_bounds
+func.func @test_add_slice_bounds() {
+  %alloc = memref.alloc() : memref<10xf32>
+  %cst = arith.constant 0.619152 : f32
+  affine.for %arg0 = 0 to 10 {
+    affine.for %arg1 = #map(%arg0) to #map1(%arg0) {
+      affine.store %cst, %alloc[%arg1] : memref<10xf32>
+    }
+  }
+  affine.for %arg0 = 0 to 3 {
+    affine.for %arg1 = 0 to 10 {
+      affine.for %arg2 = #map(%arg0) to #map1(%arg0) {
+        affine.for %arg3 = #map(%arg1) to #map1(%arg1) {
+          %0 = affine.apply #map1(%arg3)
+          %1 = affine.load %alloc[%0] : memref<10xf32>
+        }
+      }
+    }
+  }
+  return
+}

[bondhugula]

This fixes #127192 as well!