[MLIR][IntegerRangeAnalysis] Avoid crash reached when loop bound is u…

…ninitialized (llvm#74832) If the loop bound is not initialized, the analysis crashed, as it only checked for nullity. Also checking for initialization fixes the issue. Signed-off-by: Victor Perez <[email protected]> Co-authored-by: Tsang, Whitney <[email protected]>
aniragil · Dec 11, 2023 · 13c648f · 13c648f
1 parent 75193b1
commit 13c648f
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 1 deletion.
diff --git a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
@@ -180,7 +180,7 @@ void IntegerRangeAnalysis::visitNonControlFlowArguments(
       } else if (auto value = llvm::dyn_cast_if_present<Value>(*loopBound)) {
         const IntegerValueRangeLattice *lattice =
             getLatticeElementFor(op, value);
-        if (lattice != nullptr)
+        if (lattice != nullptr && !lattice->getValue().isUninitialized())
           return getUpper ? lattice->getValue().getValue().smax()
                           : lattice->getValue().getValue().smin();
       }

diff --git a/mlir/test/Dialect/Arith/int-range-interface.mlir b/mlir/test/Dialect/Arith/int-range-interface.mlir
@@ -730,3 +730,29 @@ func.func @extui_uses_unsigned(%arg0 : i32) -> i1 {
     %4 = arith.andi %2, %3 : i1
     func.return %4 : i1
 }
+
+/// Catch a bug that caused a crash in getLoopBoundFromFold when
+/// SparseConstantPropagation is loaded in the solver.
+
+// CHECK-LABEL:   func.func @caller(
+// CHECK-SAME:                      %[[VAL_0:.*]]: memref<?xindex, 4>) {
+// CHECK:           call @callee(%[[VAL_0]]) : (memref<?xindex, 4>) -> ()
+// CHECK:           return
+// CHECK:         }
+func.func @caller(%arg0: memref<?xindex, 4>) {
+  call @callee(%arg0) : (memref<?xindex, 4>) -> ()
+  return
+}
+
+// CHECK-LABEL:   func.func private @callee(
+// CHECK-SAME:                              %[[VAL_0:.*]]: memref<?xindex, 4>) {
+// CHECK:           return
+// CHECK:         }
+func.func private @callee(%arg0: memref<?xindex, 4>) {
+  %c1 = arith.constant 1 : index
+  %c0 = arith.constant 0 : index
+  %0 = affine.load %arg0[0] : memref<?xindex, 4>
+  scf.for %arg1 = %c0 to %0 step %c1 {
+  }
+  return
+}
diff --git a/mlir/test/lib/Transforms/TestIntRangeInference.cpp b/mlir/test/lib/Transforms/TestIntRangeInference.cpp
@@ -9,6 +9,7 @@
 // functionality has been integrated into SCCP.
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h"
 #include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
 #include "mlir/Analysis/DataFlow/IntegerRangeAnalysis.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
@@ -107,6 +108,7 @@ struct TestIntRangeInference
     Operation *op = getOperation();
     DataFlowSolver solver;
     solver.load<DeadCodeAnalysis>();
+    solver.load<SparseConstantPropagation>();
     solver.load<IntegerRangeAnalysis>();
     if (failed(solver.initializeAndRun(op)))
       return signalPassFailure();