[Mlir-commits] [mlir] [MLIR][IntegerRangeAnalysis] Avoid crash reached when loop bound is uninitialized (PR #74832)

[Victor Perez]

https://github.com/victor-eds created https://github.com/llvm/llvm-project/pull/74832

If the loop bound is not initialized, the analysis crashed, as it only checks for nullity. Also checking for initialization fixes the issue.

>From 538668ab40127bf00b9e89aa3c07a551ed838fd2 Mon Sep 17 00:00:00 2001
From: Victor Perez <victor.perez at codeplay.com>
Date: Fri, 8 Dec 2023 12:17:56 +0000
Subject: [PATCH] [MLIR][IntegerRangeAnalysis] Avoid crash reached when loop
 bound is uninitialized

If the loop bound is not initialized, the analysis crashed, as it only
checks for nullity. Also checking for initialization fixes the issue.

Co-authored-by: Tsang, Whitney <whitney.tsang at intel.com>
Signed-off-by: Victor Perez <victor.perez at codeplay.com>
---
 .../DataFlow/IntegerRangeAnalysis.cpp         |  2 +-
 .../Dialect/Arith/int-range-interface.mlir    | 26 +++++++++++++++++++
 .../lib/Transforms/TestIntRangeInference.cpp  |  2 ++
 3 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
index 39b3a0996396ff..a82c30717e275b 100644
--- 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
index 4c5f4095225d3a..02a9827d19d8f8 100644
--- 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
index d1978b6099f045..2f6dd5b8095dfa 100644
--- 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();