[Mlir-commits] [mlir] [MLIR][Linalg] Fix crash in Generic, Map, Reduce Ops for `getAsmBlockArgumentNames` when region is empty (PR #184743)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Arjun Bhamra (abhamra)

<details>
<summary>Changes</summary>

Originally, when calling  `--test-ir-visitors --mlir-print-assume-verified`, in the `testNoSkipErasureCallbacks` case, we would call `printBlock` which eventually constructs an `AsmState` object, using `getAsmBlockArgumentNames` internally.

For Linalg's `GenericOp`, `MapOp`, and `ReduceOp`, the implementations of `getAsmBlockArgumentNames` implicitly assume that the region is nonempty before calling `getRegionInputArgs`.

Notably, during the `--test-ir-visitors --mlir-print-assume-verified` pass call, `getAsmBlockArgumentNames` is called on ops whose regions have already been handled by the walk (since we handle ops before blocks), and this triggered the empty region error.

This PR introduces a simple empty region check and early return in each of the three Linalg ops using `getRegionInputArgs` in their `getAsmBlockArgumentNames` implementation, and adds a regression test for the same case.

Closes #<!-- -->128346

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


4 Files Affected:

- (modified) mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp (+1) 
- (modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+6) 
- (modified) mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp (+1) 
- (added) mlir/test/IR/test-visitors-assume-verified.mlir (+23) 


``````````diff
diff --git a/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp b/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp
index 49f89e1bd17f3..e0f43aa420e71 100644
--- a/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp
+++ b/mlir/lib/Dialect/Arith/IR/InferIntRangeInterfaceImpls.cpp
@@ -243,6 +243,7 @@ void arith::ExtSIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
 
 void arith::TruncIOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                         SetIntRangeFn setResultRange) {
+  // NOTE: ISSUE HERE, BE CAREFUL
   unsigned destWidth =
       ConstantIntRanges::getStorageBitwidth(getResult().getType());
   setResultRange(getResult(), truncRange(argRanges[0], destWidth));
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index bfc03cc7436df..9b68a540c52e6 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -1102,6 +1102,8 @@ static void buildGenericRegion(
 
 void GenericOp::getAsmBlockArgumentNames(Region &region,
                                          OpAsmSetValueNameFn setNameFn) {
+  if (region.empty())
+    return;
   for (Value v : getRegionInputArgs())
     setNameFn(v, "in");
   for (Value v : getRegionOutputArgs())
@@ -1481,6 +1483,8 @@ static ParseResult parseDstStyleOp(
 
 void MapOp::getAsmBlockArgumentNames(Region &region,
                                      OpAsmSetValueNameFn setNameFn) {
+  if (region.empty())
+    return;
   for (Value v : getRegionInputArgs())
     setNameFn(v, "in");
   for (Value v : getRegionOutputArgs())
@@ -1738,6 +1742,8 @@ Speculation::Speculatability MapOp::getSpeculatability() {
 
 void ReduceOp::getAsmBlockArgumentNames(Region &region,
                                         OpAsmSetValueNameFn setNameFn) {
+  if (region.empty())
+    return;
   for (Value v : getRegionInputArgs())
     setNameFn(v, "in");
   for (Value v : getRegionOutputArgs())
diff --git a/mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp b/mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp
index 21f07ddce4495..68441ff6c2d13 100644
--- a/mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp
+++ b/mlir/lib/Interfaces/Utils/InferIntRangeCommon.cpp
@@ -147,6 +147,7 @@ ConstantIntRanges mlir::intrange::truncRange(const ConstantIntRanges &range,
   // the range of the resulting value is not contiguous ind includes 0.
   // Ex. If you truncate [256, 258] from i16 to i8, you validly get [0, 2],
   // but you can't truncate [255, 257] similarly.
+  // NOTE: Issue ends here, in the trunc call when destwidth = 0
   bool hasUnsignedRollover =
       range.umin().lshr(destWidth) != range.umax().lshr(destWidth);
   APInt umin = hasUnsignedRollover ? APInt::getZero(destWidth)
diff --git a/mlir/test/IR/test-visitors-assume-verified.mlir b/mlir/test/IR/test-visitors-assume-verified.mlir
new file mode 100644
index 0000000000000..8178be1187b16
--- /dev/null
+++ b/mlir/test/IR/test-visitors-assume-verified.mlir
@@ -0,0 +1,23 @@
+// RUN: mlir-opt -test-ir-visitors --mlir-print-assume-verified %s | FileCheck %s
+
+// Regression test: linalg ops implementing getAsmBlockArgumentNames via
+// getRegionInputArgs() used to crash during block erasure in no-skip walks
+// when combined with --mlir-print-assume-verified, because AsmState
+// construction would call getAsmBlockArgumentNames on ops whose regions
+// had already been emptied.
+
+func.func @test_no_skip_block_erasure_linalg_map(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = tensor.empty() : tensor<4xf32>
+  %1 = linalg.map ins(%arg0, %arg1 : tensor<4xf32>, tensor<4xf32>)
+                  outs(%0 : tensor<4xf32>)
+    (%in0: f32, %in1: f32, %out: f32) {
+      %2 = arith.addf %in0, %in1 : f32
+      linalg.yield %2 : f32
+    }
+  return %1 : tensor<4xf32>
+}
+
+// CHECK-LABEL: Block post-order erasures (no skip)
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'linalg.map'
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'func.func'
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'builtin.module'

``````````

</details>


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