[MLIR][NVVM] Add lg2.approx.f32 intrinsic with optional FTZ support

Author: SwarnaDharshiniS

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

@@ -449,16 +449,35 @@ def NVVM_RcpApproxFtzF32Op : NVVM_IntrOp<"rcp.approx.ftz.f", [Pure], 1> {
   let assemblyFormat = "$arg attr-dict `:` type($res)";
 }
 
-def NVVM_lg2ApproxF32Op : NVVM_IntrOp<"lg2.approx.f", [Pure], 1> {
-  let arguments = (ins F32:$arg);
-  let results = (outs F32:$res);
-  let assemblyFormat = "$arg attr-dict `:` type($res)";
-}
+//===----------------------------------------------------------------------===//
+// NVVM lg2 approximate op
+//===----------------------------------------------------------------------===//
 
-def NVVM_lg2ApproxFtzF32Op : NVVM_IntrOp<"lg2.approx.ftz.f", [Pure], 1> {
-  let arguments = (ins F32:$arg);
+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});
+  }];
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/LLVMIR/nvvm.mlir

@@ -36,17 +36,17 @@ func.func @nvvm_rcp(%arg0: f32) -> f32 {
   llvm.return %0 : f32
 }
 
-// CHECK-LABEL: @nvvm_lg2_f32
-func.func @nvvm_lg2_f32(%arg0: f32) -> 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_ftz_f32
-func.func @nvvm_lg2_ftz_f32(%arg0: f32) -> f32 {
-  // CHECK: nvvm.lg2.approx.ftz.f %arg0 : f32
-  %0 = nvvm.lg2.approx.ftz.f %arg0 : 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
 }