[Mlir-commits] [mlir] [MLIR] Fix out-of-bounds crash in matchReduction for ops with fewer yield operands (PR #183555)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

`mlir::matchReduction()` accessed `terminatorOp->getOperand(redPos)` without checking that `redPos` is within the terminator's operand count. This caused an assertion failure when the region's block-argument count exceeds the terminator's yield count, e.g. for `linalg.pooling_nhwc_sum` whose kernel region has three block args but yields only one value.

Add a bounds check before the operand access so the function returns nullptr gracefully instead of crashing.

Fixes #<!-- -->131437

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


2 Files Affected:

- (modified) mlir/lib/Analysis/SliceAnalysis.cpp (+2-1) 
- (modified) mlir/test/Analysis/test-match-reduction.mlir (+16) 


``````````diff
diff --git a/mlir/lib/Analysis/SliceAnalysis.cpp b/mlir/lib/Analysis/SliceAnalysis.cpp
index 12dff19ed31d3..f388cd8041d6b 100644
--- a/mlir/lib/Analysis/SliceAnalysis.cpp
+++ b/mlir/lib/Analysis/SliceAnalysis.cpp
@@ -329,7 +329,8 @@ Value mlir::matchReduction(ArrayRef<BlockArgument> iterCarriedArgs,
   // Check that the yielded value is in the same position as in
   // `iterCarriedArgs`.
   Operation *terminatorOp = combinerOp;
-  if (terminatorOp->getOperand(redPos) != combinerOps.back()->getResults()[0])
+  if (redPos >= terminatorOp->getNumOperands() ||
+      terminatorOp->getOperand(redPos) != combinerOps.back()->getResults()[0])
     return nullptr;
 
   return reducedVal;
diff --git a/mlir/test/Analysis/test-match-reduction.mlir b/mlir/test/Analysis/test-match-reduction.mlir
index b5902db77e899..7faddf902521b 100644
--- a/mlir/test/Analysis/test-match-reduction.mlir
+++ b/mlir/test/Analysis/test-match-reduction.mlir
@@ -112,3 +112,19 @@ func.func @affine_output_dep(%in: memref<512xf32>) {
  return
 }
 
+// -----
+
+// Verify that matchReduction does not crash when the terminator has fewer
+// operands than the number of iteration-carried block arguments (issue #131437).
+// expected-remark at below {{Testing function}}
+func.func @pooling_nhwc_sum_no_crash(%arg0: tensor<1x1x1x1xf32>,
+                                     %arg1: tensor<1x3x3x1xf32>) {
+  %0 = tensor.empty() : tensor<1x1xf32>
+  // expected-remark at below {{Reduction NOT found in output #0!}}
+  // expected-remark at below {{Reduction NOT found in output #1!}}
+  %1 = linalg.pooling_nhwc_sum
+         {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>}
+         ins(%arg1, %0 : tensor<1x3x3x1xf32>, tensor<1x1xf32>)
+         outs(%arg0 : tensor<1x1x1x1xf32>) -> tensor<1x1x1x1xf32>
+  return
+}

``````````

</details>


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