[Mlir-commits] [mlir] [mlir][arith] Fix crash in ConstantOp range inference for zero-element constants (PR #204180)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: shubhamnarlawar

<details>
<summary>Changes</summary>

arith::ConstantOp::inferResultRanges computes the union of element ranges by iterating a DenseIntElementsAttr. For a zero-element constant (e.g. ) the loop body never runs,
leaving the std::optional result unset.

-> This fixes #<!-- -->202531 

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp (+3-1) 
- (modified) mlir/test/Dialect/Arith/int-range-narrowing.mlir (+8) 


``````````diff
diff --git a/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp b/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp
index 49f89e1bd17f3..6e0269d1c92d8 100644
--- a/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp
+++ b/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp
@@ -41,6 +41,9 @@ void arith::ConstantOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
   }
   if (auto arrayCstAttr =
           llvm::dyn_cast_or_null<DenseIntElementsAttr>(getValue())) {
+    if (arrayCstAttr.empty())
+      return;
+
     if (arrayCstAttr.isSplat()) {
       setResultRange(getResult(), ConstantIntRanges::constant(
                                       arrayCstAttr.getSplatValue<APInt>()));
@@ -53,7 +56,6 @@ void arith::ConstantOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
       result = (result ? result->rangeUnion(range) : range);
     }
 
-    assert(result && "Zero-sized vectors are not allowed");
     setResultRange(getResult(), *result);
     return;
   }
diff --git a/mlir/test/Dialect/Arith/int-range-narrowing.mlir b/mlir/test/Dialect/Arith/int-range-narrowing.mlir
index 7ba22af0c0f1b..5938f24868cb9 100644
--- a/mlir/test/Dialect/Arith/int-range-narrowing.mlir
+++ b/mlir/test/Dialect/Arith/int-range-narrowing.mlir
@@ -528,3 +528,11 @@ func.func @affine_for_dynamic_bound(%n: index) {
   }
   return
 }
+
+// Ensure a zero-element constant tensor doesn't crash range inference.
+// CHECK-LABEL: func @zero_element_tensor_constant
+func.func @zero_element_tensor_constant() -> tensor<2x0xi32> {
+  // CHECK: arith.constant dense<> : tensor<2x0xi32>
+  %0 = arith.constant dense<> : tensor<2x0xi32>
+  return %0 : tensor<2x0xi32>
+}

``````````

</details>


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