[Mlir-commits] [mlir] [MLIR][NVVM] Add sqrt Ops (PR #197422)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed May 13 04:25:30 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-llvm
Author: Varad Rahul Kamthe (varadk27)
<details>
<summary>Changes</summary>
Adds two NVVM dialect ops covering all 14 floating-point `sqrt` forms:
- `nvvm.sqrt` -- IEEE-compliant sqrt with explicit rounding mode
(`sqrt.<RM>[.ftz].{f32,f64}`), 12 forms.
- `nvvm.sqrt.approx` -- fast approximate sqrt (`sqrt.approx[.ftz].f32`),
2 forms; uses the `NVVM_F32UnaryApproxOp` base class.
The two ops are split because the rounded forms require an explicit
rounding mode and support both f32 and f64, while the approx forms have
no rounding mode and are f32-only.
---
Full diff: https://github.com/llvm/llvm-project/pull/197422.diff
4 Files Affected:
- (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td (+35)
- (modified) mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp (+61)
- (added) mlir/test/Target/LLVMIR/nvvm/sqrt/sqrt.mlir (+47)
- (added) mlir/test/Target/LLVMIR/nvvm/sqrt/sqrt_invalid.mlir (+17)
``````````diff
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 70aad7ac095b9..ad7bfcadfba41 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -6544,6 +6544,41 @@ def NVVM_FmaOp : NVVM_Op<"fma", [Pure, SameOperandsAndResultType]> {
}];
}
+//===----------------------------------------------------------------------===//
+// NVVM sqrt op definitions
+//===----------------------------------------------------------------------===//
+
+def NVVM_SqrtOp
+ : NVVM_SingleResultIntrinsicOp<"sqrt", [Pure, SameOperandsAndResultType]> {
+ let summary = "Take the square root of a value";
+ let description = [{
+ Compute sqrt(a) and store the result in d.
+
+ For more information, see PTX ISA:
+ [sqrt](https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-sqrt)
+ }];
+ let arguments = (ins AnyTypeOf<[F32, F64]>:$src, FPArithRoundingMode:$rnd,
+ DefaultValuedAttr<BoolAttr, "false">:$ftz);
+ let results = (outs AnyTypeOf<[F32, F64]>:$res);
+ let assemblyFormat = "$src attr-dict `:` type($src)";
+ let hasVerifier = 1;
+}
+
+def NVVM_SqrtApproxOp : NVVM_F32UnaryApproxOp<"sqrt.approx"> {
+ let summary = "Square root (fast approximation)";
+ let description = [{
+ Computes a fast approximation of the square root of the input value
+ (`d = sqrt(a)`). The maximum relative error over the entire positive
+ finite range is 2^-23.
+
+ The `ftz` attribute, when set, flushes subnormal inputs and results to
+ sign-preserving zero.
+
+ For more information, see PTX ISA:
+ [sqrt](https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-sqrt)
+ }];
+}
+
//===----------------------------------------------------------------------===//
// NVVM tensormap.replace Op
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 943de71d34a96..3100fa7300684 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3339,6 +3339,16 @@ LogicalResult NVVM::FmaOp::verify() {
return success();
}
+LogicalResult NVVM::SqrtOp::verify() {
+ if (getRnd() == NVVM::FPRoundingMode::NONE)
+ return emitOpError("rounding mode must be specified");
+
+ if (getRes().getType().isF64() && getFtz())
+ return emitOpError("FTZ is not supported for f64");
+
+ return success();
+}
+
/// Packs the given `field` into the `result`.
/// The `result` is 64-bits and each `field` can be 32-bits or narrower.
static llvm::Value *
@@ -3535,6 +3545,57 @@ RsqrtOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
return {id, {mt.lookupValue(thisOp.getSrc())}};
}
+mlir::NVVM::IDArgPair
+SqrtOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
+ llvm::IRBuilderBase &builder) {
+ auto thisOp = cast<NVVM::SqrtOp>(op);
+ Type t = thisOp.getRes().getType();
+ NVVM::FPRoundingMode rndMode = thisOp.getRnd();
+ bool isFtz = thisOp.getFtz();
+
+ // f32: 5 RM slots (NONE→RN, RN, RM, RP, RZ) × 2 ftz states.
+ static constexpr llvm::Intrinsic::ID f32IDs[] = {
+ llvm::Intrinsic::nvvm_sqrt_rn_f, // NONE → default RN (verifier rejects
+ // NONE anyway)
+ llvm::Intrinsic::nvvm_sqrt_rn_f,
+ llvm::Intrinsic::nvvm_sqrt_rm_f,
+ llvm::Intrinsic::nvvm_sqrt_rp_f,
+ llvm::Intrinsic::nvvm_sqrt_rz_f,
+ llvm::Intrinsic::nvvm_sqrt_rn_ftz_f,
+ llvm::Intrinsic::nvvm_sqrt_rn_ftz_f,
+ llvm::Intrinsic::nvvm_sqrt_rm_ftz_f,
+ llvm::Intrinsic::nvvm_sqrt_rp_ftz_f,
+ llvm::Intrinsic::nvvm_sqrt_rz_ftz_f,
+ };
+ // f64: 5 RM slots, no ftz.
+ static constexpr llvm::Intrinsic::ID f64IDs[] = {
+ llvm::Intrinsic::nvvm_sqrt_rn_d, llvm::Intrinsic::nvvm_sqrt_rn_d,
+ llvm::Intrinsic::nvvm_sqrt_rm_d, llvm::Intrinsic::nvvm_sqrt_rp_d,
+ llvm::Intrinsic::nvvm_sqrt_rz_d,
+ };
+
+ llvm::Intrinsic::ID id = [&] {
+ if (t.isF32()) {
+ unsigned index = (isFtz * 5) + static_cast<unsigned>(rndMode);
+ return f32IDs[index];
+ }
+ // f64
+ return f64IDs[static_cast<unsigned>(rndMode)];
+ }();
+
+ return {id, {mt.lookupValue(thisOp.getSrc())}};
+}
+
+mlir::NVVM::IDArgPair
+SqrtApproxOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
+ llvm::IRBuilderBase &builder) {
+ auto thisOp = cast<NVVM::SqrtApproxOp>(op);
+ llvm::Intrinsic::ID id = thisOp.getFtz()
+ ? llvm::Intrinsic::nvvm_sqrt_approx_ftz_f
+ : llvm::Intrinsic::nvvm_sqrt_approx_f;
+ return {id, {mt.lookupValue(thisOp.getSrc())}};
+}
+
mlir::NVVM::IDArgPair
PMEventOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
llvm::IRBuilderBase &builder) {
diff --git a/mlir/test/Target/LLVMIR/nvvm/sqrt/sqrt.mlir b/mlir/test/Target/LLVMIR/nvvm/sqrt/sqrt.mlir
new file mode 100644
index 0000000000000..1ed9d1b5196ac
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/sqrt/sqrt.mlir
@@ -0,0 +1,47 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// f32 sqrt — all 8 forms (4 rounding modes × 2 ftz states).
+llvm.func @sqrt_f32(%a : f32) -> f32 {
+ // CHECK-LABEL: define float @sqrt_f32(float %0) {
+ // CHECK: call float @llvm.nvvm.sqrt.rn.f(float %{{.*}})
+ // CHECK: call float @llvm.nvvm.sqrt.rz.f(float %{{.*}})
+ // CHECK: call float @llvm.nvvm.sqrt.rm.f(float %{{.*}})
+ // CHECK: call float @llvm.nvvm.sqrt.rp.f(float %{{.*}})
+ // CHECK: call float @llvm.nvvm.sqrt.rn.ftz.f(float %{{.*}})
+ // CHECK: call float @llvm.nvvm.sqrt.rz.ftz.f(float %{{.*}})
+ // CHECK: call float @llvm.nvvm.sqrt.rm.ftz.f(float %{{.*}})
+ // CHECK: call float @llvm.nvvm.sqrt.rp.ftz.f(float %{{.*}})
+ %r1 = nvvm.sqrt %a {rnd = #nvvm.fp_rnd_mode<rn>} : f32
+ %r2 = nvvm.sqrt %r1 {rnd = #nvvm.fp_rnd_mode<rz>} : f32
+ %r3 = nvvm.sqrt %r2 {rnd = #nvvm.fp_rnd_mode<rm>} : f32
+ %r4 = nvvm.sqrt %r3 {rnd = #nvvm.fp_rnd_mode<rp>} : f32
+ %r5 = nvvm.sqrt %r4 {rnd = #nvvm.fp_rnd_mode<rn>, ftz = true} : f32
+ %r6 = nvvm.sqrt %r5 {rnd = #nvvm.fp_rnd_mode<rz>, ftz = true} : f32
+ %r7 = nvvm.sqrt %r6 {rnd = #nvvm.fp_rnd_mode<rm>, ftz = true} : f32
+ %r8 = nvvm.sqrt %r7 {rnd = #nvvm.fp_rnd_mode<rp>, ftz = true} : f32
+ llvm.return %r8 : f32
+}
+
+// f64 sqrt — all 4 forms (4 rounding modes, no ftz).
+llvm.func @sqrt_f64(%a : f64) -> f64 {
+ // CHECK-LABEL: define double @sqrt_f64(double %0) {
+ // CHECK: call double @llvm.nvvm.sqrt.rn.d(double %{{.*}})
+ // CHECK: call double @llvm.nvvm.sqrt.rz.d(double %{{.*}})
+ // CHECK: call double @llvm.nvvm.sqrt.rm.d(double %{{.*}})
+ // CHECK: call double @llvm.nvvm.sqrt.rp.d(double %{{.*}})
+ %r1 = nvvm.sqrt %a {rnd = #nvvm.fp_rnd_mode<rn>} : f64
+ %r2 = nvvm.sqrt %r1 {rnd = #nvvm.fp_rnd_mode<rz>} : f64
+ %r3 = nvvm.sqrt %r2 {rnd = #nvvm.fp_rnd_mode<rm>} : f64
+ %r4 = nvvm.sqrt %r3 {rnd = #nvvm.fp_rnd_mode<rp>} : f64
+ llvm.return %r4 : f64
+}
+
+// sqrt.approx — 2 forms.
+llvm.func @sqrt_approx(%a : f32) -> f32 {
+ // CHECK-LABEL: define float @sqrt_approx(float %0) {
+ // CHECK: call float @llvm.nvvm.sqrt.approx.f(float %{{.*}})
+ // CHECK: call float @llvm.nvvm.sqrt.approx.ftz.f(float %{{.*}})
+ %r1 = nvvm.sqrt.approx %a : f32
+ %r2 = nvvm.sqrt.approx %r1 {ftz = true} : f32
+ llvm.return %r2 : f32
+}
\ No newline at end of file
diff --git a/mlir/test/Target/LLVMIR/nvvm/sqrt/sqrt_invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/sqrt/sqrt_invalid.mlir
new file mode 100644
index 0000000000000..486ae2aa85b95
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/sqrt/sqrt_invalid.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-translate --mlir-to-llvmir --split-input-file --verify-diagnostics %s
+
+// -----
+
+llvm.func @sqrt_invalid_no_rnd(%a : f32) -> f32 {
+ // expected-error at +1 {{rounding mode must be specified}}
+ %0 = nvvm.sqrt %a {rnd = #nvvm.fp_rnd_mode<none>} : f32
+ llvm.return %0 : f32
+}
+
+// -----
+
+llvm.func @sqrt_invalid_f64_ftz(%a : f64) -> f64 {
+ // expected-error at +1 {{FTZ is not supported for f64}}
+ %0 = nvvm.sqrt %a {rnd = #nvvm.fp_rnd_mode<rn>, ftz = true} : f64
+ llvm.return %0 : f64
+}
\ No newline at end of file
``````````
</details>
https://github.com/llvm/llvm-project/pull/197422
More information about the Mlir-commits
mailing list