[Mlir-commits] [mlir] [MLIR][NVVM] Add `nvvm.lg2` OP (PR #193789)

[Guray Ozen]

https://github.com/grypp updated https://github.com/llvm/llvm-project/pull/193789

>From e333b00e706ae6fc26850fbda03cc75f71202d0a Mon Sep 17 00:00:00 2001
From: Guray Ozen <gozen at nvidia.com>
Date: Fri, 24 Apr 2026 12:44:06 +0200
Subject: [PATCH] [MLIR][NVVM] Add `nvvm.lg2` OP

Implement `nvvm.lg2` with ftz flag

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply at anthropic.com>
---
 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 9f01d44cadd76..875b7707ca811 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)";
+}
+
 //===----------------------------------------------------------------------===//
 // NVVM redux op definitions
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 8b4f832997a0f..90f283ba6f561 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3484,6 +3484,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
 PMEventOp::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 777152e0fd797..be8ebb60a4fa3 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm-transcendentals.mlir
@@ -28,3 +28,17 @@ func.func @nvvm_sin_ftz_f32(%arg0: f32) -> f32 {
   %0 = nvvm.sin %arg0 {ftz = true} : 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
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
index 06da78820b25f..92f8020515858 100644
--- a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
@@ -27,3 +27,17 @@ llvm.func @nvvm_cos_ftz(%arg0: f32) -> f32 {
   %0 = nvvm.cos %arg0 {ftz = true} : 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
+}