[Mlir-commits] [mlir] Fix 228315 (PR #228861)
Madhur Kumar
llvmlistbot at llvm.org
Sun Oct 4 01:54:17 PDT 2026
https://github.com/MadhurKumar004 created https://github.com/llvm/llvm-project/pull/228861
fixes:https://github.com/llvm/llvm-project/issues/228315
>From 7da366c333f844ed54aa8eb820dee990e1f8180b Mon Sep 17 00:00:00 2001
From: Madhur Kumar <madhurkumar004 at gmail.com>
Date: Sun, 4 Oct 2026 14:04:54 +0530
Subject: [PATCH 1/2] [MLIR][NVVM] Pre-commit Test
---
.../Dialect/LLVMIR/optimize-for-nvvm.mlir | 23 +++++++++++++++++++
1 file changed, 23 insertions(+)
diff --git a/mlir/test/Dialect/LLVMIR/optimize-for-nvvm.mlir b/mlir/test/Dialect/LLVMIR/optimize-for-nvvm.mlir
index 2ef1982c40214d5..21c3b4ab5ba4cbb 100644
--- a/mlir/test/Dialect/LLVMIR/optimize-for-nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/optimize-for-nvvm.mlir
@@ -22,3 +22,26 @@ llvm.func @fdiv_fp16(%arg0 : f16, %arg1 : f16) -> f16 {
// CHECK: llvm.return %[[result]] : f16
llvm.return %result : f16
}
+
+// CHECK-LABEL: llvm.func @fdiv_fp16_fastmath
+llvm.func @fdiv_fp16_fastmath(%arg0 : f16, %arg1 : f16) -> f16 {
+ // CHECK-DAG: %[[c0:.*]] = llvm.mlir.constant(0 : i32) : i32
+ // CHECK-DAG: %[[mask:.*]] = llvm.mlir.constant(2139095040 : i32) : i32
+ // CHECK-DAG: %[[lhs:.*]] = llvm.fpext %arg0 : f16 to f32
+ // CHECK-DAG: %[[rhs:.*]] = llvm.fpext %arg1 : f16 to f32
+ // CHECK-DAG: %[[rcp:.*]] = nvvm.rcp.approx.ftz.f %[[rhs]] : f32
+ // CHECK-DAG: %[[approx:.*]] = llvm.fmul %[[lhs]], %[[rcp]] : f32
+ // CHECK-DAG: %[[neg:.*]] = llvm.fneg %[[rhs]] : f32
+ // CHECK-DAG: %[[err:.*]] = llvm.intr.fma(%[[approx]], %[[neg]], %[[lhs]]) : (f32, f32, f32) -> f32
+ // CHECK-DAG: %[[refined:.*]] = llvm.intr.fma(%[[err]], %[[rcp]], %[[approx]]) : (f32, f32, f32) -> f32
+ // CHECK-DAG: %[[cast:.*]] = llvm.bitcast %[[approx]] : f32 to i32
+ // CHECK-DAG: %[[exp:.*]] = llvm.and %[[cast]], %[[mask]] : i32
+ // CHECK-DAG: %[[is_zero:.*]] = llvm.icmp "eq" %[[exp]], %[[c0]] : i32
+ // CHECK-DAG: %[[is_mask:.*]] = llvm.icmp "eq" %[[exp]], %[[mask]] : i32
+ // CHECK-DAG: %[[pred:.*]] = llvm.or %[[is_zero]], %[[is_mask]] : i1
+ // CHECK-DAG: %[[select:.*]] = llvm.select %[[pred]], %[[approx]], %[[refined]] : i1, f32
+ // CHECK-DAG: %[[result:.*]] = llvm.fptrunc %[[select]] : f32 to f16
+ %result = llvm.fdiv %arg0, %arg1 fastmath<nnan> : f16
+ // CHECK: llvm.return %[[result]] : f16
+ llvm.return %result : f16
+}
>From b8d81ab9773315e547268d973427a49ef01d6d5b Mon Sep 17 00:00:00 2001
From: Madhur Kumar <madhurkumar004 at gmail.com>
Date: Sun, 4 Oct 2026 14:22:26 +0530
Subject: [PATCH 2/2] [MLIR][NVVM] Preserve FastMathFlags in ExpandDivF16
---
.../Dialect/NVVM/Transforms/OptimizeForNVVM.cpp | 17 ++++++++++-------
mlir/test/Dialect/LLVMIR/optimize-for-nvvm.mlir | 14 +++++++-------
2 files changed, 17 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Dialect/NVVM/Transforms/OptimizeForNVVM.cpp b/mlir/lib/Dialect/NVVM/Transforms/OptimizeForNVVM.cpp
index db6e746d3ecc682..9499050d458b370 100644
--- a/mlir/lib/Dialect/NVVM/Transforms/OptimizeForNVVM.cpp
+++ b/mlir/lib/Dialect/NVVM/Transforms/OptimizeForNVVM.cpp
@@ -55,22 +55,25 @@ LogicalResult ExpandDivF16::matchAndRewrite(LLVM::FDivOp op,
return rewriter.notifyMatchFailure(op, "not f16");
Location loc = op.getLoc();
+ LLVM::FastmathFlags fmf = op.getFastmathFlags();
+
Type f32Type = rewriter.getF32Type();
Type i32Type = rewriter.getI32Type();
// Extend lhs and rhs to fp32.
- Value lhs = LLVM::FPExtOp::create(rewriter, loc, f32Type, op.getLhs());
- Value rhs = LLVM::FPExtOp::create(rewriter, loc, f32Type, op.getRhs());
+ Value lhs = LLVM::FPExtOp::create(rewriter, loc, f32Type, op.getLhs(), fmf);
+ Value rhs = LLVM::FPExtOp::create(rewriter, loc, f32Type, op.getRhs(), fmf);
// float rcp = rcp.approx.ftz.f32(rhs), approx = lhs * rcp.
Value rcp = NVVM::RcpApproxFtzF32Op::create(rewriter, loc, f32Type, rhs);
- Value approx = LLVM::FMulOp::create(rewriter, loc, lhs, rcp);
+ Value approx = LLVM::FMulOp::create(rewriter, loc, lhs, rcp, fmf);
// Refine the approximation with one Newton iteration:
// float refined = approx + (lhs - approx * rhs) * rcp;
- Value err = LLVM::FMAOp::create(
- rewriter, loc, approx, LLVM::FNegOp::create(rewriter, loc, rhs), lhs);
- Value refined = LLVM::FMAOp::create(rewriter, loc, err, rcp, approx);
+ Value err = LLVM::FMAOp::create(rewriter, loc, approx,
+ LLVM::FNegOp::create(rewriter, loc, rhs, fmf),
+ lhs, fmf);
+ Value refined = LLVM::FMAOp::create(rewriter, loc, err, rcp, approx, fmf);
// Use refined value if approx is normal (exponent neither all 0 or all 1).
Value mask = LLVM::ConstantOp::create(rewriter, loc, i32Type,
@@ -87,7 +90,7 @@ LogicalResult ExpandDivF16::matchAndRewrite(LLVM::FDivOp op,
LLVM::SelectOp::create(rewriter, loc, f32Type, pred, approx, refined);
// Replace with trucation back to fp16.
- rewriter.replaceOpWithNewOp<LLVM::FPTruncOp>(op, op.getType(), result);
+ rewriter.replaceOpWithNewOp<LLVM::FPTruncOp>(op, op.getType(), result, fmf);
return success();
}
diff --git a/mlir/test/Dialect/LLVMIR/optimize-for-nvvm.mlir b/mlir/test/Dialect/LLVMIR/optimize-for-nvvm.mlir
index 21c3b4ab5ba4cbb..0f7419066c83dec 100644
--- a/mlir/test/Dialect/LLVMIR/optimize-for-nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/optimize-for-nvvm.mlir
@@ -27,20 +27,20 @@ llvm.func @fdiv_fp16(%arg0 : f16, %arg1 : f16) -> f16 {
llvm.func @fdiv_fp16_fastmath(%arg0 : f16, %arg1 : f16) -> f16 {
// CHECK-DAG: %[[c0:.*]] = llvm.mlir.constant(0 : i32) : i32
// CHECK-DAG: %[[mask:.*]] = llvm.mlir.constant(2139095040 : i32) : i32
- // CHECK-DAG: %[[lhs:.*]] = llvm.fpext %arg0 : f16 to f32
- // CHECK-DAG: %[[rhs:.*]] = llvm.fpext %arg1 : f16 to f32
+ // CHECK-DAG: %[[lhs:.*]] = llvm.fpext %arg0 fastmath<nnan> : f16 to f32
+ // CHECK-DAG: %[[rhs:.*]] = llvm.fpext %arg1 fastmath<nnan> : f16 to f32
// CHECK-DAG: %[[rcp:.*]] = nvvm.rcp.approx.ftz.f %[[rhs]] : f32
- // CHECK-DAG: %[[approx:.*]] = llvm.fmul %[[lhs]], %[[rcp]] : f32
- // CHECK-DAG: %[[neg:.*]] = llvm.fneg %[[rhs]] : f32
- // CHECK-DAG: %[[err:.*]] = llvm.intr.fma(%[[approx]], %[[neg]], %[[lhs]]) : (f32, f32, f32) -> f32
- // CHECK-DAG: %[[refined:.*]] = llvm.intr.fma(%[[err]], %[[rcp]], %[[approx]]) : (f32, f32, f32) -> f32
+ // CHECK-DAG: %[[approx:.*]] = llvm.fmul %[[lhs]], %[[rcp]] fastmath<nnan> : f32
+ // CHECK-DAG: %[[neg:.*]] = llvm.fneg %[[rhs]] fastmath<nnan> : f32
+ // CHECK-DAG: %[[err:.*]] = llvm.intr.fma(%[[approx]], %[[neg]], %[[lhs]]) fastmath<nnan> : (f32, f32, f32) -> f32
+ // CHECK-DAG: %[[refined:.*]] = llvm.intr.fma(%[[err]], %[[rcp]], %[[approx]]) fastmath<nnan> : (f32, f32, f32) -> f32
// CHECK-DAG: %[[cast:.*]] = llvm.bitcast %[[approx]] : f32 to i32
// CHECK-DAG: %[[exp:.*]] = llvm.and %[[cast]], %[[mask]] : i32
// CHECK-DAG: %[[is_zero:.*]] = llvm.icmp "eq" %[[exp]], %[[c0]] : i32
// CHECK-DAG: %[[is_mask:.*]] = llvm.icmp "eq" %[[exp]], %[[mask]] : i32
// CHECK-DAG: %[[pred:.*]] = llvm.or %[[is_zero]], %[[is_mask]] : i1
// CHECK-DAG: %[[select:.*]] = llvm.select %[[pred]], %[[approx]], %[[refined]] : i1, f32
- // CHECK-DAG: %[[result:.*]] = llvm.fptrunc %[[select]] : f32 to f16
+ // CHECK-DAG: %[[result:.*]] = llvm.fptrunc %[[select]] fastmath<nnan> : f32 to f16
%result = llvm.fdiv %arg0, %arg1 fastmath<nnan> : f16
// CHECK: llvm.return %[[result]] : f16
llvm.return %result : f16
More information about the Mlir-commits
mailing list