[Mlir-commits] [mlir] [mlir][NVVM] Add packed EX2 operations (PR #220501)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Bastian Hagedorn (bastianhagedorn)

<details>
<summary>Changes</summary>

Add NVVM dialect operations for the public packed EX2 PTX instructions:
* `nvvm.ex2.f16x2` lowers to `ex2.approx.f16x2`
* `nvvm.ex2.bf16x2` lowers to `ex2.approx.ftz.bf16x2`
The existing scalar `nvvm.ex2` operation is unchanged. Includes LLVM IR translation tests for both overloads.

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


2 Files Affected:

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


``````````diff
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 634918b1b21f2..6b30127463839 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -512,6 +512,46 @@ def NVVM_Ex2Op : NVVM_F32UnaryApproxOp<"ex2"> {
   }];
 }
 
+def NVVM_Ex2F16x2Op
+    : NVVM_Op<"ex2.f16x2", [Pure, SameOperandsAndResultType]> {
+  let summary = "Packed f16x2 base-2 exponential (fast approximation)";
+  let description = [{
+    Computes a fast approximation of 2 raised to the power of each f16 element
+    in the input vector. This operation lowers to the `ex2.approx.f16x2` PTX
+    instruction.
+
+    For more information, see PTX ISA:
+    [ex2](https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-ex2)
+  }];
+  let arguments = (ins VectorOfLengthAndType<[2], [F16]>:$src);
+  let results = (outs VectorOfLengthAndType<[2], [F16]>:$res);
+  let assemblyFormat = "$src attr-dict `:` type($src)";
+  string llvmBuilder = [{
+    $res = createIntrinsicCall(builder, llvm::Intrinsic::nvvm_ex2_approx,
+                               {$src});
+  }];
+}
+
+def NVVM_Ex2BF16x2Op
+    : NVVM_Op<"ex2.bf16x2", [Pure, SameOperandsAndResultType]> {
+  let summary = "Packed bf16x2 base-2 exponential (fast approximation)";
+  let description = [{
+    Computes a fast approximation of 2 raised to the power of each bf16
+    element in the input vector. This operation lowers to the
+    `ex2.approx.ftz.bf16x2` PTX instruction.
+
+    For more information, see PTX ISA:
+    [ex2](https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-ex2)
+  }];
+  let arguments = (ins VectorOfLengthAndType<[2], [BF16]>:$src);
+  let results = (outs VectorOfLengthAndType<[2], [BF16]>:$res);
+  let assemblyFormat = "$src attr-dict `:` type($src)";
+  string llvmBuilder = [{
+    $res = createIntrinsicCall(builder,
+                               llvm::Intrinsic::nvvm_ex2_approx_ftz, {$src});
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // NVVM rsqrt op definitions
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
index 78d1acc190193..b6fc4eaa0270d 100644
--- a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
@@ -55,3 +55,17 @@ llvm.func @nvvm_ex2_ftz(%arg0: f32) -> f32 {
   %0 = nvvm.ex2 %arg0 ftz = true : f32
   llvm.return %0 : f32
 }
+
+// CHECK-LABEL: @nvvm_ex2_f16x2
+llvm.func @nvvm_ex2_f16x2(%arg0: vector<2xf16>) -> vector<2xf16> {
+  // CHECK: call <2 x half> @llvm.nvvm.ex2.approx.v2f16(<2 x half> %{{.*}})
+  %0 = nvvm.ex2.f16x2 %arg0 : vector<2xf16>
+  llvm.return %0 : vector<2xf16>
+}
+
+// CHECK-LABEL: @nvvm_ex2_bf16x2
+llvm.func @nvvm_ex2_bf16x2(%arg0: vector<2xbf16>) -> vector<2xbf16> {
+  // CHECK: call <2 x bfloat> @llvm.nvvm.ex2.approx.ftz.v2bf16(<2 x bfloat> %{{.*}})
+  %0 = nvvm.ex2.bf16x2 %arg0 : vector<2xbf16>
+  llvm.return %0 : vector<2xbf16>
+}

``````````

</details>


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