[Mlir-commits] [mlir] [MLIR][Affine] Fix getSliceBounds for missing handling of no lower/upper bound in certain cases (PR #127192)

[Uday Bondhugula]

https://github.com/bondhugula updated https://github.com/llvm/llvm-project/pull/127192

>From 0a714d469b4946af26ebddfee2b6e4a4b06e853e Mon Sep 17 00:00:00 2001
From: Uday Bondhugula <uday at polymagelabs.com>
Date: Fri, 14 Feb 2025 11:46:26 +0530
Subject: [PATCH] [MLIR][Affine] Fix getSliceBounds for missing handling of no
 lower/upper bound in certain cases

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

Fixes: https://github.com/llvm/llvm-project/issues/119525
Fixes: https://github.com/llvm/llvm-project/issues/108374
---
 .../Analysis/FlatLinearValueConstraints.cpp   |  4 +-
 mlir/test/Dialect/Affine/loop-fusion-4.mlir   | 70 +++++++++++++++++++
 2 files changed, 72 insertions(+), 2 deletions(-)

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 3a7cbcfecf2eb..9cc2dc1520623 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -425,3 +425,73 @@ func.func @non_int_memory_space() {
   // PRODUCER-CONSUMER-NEXT:   affine.for %{{.*}} = 0 to 8
   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
+}