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:$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