[Mlir-commits] [mlir] 78eccec - [MLIR][NVVM] Add `nvvm.log2` OP (#193789)

[llvmlistbot at llvm.org]

Author: Guray Ozen
Date: 2026-04-27T14:20:00Z
New Revision: 78eccec0db8a7739151c6de12de0e130c820e20d

URL: https://github.com/llvm/llvm-project/commit/78eccec0db8a7739151c6de12de0e130c820e20d
DIFF: https://github.com/llvm/llvm-project/commit/78eccec0db8a7739151c6de12de0e130c820e20d.diff

LOG: [MLIR][NVVM] Add `nvvm.log2` OP (#193789)

Implement `nvvm.log2` with ftz flag

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
    mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
    mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir
    mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 7f1f9774abf52..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,6 +558,20 @@ def NVVM_CosOp : NVVM_SingleResultIntrinsicOp<"cos",
   let assemblyFormat = "$src attr-dict `:` type($src)";
 }
 
+def NVVM_Log2Op : NVVM_SingleResultIntrinsicOp<"log2",
+    [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..7770f8f2dba0a 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
+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;
+  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..aa5439a70395e 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.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.log2 {{.*}} {ftz = true} : f32
+  %0 = nvvm.log2 %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..5b7ecc9e1c87c 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.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.log2 %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 %{{.*}})