[mlir][arith] Add mulf(x, 0) -> 0 to mulf folder #161395

python3kgae · 2025-09-30T15:59:07Z

Fold mulf(x, 0) -> 0 when (nnan | ninf)

Fold `mulf(x, 0) -> 0`. Updated the yield_constant_loop test in mlir/test/Dialect/SCF/loop-pipelining.mlir to workaround [TODO](https://github.com/llvm/llvm-project/blob/main/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp#L163) in TestSCFUtils.cpp

llvmbot · 2025-09-30T15:59:36Z

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-scf

@llvm/pr-subscribers-mlir-arith

Author: Xiang Li (python3kgae)

Changes

Fold mulf(x, 0) -> 0.

Updated the yield_constant_loop test in mlir/test/Dialect/SCF/loop-pipelining.mlir
to workaround TODO in TestSCFUtils.cpp

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

3 Files Affected:

(modified) mlir/lib/Dialect/Arith/IR/ArithOps.cpp (+3)
(modified) mlir/test/Dialect/Arith/canonicalize.mlir (+10)
(modified) mlir/test/Dialect/SCF/loop-pipelining.mlir (+6-6)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp index 7cfd6d3a98df8..676297f56ac0f 100644 --- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp +++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp @@ -1281,6 +1281,9 @@ OpFoldResult arith::MulFOp::fold(FoldAdaptor adaptor) { // mulf(x, 1) -> x if (matchPattern(adaptor.getRhs(), m_OneFloat())) return getLhs(); + // mulf(x, 0) -> 0 + if (matchPattern(adaptor.getRhs(), m_AnyZeroFloat())) + return getRhs(); return constFoldBinaryOp<FloatAttr>( adaptor.getOperands(), diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir index ca3de3a2d7703..4c72a1bb27b01 100644 --- a/mlir/test/Dialect/Arith/canonicalize.mlir +++ b/mlir/test/Dialect/Arith/canonicalize.mlir @@ -2216,6 +2216,16 @@ func.func @test_mulf1(%arg0 : f32, %arg1 : f32) -> (f32) { return %2 : f32 } +// CHECK-LABEL: @test_mulf2( +func.func @test_mulf2(%arg0 : f32, %arg1 : f32) -> (f32, f32) { + // CHECK-NEXT: %[[C0:.+]] = arith.constant 0.000000e+00 : f32 + // CHECK-NEXT: return %[[C0]], %[[C0]] + %c0 = arith.constant 0.0 : f32 + %0 = arith.mulf %arg0, %c0 : f32 + %1 = arith.mulf %c0, %arg1 : f32 + return %0, %1 : f32, f32 +} + // ----- // CHECK-LABEL: @test_divf( diff --git a/mlir/test/Dialect/SCF/loop-pipelining.mlir b/mlir/test/Dialect/SCF/loop-pipelining.mlir index 86af637fc05d7..11dc55c7ebb17 100644 --- a/mlir/test/Dialect/SCF/loop-pipelining.mlir +++ b/mlir/test/Dialect/SCF/loop-pipelining.mlir @@ -930,7 +930,7 @@ func.func @dynamic_loop_result(%A: memref<?xf32>, %result: memref<?xf32>, %lb: i // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index -// CHECK-DAG: %[[CST0:.*]] = arith.constant 0.000000e+00 : f32 +// CHECK-DAG: %[[CST10:.*]] = arith.constant 1.000000e+01 : f32 // CHECK-DAG: %[[CST2:.*]] = arith.constant 2.000000e+00 : f32 // Prologue: // CHECK: %[[L0:.*]] = memref.load %[[A]][%[[C0]]] : memref<?xf32> @@ -938,15 +938,15 @@ func.func @dynamic_loop_result(%A: memref<?xf32>, %result: memref<?xf32>, %lb: i // CHECK-NEXT: %[[L1:.*]]:2 = scf.for %[[IV:.*]] = %[[C0]] to %[[C3]] // CHECK-SAME: step %[[C1]] iter_args(%[[ARG0:.*]] = %[[CST2]], %[[ARG1:.*]] = %[[L0]]) -> (f32, f32) { // CHECK-NEXT: %[[ADD0:.*]] = arith.addf %[[ARG1]], %[[ARG0]] : f32 -// CHECK-NEXT: %[[MUL0:.*]] = arith.mulf %[[ADD0]], %[[CST0]] : f32 +// CHECK-NEXT: %[[MUL0:.*]] = arith.mulf %[[ADD0]], %[[CST10]] : f32 // CHECK-NEXT: memref.store %[[MUL0]], %[[A]][%[[IV]]] : memref<?xf32> // CHECK-NEXT: %[[IV1:.*]] = arith.addi %[[IV]], %[[C1]] : index // CHECK-NEXT: %[[L2:.*]] = memref.load %[[A]][%[[IV1]]] : memref<?xf32> -// CHECK-NEXT: scf.yield %[[CST0]], %[[L2]] : f32 +// CHECK-NEXT: scf.yield %[[CST10]], %[[L2]] : f32 // CHECK-NEXT: } // Epilogue: -// CHECK-NEXT: %[[ADD1:.*]] = arith.addf %[[L1]]#1, %[[CST0]] : f32 -// CHECK-NEXT: %[[MUL1:.*]] = arith.mulf %[[ADD1]], %[[CST0]] : f32 +// CHECK-NEXT: %[[ADD1:.*]] = arith.addf %[[L1]]#1, %[[CST10]] : f32 +// CHECK-NEXT: %[[MUL1:.*]] = arith.mulf %[[ADD1]], %[[CST10]] : f32 // CHECK-NEXT: memref.store %[[MUL1]], %[[A]][%[[C3]]] : memref<?xf32> // CHECK-NEXT: return %[[L1]]#0 : f32 @@ -954,7 +954,7 @@ func.func @yield_constant_loop(%A: memref<?xf32>) -> f32 { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c4 = arith.constant 4 : index - %cf0 = arith.constant 0.0 : f32 + %cf0 = arith.constant 10.0 : f32 %cf2 = arith.constant 2.0 : f32 %r = scf.for %i0 = %c0 to %c4 step %c1 iter_args(%arg0 = %cf2) -> f32 { %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 3 } : memref<?xf32>

ThomasRaoux · 2025-09-30T16:02:28Z

mlir/test/Dialect/Arith/canonicalize.mlir

+// CHECK-LABEL: @test_mulf2(
+func.func @test_mulf2(%arg0 : f32, %arg1 : f32) -> (f32, f32) {
+ // CHECK-NEXT: %[[C0:.+]] = arith.constant 0.000000e+00 : f32
+ // CHECK-NEXT: return %[[C0]], %[[C0]]
+ %c0 = arith.constant 0.0 : f32
+ %0 = arith.mulf %arg0, %c0 : f32
+ %1 = arith.mulf %c0, %arg1 : f32
+ return %0, %1 : f32, f32
+}


that's not correct for Nan

Also does not preserve the sign of the operand.

We could do all this with fast-math flags.

Fixed by fold NaN before 0.

that's still not correct, the value may dynamically be Nan and not be a constant Nan

I see.
Updated with fast-math flags.

Is nnan/ninf enough to lose the sign of the input?

I don't think so :( .
Will FastMathFlags::nsz cover it or we'll have to go FastMathFlags::fast?

Added FastMathFlags::nsz.

mlir/test/Dialect/Arith/canonicalize.mlir

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

kuhar

Since you match anyZeroFloat, maybe also add a test case with -0.0? Looks good otherwise.

python3kgae · 2025-10-01T13:20:41Z

Since you match anyZeroFloat, maybe also add a test case with -0.0? Looks good otherwise.

Done.

mlir/test/Dialect/Arith/canonicalize.mlir

ThomasRaoux · 2025-10-02T02:59:06Z

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

+ if (arith::bitEnumContainsAll(getFastmath(), arith::FastMathFlags::nnan |
+ arith::FastMathFlags::nsz)) {


doesn't it need also ninf? inf * 0 -> Nan

I tried to check this with Alive: https://alive2.llvm.org/ce/z/wvNkdy

It's because nnan applies to the result as well:

nnan
No NaNs - Allow optimizations to assume the arguments and result are not NaN.

Fold `mulf(x, 0) -> 0` when (nnan | nsz)

python3kgae added the mlir:arith label Sep 30, 2025

llvmbot added mlir mlir:scf labels Sep 30, 2025

ThomasRaoux requested changes Sep 30, 2025

View reviewed changes

mulf(NaN, x) -> NaN

8d49f3f

python3kgae changed the title ~~[mlir][arith] Add mulf(x, 0) -> 0 to mulf folder~~ [mlir][arith] Add more patterns to mulf folder Sep 30, 2025

kuhar self-requested a review September 30, 2025 17:45

Limit to fastmask nnan | ninf

8c72844

python3kgae changed the title ~~[mlir][arith] Add more patterns to mulf folder~~ [mlir][arith] Add mulf(x, 0) -> 0 to mulf folder Sep 30, 2025

joker-eph reviewed Sep 30, 2025

View reviewed changes

mlir/test/Dialect/Arith/canonicalize.mlir Outdated Show resolved Hide resolved

Xiang Li added 2 commits September 30, 2025 18:39

Update per comment.

e86191a

Add nsz.

bdb0f61

kuhar reviewed Oct 1, 2025

View reviewed changes

mlir/lib/Dialect/Arith/IR/ArithOps.cpp Outdated Show resolved Hide resolved

mlir/lib/Dialect/Arith/IR/ArithOps.cpp Outdated Show resolved Hide resolved

Remove ninf which is not needed.

18e9504

kuhar reviewed Oct 1, 2025

View reviewed changes

Add test for neg zero.

f43001f

kuhar approved these changes Oct 1, 2025

View reviewed changes

mlir/test/Dialect/Arith/canonicalize.mlir Outdated Show resolved Hide resolved

Use CHECK_DAG.

4019c36

python3kgae merged commit 2d06374 into llvm:main Oct 2, 2025
9 checks passed

python3kgae deleted the fold_mulf_0 branch October 2, 2025 02:47

ThomasRaoux reviewed Oct 2, 2025

View reviewed changes

mahesh-attarde pushed a commit to mahesh-attarde/llvm-project that referenced this pull request Oct 3, 2025

[mlir][arith] Add mulf(x, 0) -> 0 to mulf folder (llvm#161395)

58f6605

Fold `mulf(x, 0) -> 0` when (nnan | nsz)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][arith] Add mulf(x, 0) -> 0 to mulf folder #161395

[mlir][arith] Add mulf(x, 0) -> 0 to mulf folder #161395

Uh oh!

python3kgae commented Sep 30, 2025 •

edited

Loading

llvmbot commented Sep 30, 2025 •

edited

Loading

ThomasRaoux Sep 30, 2025

joker-eph Sep 30, 2025

python3kgae Sep 30, 2025

ThomasRaoux Sep 30, 2025

python3kgae Sep 30, 2025

joker-eph Sep 30, 2025

python3kgae Sep 30, 2025

python3kgae Sep 30, 2025

Uh oh!

Uh oh!

Uh oh!

kuhar left a comment

python3kgae commented Oct 1, 2025

Uh oh!

Uh oh!

ThomasRaoux Oct 2, 2025

kuhar Oct 2, 2025

kuhar Oct 2, 2025 •

edited

Loading

Labels

5 participants

		if (arith::bitEnumContainsAll(getFastmath(), arith::FastMathFlags::nnan \|
		arith::FastMathFlags::nsz)) {

[mlir][arith] Add mulf(x, 0) -> 0 to mulf folder #161395

[mlir][arith] Add mulf(x, 0) -> 0 to mulf folder #161395

Uh oh!

Conversation

python3kgae commented Sep 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

llvmbot commented Sep 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

kuhar left a comment

Choose a reason for hiding this comment

python3kgae commented Oct 1, 2025

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kuhar Oct 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Labels

5 participants

python3kgae commented Sep 30, 2025 •

edited

Loading

llvmbot commented Sep 30, 2025 •

edited

Loading

kuhar Oct 2, 2025 •

edited

Loading