[Mlir-commits] [mlir] Fix 228315 (PR #228861)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-llvm

Author: Madhur Kumar (MadhurKumar004)

<details>
<summary>Changes</summary>

fixes:https://github.com/llvm/llvm-project/issues/228315

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


2 Files Affected:

- (modified) mlir/lib/Dialect/NVVM/Transforms/OptimizeForNVVM.cpp (+10-7) 
- (modified) mlir/test/Dialect/LLVMIR/optimize-for-nvvm.mlir (+23) 


``````````diff
diff --git a/mlir/lib/Dialect/NVVM/Transforms/OptimizeForNVVM.cpp b/mlir/lib/Dialect/NVVM/Transforms/OptimizeForNVVM.cpp
index db6e746d3ecc6..9499050d458b3 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 2ef1982c40214..0f7419066c83d 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 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]] 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]] fastmath<nnan> : f32 to f16
+  %result = llvm.fdiv %arg0, %arg1 fastmath<nnan> : f16
+  // CHECK: llvm.return %[[result]] : f16
+  llvm.return %result : f16
+}

``````````

</details>


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