[MLIR][NVVM] Add log2.approx.f32 intrinsic with FTZ variant

Author: SwarnaDharshiniS

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

@@ -453,30 +453,33 @@ def NVVM_RcpApproxFtzF32Op : NVVM_IntrOp<"rcp.approx.ftz.f", [Pure], 1> {
 // NVVM lg2 approximate op
 //===----------------------------------------------------------------------===//
 
-def NVVM_Lg2ApproxF32Op : NVVM_Op<"lg2.approx.f", [Pure]> {
+def NVVM_Log2ApproxF32Op : NVVM_Op<"log2.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.
+    
+    [PTX ISA Reference for log2.approx]
+    (https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#floating-point-instructions-lg2)
+
   }];
 
   let arguments = (ins 
-    F32:$arg,
-    OptionalAttr<BoolAttr>:$ftz);
+    F32:$input,
+    DefaultValuedAttr<BoolAttr, "false">:$ftz);
 
   let results = (outs F32:$res);
 
-  let assemblyFormat = "$arg attr-dict `:` type($res)";
+  let assemblyFormat = "$input attr-dict `:` type($res)";
 
   string llvmBuilder = [{
-    bool ftz = $ftz && *$ftz;
     llvm::Intrinsic::ID intrinsicID = 
-      ftz ? llvm::Intrinsic::nvvm_lg2_approx_ftz_f
+      $ftz ? llvm::Intrinsic::nvvm_lg2_approx_ftz_f
           : llvm::Intrinsic::nvvm_lg2_approx_f;
 
-    $res = createIntrinsicCall(builder, intrinsicID, {$arg});
+    $res = createIntrinsicCall(builder, intrinsicID, {$input});
   }];
 }
 

mlir/test/Dialect/LLVMIR/nvvm.mlir

@@ -36,20 +36,6 @@ 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

mlir/test/Target/LLVMIR/nvvm/log2_approx.mlir

@@ -0,0 +1,15 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// CHECK-LABEL: define float @nvvm_log2_approx_f(float %0)
+// CHECK: call float @llvm.nvvm.lg2.approx.f(float %0)
+llvm.func @nvvm_log2_approx_f(%arg0: f32) -> f32 {
+  %0 = nvvm.log2.approx.f %arg0 : f32
+  llvm.return %0 : f32
+}
+
+// CHECK-LABEL: define float @nvvm_log2_approx_ftz_f(float %0)
+// CHECK: call float @llvm.nvvm.lg2.approx.ftz.f(float %0)
+llvm.func @nvvm_log2_approx_ftz_f(%arg0: f32) -> f32 {
+  %0 = nvvm.log2.approx.f %arg0 {ftz = true} : f32
+  llvm.return %0 : f32
+}