[Mlir-commits] [mlir] [MLIR][NVVM] Add lg2.approx.f32 intrinsic with optional FTZ support (PR #167330)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir

Author: Swarna Dharshini S (SwarnaDharshiniS)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/167330.diff


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td (+31) 
- (modified) mlir/test/Dialect/LLVMIR/nvvm.mlir (+14) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 4f483859ac18d..4c32f4c3a4d30 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -449,6 +449,37 @@ def NVVM_RcpApproxFtzF32Op : NVVM_IntrOp<"rcp.approx.ftz.f", [Pure], 1> {
   let assemblyFormat = "$arg attr-dict `:` type($res)";
 }
 
+//===----------------------------------------------------------------------===//
+// NVVM lg2 approximate op
+//===----------------------------------------------------------------------===//
+
+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});
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // NVVM redux op definitions
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index 0243f5eb8c862..f4fbf5d02e542 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_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_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
+}
+
 // CHECK-LABEL: @llvm_nvvm_barrier0
 func.func @llvm_nvvm_barrier0() {
   // CHECK: nvvm.barrier0

``````````

</details>


https://github.com/llvm/llvm-project/pull/167330