[Mlir-commits] [mlir] [MLIR][NVVM] Add `nvvm.log2` OP (PR #193789)
Guray Ozen
llvmlistbot at llvm.org
Mon Apr 27 07:02:37 PDT 2026
https://github.com/grypp updated https://github.com/llvm/llvm-project/pull/193789
>From cf6f28df250a0173152515f6201c96a518dd309e Mon Sep 17 00:00:00 2001
From: Guray Ozen <gozen at nvidia.com>
Date: Fri, 24 Apr 2026 17:52:20 +0200
Subject: [PATCH 1/2] [MLIR][NVVM] Add `nvvm.lg2` OP
Implement `nvvm.lg2` with ftz flag
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 14 ++++++++++++++
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 10 ++++++++++
mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir | 14 ++++++++++++++
mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir | 14 ++++++++++++++
4 files changed, 52 insertions(+)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 7f1f9774abf52..d6a1d358b265f 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -558,6 +558,20 @@ def NVVM_CosOp : NVVM_SingleResultIntrinsicOp<"cos",
let assemblyFormat = "$src attr-dict `:` type($src)";
}
+def NVVM_Lg2Op : NVVM_SingleResultIntrinsicOp<"lg2",
+ [Pure, SameOperandsAndResultType]> {
+ let summary = "Base-2 logarithm (fast approximation)";
+ let description = [{
+ Computes a fast approximation of the base-2 logarithm of the input
+ value. The `ftz` attribute, when set, flushes subnormal inputs and
+ results to sign-preserving zero.
+ }];
+ let arguments = (ins F32:$src,
+ DefaultValuedAttr<BoolAttr, "false">:$ftz);
+ let results = (outs F32:$res);
+ let assemblyFormat = "$src attr-dict `:` type($src)";
+}
+
def NVVM_Ex2Op : NVVM_SingleResultIntrinsicOp<"ex2",
[Pure, SameOperandsAndResultType]> {
let summary = "Base-2 exponential (fast approximation)";
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index f403c952221d9..7dda2802dd659 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3495,6 +3495,16 @@ SinOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
return {id, {mt.lookupValue(thisOp.getSrc())}};
}
+mlir::NVVM::IDArgPair
+Lg2Op::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
+ llvm::IRBuilderBase &builder) {
+ auto thisOp = cast<NVVM::Lg2Op>(op);
+ llvm::Intrinsic::ID id = thisOp.getFtz()
+ ? llvm::Intrinsic::nvvm_lg2_approx_ftz_f
+ : llvm::Intrinsic::nvvm_lg2_approx_f;
+ return {id, {mt.lookupValue(thisOp.getSrc())}};
+}
+
mlir::NVVM::IDArgPair
Ex2Op::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
llvm::IRBuilderBase &builder) {
diff --git a/mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir b/mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir
index cfce0489fe0e3..54def11d02457 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir
@@ -29,6 +29,20 @@ func.func @nvvm_sin_ftz_f32(%arg0: f32) -> f32 {
return %0 : f32
}
+// CHECK-LABEL: @nvvm_lg2_f32
+func.func @nvvm_lg2_f32(%arg0: f32) -> f32 {
+ // CHECK: nvvm.lg2 {{.*}} : f32
+ %0 = nvvm.lg2 %arg0 : f32
+ return %0 : f32
+}
+
+// CHECK-LABEL: @nvvm_lg2_ftz_f32
+func.func @nvvm_lg2_ftz_f32(%arg0: f32) -> f32 {
+ // CHECK: nvvm.lg2 {{.*}} {ftz = true} : f32
+ %0 = nvvm.lg2 %arg0 {ftz = true} : f32
+ return %0 : f32
+}
+
// CHECK-LABEL: @nvvm_ex2_f32
func.func @nvvm_ex2_f32(%arg0: f32) -> f32 {
// CHECK: nvvm.ex2 {{.*}} : f32
diff --git a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
index d4fa373ad49ad..0762c403b7307 100644
--- a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
@@ -28,6 +28,20 @@ llvm.func @nvvm_cos_ftz(%arg0: f32) -> f32 {
llvm.return %0 : f32
}
+// CHECK-LABEL: @nvvm_lg2
+llvm.func @nvvm_lg2(%arg0: f32) -> f32 {
+ // CHECK: call float @llvm.nvvm.lg2.approx.f(float %{{.*}})
+ %0 = nvvm.lg2 %arg0 : f32
+ llvm.return %0 : f32
+}
+
+// CHECK-LABEL: @nvvm_lg2_ftz
+llvm.func @nvvm_lg2_ftz(%arg0: f32) -> f32 {
+ // CHECK: call float @llvm.nvvm.lg2.approx.ftz.f(float %{{.*}})
+ %0 = nvvm.lg2 %arg0 {ftz = true} : f32
+ llvm.return %0 : f32
+}
+
// CHECK-LABEL: @nvvm_ex2
llvm.func @nvvm_ex2(%arg0: f32) -> f32 {
// CHECK: call float @llvm.nvvm.ex2.approx.f32(float %{{.*}})
>From 3dccaf77ada3278c51ebeb4b516454cace2c34c2 Mon Sep 17 00:00:00 2001
From: Guray Ozen <gozen at nvidia.com>
Date: Mon, 27 Apr 2026 16:00:57 +0200
Subject: [PATCH 2/2] name change
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 6 +++---
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 6 +++---
mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir | 8 ++++----
mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir | 4 ++--
4 files changed, 12 insertions(+), 12 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index d6a1d358b265f..95fe1e0535843 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -521,8 +521,8 @@ def NVVM_RcpApproxFtzF32Op : NVVM_IntrOp<"rcp.approx.ftz.f", [Pure], 1> {
//===----------------------------------------------------------------------===//
// NVVM transcendental op definitions
//===----------------------------------------------------------------------===//
-// PTX provides only fast approximations for sin/cos/lg2/ex2/tanh. Exposing
-// these as plain `nvvm.{sin,cos,lg2,ex2,tanh}` (without a `.approx` suffix)
+// PTX provides only fast approximations for sin/cos/log2/ex2/tanh. Exposing
+// these as plain `nvvm.{sin,cos,log2,ex2,tanh}` (without a `.approx` suffix)
// keeps the NVVM dialect self-contained -- users don't need to reach for
// `math.*` ops for something that already has a PTX instruction.
@@ -558,7 +558,7 @@ def NVVM_CosOp : NVVM_SingleResultIntrinsicOp<"cos",
let assemblyFormat = "$src attr-dict `:` type($src)";
}
-def NVVM_Lg2Op : NVVM_SingleResultIntrinsicOp<"lg2",
+def NVVM_Log2Op : NVVM_SingleResultIntrinsicOp<"log2",
[Pure, SameOperandsAndResultType]> {
let summary = "Base-2 logarithm (fast approximation)";
let description = [{
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 7dda2802dd659..7770f8f2dba0a 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3496,9 +3496,9 @@ SinOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
}
mlir::NVVM::IDArgPair
-Lg2Op::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
- llvm::IRBuilderBase &builder) {
- auto thisOp = cast<NVVM::Lg2Op>(op);
+Log2Op::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
+ llvm::IRBuilderBase &builder) {
+ auto thisOp = cast<NVVM::Log2Op>(op);
llvm::Intrinsic::ID id = thisOp.getFtz()
? llvm::Intrinsic::nvvm_lg2_approx_ftz_f
: llvm::Intrinsic::nvvm_lg2_approx_f;
diff --git a/mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir b/mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir
index 54def11d02457..aa5439a70395e 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir
@@ -31,15 +31,15 @@ func.func @nvvm_sin_ftz_f32(%arg0: f32) -> f32 {
// CHECK-LABEL: @nvvm_lg2_f32
func.func @nvvm_lg2_f32(%arg0: f32) -> f32 {
- // CHECK: nvvm.lg2 {{.*}} : f32
- %0 = nvvm.lg2 %arg0 : f32
+ // CHECK: nvvm.log2 {{.*}} : f32
+ %0 = nvvm.log2 %arg0 : f32
return %0 : f32
}
// CHECK-LABEL: @nvvm_lg2_ftz_f32
func.func @nvvm_lg2_ftz_f32(%arg0: f32) -> f32 {
- // CHECK: nvvm.lg2 {{.*}} {ftz = true} : f32
- %0 = nvvm.lg2 %arg0 {ftz = true} : f32
+ // CHECK: nvvm.log2 {{.*}} {ftz = true} : f32
+ %0 = nvvm.log2 %arg0 {ftz = true} : f32
return %0 : f32
}
diff --git a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
index 0762c403b7307..5b7ecc9e1c87c 100644
--- a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
@@ -31,14 +31,14 @@ llvm.func @nvvm_cos_ftz(%arg0: f32) -> f32 {
// CHECK-LABEL: @nvvm_lg2
llvm.func @nvvm_lg2(%arg0: f32) -> f32 {
// CHECK: call float @llvm.nvvm.lg2.approx.f(float %{{.*}})
- %0 = nvvm.lg2 %arg0 : f32
+ %0 = nvvm.log2 %arg0 : f32
llvm.return %0 : f32
}
// CHECK-LABEL: @nvvm_lg2_ftz
llvm.func @nvvm_lg2_ftz(%arg0: f32) -> f32 {
// CHECK: call float @llvm.nvvm.lg2.approx.ftz.f(float %{{.*}})
- %0 = nvvm.lg2 %arg0 {ftz = true} : f32
+ %0 = nvvm.log2 %arg0 {ftz = true} : f32
llvm.return %0 : f32
}
More information about the Mlir-commits
mailing list