[Mlir-commits] [mlir] [MLIR][NVVM] Add lg2.approx.f32 intrinsic with optional FTZ support (PR #167330)
Swarna Dharshini S
llvmlistbot at llvm.org
Mon Nov 10 07:44:24 PST 2025
https://github.com/SwarnaDharshiniS created https://github.com/llvm/llvm-project/pull/167330
This patch introduces the `nvvm.lg2.approx.f` operation to the NVVM dialect.
- Adds an NVVM op that maps to the PTX instruction `lg2.approx.f32`
- Supports an optional `ftz` (flush-to-zero) attribute:
When `ftz = true`, maps to `lg2.approx.ftz.f32`
Otherwise, maps to `lg2.approx.f32`
This enables MLIR lowering for approximate base-2 logarithm computations under NVVM.
- Added tests in `mlir/test/Dialect/LLVMIR/nvvm.mlir`
- Verified with `ninja check-mlir` — all tests passed
>From f418c155e1b1e41ad58fcdfe1fb2cdafa19222ba Mon Sep 17 00:00:00 2001
From: Swarna Dharshini S <swarnadharshinis at gmail.com>
Date: Tue, 4 Nov 2025 15:28:55 +0530
Subject: [PATCH 1/3] [MLIR] Add ex2 and lg2 operations to NVVM dialect
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 12 ++++++++++++
mlir/test/Dialect/LLVMIR/nvvm.mlir | 14 ++++++++++++++
2 files changed, 26 insertions(+)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 4f483859ac18d..7ac16e0ce0ca7 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -449,6 +449,18 @@ def NVVM_RcpApproxFtzF32Op : NVVM_IntrOp<"rcp.approx.ftz.f", [Pure], 1> {
let assemblyFormat = "$arg attr-dict `:` type($res)";
}
+def NVVM_ex2ApproxFtzF32Op : NVVM_IntrOp<"ex2.approx.ftz.f", [Pure], 1> {
+ let arguments = (ins F32:$arg);
+ let results = (outs F32:$res);
+ let assemblyFormat = "$arg attr-dict `:` type($res)";
+}
+
+def NVVM_lg2ApproxFtzF32Op : NVVM_IntrOp<"lg2.approx.ftz.f", [Pure], 1> {
+ let arguments = (ins F32:$arg);
+ let results = (outs F32:$res);
+ let assemblyFormat = "$arg attr-dict `:` type($res)";
+}
+
//===----------------------------------------------------------------------===//
// NVVM redux op definitions
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index 0243f5eb8c862..54b6a12e3d413 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir
@@ -36,6 +36,20 @@ func.func @nvvm_rcp(%arg0: f32) -> f32 {
llvm.return %0 : f32
}
+// CHECK-LABEL: @nvvm_ex2
+func.func @nvvm_ex2(%arg0: f32) -> f32 {
+ // CHECK: nvvm.ex2.approx.ftz.f %arg0 : f32
+ %0 = nvvm.ex2.approx.ftz.f %arg0 : f32
+ llvm.return %0 : f32
+}
+
+// CHECK-LABEL: @nvvm_lg2
+func.func @nvvm_lg2(%arg0: f32) -> f32 {
+ // CHECK: nvvm.lg2.approx.ftz.f %arg0 : f32
+ %0 = nvvm.lg2.approx.ftz.f %arg0 : f32
+ llvm.return %0 : f32
+}
+
// CHECK-LABEL: @llvm_nvvm_barrier0
func.func @llvm_nvvm_barrier0() {
// CHECK: nvvm.barrier0
>From 841d71a2461fefe8317f15778b1c13eb6f00b26f Mon Sep 17 00:00:00 2001
From: Swarna Dharshini S <swarnadharshinis at gmail.com>
Date: Wed, 5 Nov 2025 14:42:13 +0530
Subject: [PATCH 2/3] [MLIR][NVVM] Add lg2.approx.f and lg2.approx.ftz.f ops to
NVVM dialect
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 2 +-
mlir/test/Dialect/LLVMIR/nvvm.mlir | 12 ++++++------
2 files changed, 7 insertions(+), 7 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 7ac16e0ce0ca7..79988299b4ef9 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -449,7 +449,7 @@ def NVVM_RcpApproxFtzF32Op : NVVM_IntrOp<"rcp.approx.ftz.f", [Pure], 1> {
let assemblyFormat = "$arg attr-dict `:` type($res)";
}
-def NVVM_ex2ApproxFtzF32Op : NVVM_IntrOp<"ex2.approx.ftz.f", [Pure], 1> {
+def NVVM_lg2ApproxF32Op : NVVM_IntrOp<"lg2.approx.f", [Pure], 1> {
let arguments = (ins F32:$arg);
let results = (outs F32:$res);
let assemblyFormat = "$arg attr-dict `:` type($res)";
diff --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index 54b6a12e3d413..03735f2c11c72 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir
@@ -36,15 +36,15 @@ func.func @nvvm_rcp(%arg0: f32) -> f32 {
llvm.return %0 : f32
}
-// CHECK-LABEL: @nvvm_ex2
-func.func @nvvm_ex2(%arg0: f32) -> f32 {
- // CHECK: nvvm.ex2.approx.ftz.f %arg0 : f32
- %0 = nvvm.ex2.approx.ftz.f %arg0 : f32
+// CHECK-LABEL: @nvvm_lg2_f32
+func.func @nvvm_lg2_f32(%arg0: f32) -> f32 {
+ // CHECK: nvvm.lg2.approx.f %arg0 : f32
+ %0 = nvvm.lg2.approx.f %arg0 : f32
llvm.return %0 : f32
}
-// CHECK-LABEL: @nvvm_lg2
-func.func @nvvm_lg2(%arg0: f32) -> f32 {
+// CHECK-LABEL: @nvvm_lg2_ftz_f32
+func.func @nvvm_lg2_ftz_f32(%arg0: f32) -> f32 {
// CHECK: nvvm.lg2.approx.ftz.f %arg0 : f32
%0 = nvvm.lg2.approx.ftz.f %arg0 : f32
llvm.return %0 : f32
>From e9c5740bbae884e118eba50512d48b6ea6b1b663 Mon Sep 17 00:00:00 2001
From: Swarna Dharshini S <swarnadharshinis at gmail.com>
Date: Mon, 10 Nov 2025 16:37:20 +0530
Subject: [PATCH 3/3] [MLIR][NVVM] Add lg2.approx.f32 intrinsic with optional
FTZ support
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 33 ++++++++++++++++-----
mlir/test/Dialect/LLVMIR/nvvm.mlir | 12 ++++----
2 files changed, 32 insertions(+), 13 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 79988299b4ef9..4c32f4c3a4d30 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -449,16 +449,35 @@ def NVVM_RcpApproxFtzF32Op : NVVM_IntrOp<"rcp.approx.ftz.f", [Pure], 1> {
let assemblyFormat = "$arg attr-dict `:` type($res)";
}
-def NVVM_lg2ApproxF32Op : NVVM_IntrOp<"lg2.approx.f", [Pure], 1> {
- let arguments = (ins F32:$arg);
- let results = (outs F32:$res);
- let assemblyFormat = "$arg attr-dict `:` type($res)";
-}
+//===----------------------------------------------------------------------===//
+// NVVM lg2 approximate op
+//===----------------------------------------------------------------------===//
-def NVVM_lg2ApproxFtzF32Op : NVVM_IntrOp<"lg2.approx.ftz.f", [Pure], 1> {
- let arguments = (ins F32:$arg);
+def NVVM_Lg2ApproxF32Op : NVVM_Op<"lg2.approx.f", [Pure]> {
+
+ let summary = "Compute approximate base-2 log (lg2.approx.f32)";
+
+ let description = [{
+ Compute the approximate base-2 logarithm of the input value.
+ If 'ftz' is true, subnormal numbers are flushed to zero.
+ }];
+
+ let arguments = (ins
+ F32:$arg,
+ OptionalAttr<BoolAttr>:$ftz);
+
let results = (outs F32:$res);
+
let assemblyFormat = "$arg attr-dict `:` type($res)";
+
+ string llvmBuilder = [{
+ bool ftz = $ftz && *$ftz;
+ llvm::Intrinsic::ID intrinsicID =
+ ftz ? llvm::Intrinsic::nvvm_lg2_approx_ftz_f
+ : llvm::Intrinsic::nvvm_lg2_approx_f;
+
+ $res = createIntrinsicCall(builder, intrinsicID, {$arg});
+ }];
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index 03735f2c11c72..f4fbf5d02e542 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir
@@ -36,17 +36,17 @@ func.func @nvvm_rcp(%arg0: f32) -> f32 {
llvm.return %0 : f32
}
-// CHECK-LABEL: @nvvm_lg2_f32
-func.func @nvvm_lg2_f32(%arg0: f32) -> f32 {
+// CHECK-LABEL: @nvvm_lg2_approx_f
+func.func @nvvm_lg2_approx_f(%arg0: f32) -> f32 {
// CHECK: nvvm.lg2.approx.f %arg0 : f32
%0 = nvvm.lg2.approx.f %arg0 : f32
llvm.return %0 : f32
}
-// CHECK-LABEL: @nvvm_lg2_ftz_f32
-func.func @nvvm_lg2_ftz_f32(%arg0: f32) -> f32 {
- // CHECK: nvvm.lg2.approx.ftz.f %arg0 : f32
- %0 = nvvm.lg2.approx.ftz.f %arg0 : f32
+// CHECK-LABEL: @nvvm_lg2_approx_ftz_f
+func.func @nvvm_lg2_approx_ftz_f(%arg0: f32) -> f32 {
+ // CHECK: nvvm.lg2.approx.f %arg0 {ftz = true} : f32
+ %0 = nvvm.lg2.approx.f %arg0 {ftz = true} : f32
llvm.return %0 : f32
}
More information about the Mlir-commits
mailing list