[Mlir-commits] [mlir] [MLIR][NVVM] Add sqrt Ops (PR #197422)
Varad Rahul Kamthe
llvmlistbot at llvm.org
Wed May 13 04:26:06 PDT 2026
https://github.com/varadk27 updated https://github.com/llvm/llvm-project/pull/197422
>From 3dede0dcdb761164d97fb8d3a4cccbd3ccde18b8 Mon Sep 17 00:00:00 2001
From: Varad Rahul Kamthe <vkamthe at nvidia.com>
Date: Thu, 7 May 2026 06:25:10 +0000
Subject: [PATCH] Add sqrt Ops
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 35 +++++++++++
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 61 +++++++++++++++++++
mlir/test/Target/LLVMIR/nvvm/sqrt/sqrt.mlir | 47 ++++++++++++++
.../Target/LLVMIR/nvvm/sqrt/sqrt_invalid.mlir | 17 ++++++
4 files changed, 160 insertions(+)
create mode 100644 mlir/test/Target/LLVMIR/nvvm/sqrt/sqrt.mlir
create mode 100644 mlir/test/Target/LLVMIR/nvvm/sqrt/sqrt_invalid.mlir
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..4feddaeed22c1
--- /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
+}
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..d35e0a5e708df
--- /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
+}
More information about the Mlir-commits
mailing list