[Mlir-commits] [mlir] [mlir][IntegerRangeAnalysis] Fix crash when loop bounds are unavailable (PR #180330)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-arith

@llvm/pr-subscribers-mlir

Author: puneeth_aditya_5656 (mugiwaraluffy56)

<details>
<summary>Changes</summary>

Some operations implement `LoopLikeOpInterface` but return `nullopt` for `getLoopLowerBounds()`, `getLoopUpperBounds()`, or `getLoopSteps()` (e.g., `tensor.generate`). The code was dereferencing these optionals without checking if they have values, causing a crash with the `--arith-int-range-narrowing` pass.

Add null checks for all three optional results and fall back to the parent class implementation when bounds information is unavailable.

Fixes #<!-- -->180312 

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


2 Files Affected:

- (modified) mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp (+13-4) 
- (added) mlir/test/Dialect/Arith/int-range-narrowing-no-crash.mlir (+28) 


``````````diff
diff --git a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
index 7b567f043577a..ac7e0acb3b58a 100644
--- a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
@@ -217,10 +217,19 @@ void IntegerRangeAnalysis::visitNonControlFlowArguments(
       return SparseForwardDataFlowAnalysis ::visitNonControlFlowArguments(
           op, successor, nonSuccessorInputs, nonSuccessorInputLattices);
     }
-    // This shouldn't be returning nullopt if there are indunction variables.
-    SmallVector<OpFoldResult> lowerBounds = *loop.getLoopLowerBounds();
-    SmallVector<OpFoldResult> upperBounds = *loop.getLoopUpperBounds();
-    SmallVector<OpFoldResult> steps = *loop.getLoopSteps();
+    std::optional<SmallVector<OpFoldResult>> maybeLowerBounds =
+        loop.getLoopLowerBounds();
+    std::optional<SmallVector<OpFoldResult>> maybeUpperBounds =
+        loop.getLoopUpperBounds();
+    std::optional<SmallVector<OpFoldResult>> maybeSteps = loop.getLoopSteps();
+    // Not all loops provide bounds information (e.g., some region-based ops).
+    if (!maybeLowerBounds || !maybeUpperBounds || !maybeSteps) {
+      return SparseForwardDataFlowAnalysis::visitNonControlFlowArguments(
+          op, successor, nonSuccessorInputs, nonSuccessorInputLattices);
+    }
+    SmallVector<OpFoldResult> &lowerBounds = *maybeLowerBounds;
+    SmallVector<OpFoldResult> &upperBounds = *maybeUpperBounds;
+    SmallVector<OpFoldResult> &steps = *maybeSteps;
     for (auto [iv, lowerBound, upperBound, step] :
          llvm::zip_equal(*maybeIvs, lowerBounds, upperBounds, steps)) {
       Block *block = iv.getParentBlock();
diff --git a/mlir/test/Dialect/Arith/int-range-narrowing-no-crash.mlir b/mlir/test/Dialect/Arith/int-range-narrowing-no-crash.mlir
new file mode 100644
index 0000000000000..eddbe21958129
--- /dev/null
+++ b/mlir/test/Dialect/Arith/int-range-narrowing-no-crash.mlir
@@ -0,0 +1,28 @@
+// RUN: mlir-opt --arith-int-range-narrowing="int-bitwidths-supported=1,8,16,32" %s | FileCheck %s
+
+// Test that the pass doesn't crash on operations that implement
+// LoopLikeOpInterface but don't provide loop bounds (e.g., tensor.generate).
+// See https://github.com/llvm/llvm-project/issues/180312
+
+// CHECK-LABEL: func @tensor_generate_no_crash
+func.func @tensor_generate_no_crash(%arg0: index) -> tensor<?xf32> {
+  %cst = arith.constant 1.0 : f32
+  // tensor.generate implements LoopLikeOpInterface but getLoopLowerBounds(),
+  // getLoopUpperBounds(), and getLoopSteps() return nullopt.
+  %0 = tensor.generate %arg0 {
+  ^bb0(%i: index):
+    tensor.yield %cst : f32
+  } : tensor<?xf32>
+  return %0 : tensor<?xf32>
+}
+
+// CHECK-LABEL: func @tensor_generate_with_arith
+func.func @tensor_generate_with_arith(%arg0: index) -> tensor<?xindex> {
+  %c1 = arith.constant 1 : index
+  %0 = tensor.generate %arg0 {
+  ^bb0(%i: index):
+    %sum = arith.addi %i, %c1 : index
+    tensor.yield %sum : index
+  } : tensor<?xindex>
+  return %0 : tensor<?xindex>
+}

``````````

</details>


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