[Mlir-commits] [mlir] [mlir][affine] Fix crash in AffineParallelLowering for unsupported reductions (PR #186189)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-arith

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

When lowering affine.parallel with a reduction kind that has no identity value (e.g. "assign"), getIdentityValueAttr() returns nullptr. The caller getIdentityValue() then passed this null TypedAttr to arith::ConstantOp::create(), triggering an LLVM_ERROR crash:
  "Failed to infer result type(s): arith.constant"

Fixes #<!-- -->185250

Assisted-by: Claude Code

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


3 Files Affected:

- (modified) mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp (+7-3) 
- (modified) mlir/lib/Dialect/Arith/IR/ArithOps.cpp (+2) 
- (added) mlir/test/Conversion/AffineToStandard/lower-affine-invalid.mlir (+17) 


``````````diff
diff --git a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
index 3b148f9021666..826d7547716e4 100644
--- a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
+++ b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
@@ -134,7 +134,7 @@ class AffineYieldOpLowering : public OpRewritePattern<AffineYieldOp> {
 
   LogicalResult matchAndRewrite(AffineYieldOp op,
                                 PatternRewriter &rewriter) const override {
-    if (isa<scf::ParallelOp>(op->getParentOp())) {
+    if (isa<scf::ParallelOp, AffineParallelOp>(op->getParentOp())) {
       // Terminator is rewritten as part of the "affine.parallel" lowering
       // pattern.
       return failure();
@@ -230,8 +230,12 @@ class AffineParallelLowering : public OpRewritePattern<AffineParallelOp> {
               static_cast<uint64_t>(cast<IntegerAttr>(reduction).getInt()));
       assert(reductionOp && "Reduction operation cannot be of None Type");
       arith::AtomicRMWKind reductionOpValue = *reductionOp;
-      identityVals.push_back(
-          arith::getIdentityValue(reductionOpValue, resultType, rewriter, loc));
+      Value identityVal =
+          arith::getIdentityValue(reductionOpValue, resultType, rewriter, loc);
+      if (!identityVal)
+        return rewriter.notifyMatchFailure(
+            op, "unsupported reduction kind for identity value");
+      identityVals.push_back(identityVal);
     }
     parOp = scf::ParallelOp::create(rewriter, loc, lowerBoundTuple,
                                     upperBoundTuple, steps, identityVals,
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 6f368604df65a..d3e2e77d10341 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -2801,6 +2801,8 @@ Value mlir::arith::getIdentityValue(AtomicRMWKind op, Type resultType,
                                     bool useOnlyFiniteValue) {
   auto attr =
       getIdentityValueAttr(op, resultType, builder, loc, useOnlyFiniteValue);
+  if (!attr)
+    return {};
   return arith::ConstantOp::create(builder, loc, attr);
 }
 
diff --git a/mlir/test/Conversion/AffineToStandard/lower-affine-invalid.mlir b/mlir/test/Conversion/AffineToStandard/lower-affine-invalid.mlir
new file mode 100644
index 0000000000000..a357b27a2430e
--- /dev/null
+++ b/mlir/test/Conversion/AffineToStandard/lower-affine-invalid.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt --lower-affine %s 2>&1 | FileCheck %s
+
+// Test that affine.parallel with an unsupported reduction kind ("assign")
+// does not crash but emits a proper error message. Previously,
+// getIdentityValue would be called with a null TypedAttr and crash inside
+// arith::ConstantOp::build with "Failed to infer result type(s)".
+
+// CHECK: Reduction operation type not supported
+// CHECK-NOT: Failed to infer result type
+
+func.func @affine_parallel_assign_reduction_no_crash(%n: index) -> i32 {
+  %0 = affine.parallel (%i) = (0) to (%n) reduce ("assign") -> i32 {
+    %c0 = arith.constant 0 : i32
+    affine.yield %c0 : i32
+  }
+  return %0 : i32
+}

``````````

</details>


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