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

[Bastian Hagedorn]

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

>From a58764a43a17bfc8be067bed0790b73b331ed5f6 Mon Sep 17 00:00:00 2001
From: Bastian Hagedorn <bhagedorn at nvidia.com>
Date: Fri, 11 Sep 2026 12:14:26 +0000
Subject: [PATCH] [mlir][NVVM] Add packed EX2 operations

---
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td   | 20 +++++++++++++++----
 mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp    | 12 +++++++++++
 mlir/test/Target/LLVMIR/nvvm/ex2_invalid.mlir | 17 ++++++++++++++++
 .../Target/LLVMIR/nvvm/transcendentals.mlir   | 14 +++++++++++++
 4 files changed, 59 insertions(+), 4 deletions(-)
 create mode 100644 mlir/test/Target/LLVMIR/nvvm/ex2_invalid.mlir

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 51396947fad4e..88b4e53e2d74b 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -569,13 +569,25 @@ def NVVM_Log2Op : NVVM_F32UnaryApproxOp<"log2"> {
   }];
 }
 
-def NVVM_Ex2Op : NVVM_F32UnaryApproxOp<"ex2"> {
+def NVVM_Ex2Op
+    : NVVM_SingleResultIntrinsicOp<"ex2", [Pure, SameOperandsAndResultType]> {
   let summary = "Base-2 exponential (fast approximation)";
   let description = [{
-    Computes a fast approximation of 2 raised to the power of the input
-    value. The `ftz` attribute, when set, flushes subnormal inputs and
-    results to sign-preserving zero.
+    Computes a fast approximation of 2 raised to the power of the input value.
+    Supports f32 and packed f16x2 and bf16x2 values. The `ftz` attribute is
+    optional for f32, unavailable for f16x2, and required for bf16x2.
+  }];
+  let arguments = (ins
+    AnyTypeOf<[F32, VectorOfLengthAndType<[2], [F16]>,
+               VectorOfLengthAndType<[2], [BF16]>]>:$src,
+    DefaultValuedAttr<BoolAttr, "false">:$ftz);
+  let results = (outs
+    AnyTypeOf<[F32, VectorOfLengthAndType<[2], [F16]>,
+               VectorOfLengthAndType<[2], [BF16]>]>:$res);
+  let assemblyFormat = [{
+    $src oilist(`ftz` `=` $ftz) attr-dict `:` type($src)
   }];
+  let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 7770f8f2dba0a..dfd688cc84b52 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3515,6 +3515,18 @@ Ex2Op::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
   return {id, {mt.lookupValue(thisOp.getSrc())}};
 }
 
+LogicalResult NVVM::Ex2Op::verify() {
+  auto vectorType = dyn_cast<VectorType>(getSrc().getType());
+  if (!vectorType)
+    return success();
+
+  if (vectorType.getElementType().isF16() && getFtz())
+    return emitOpError("FTZ is not supported for vector<2xf16>");
+  if (vectorType.getElementType().isBF16() && !getFtz())
+    return emitOpError("FTZ is required for vector<2xbf16>");
+  return success();
+}
+
 mlir::NVVM::IDArgPair
 PMEventOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
                                  llvm::IRBuilderBase &builder) {
diff --git a/mlir/test/Target/LLVMIR/nvvm/ex2_invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/ex2_invalid.mlir
new file mode 100644
index 0000000000000..0ff00a258b961
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/ex2_invalid.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-translate --mlir-to-llvmir --split-input-file --verify-diagnostics %s
+
+// -----
+
+llvm.func @ex2_invalid_ftz_f16x2(%arg0: vector<2xf16>) -> vector<2xf16> {
+  // expected-error at +1 {{FTZ is not supported for vector<2xf16>}}
+  %0 = nvvm.ex2 %arg0 ftz = true : vector<2xf16>
+  llvm.return %0 : vector<2xf16>
+}
+
+// -----
+
+llvm.func @ex2_invalid_no_ftz_bf16x2(%arg0: vector<2xbf16>) -> vector<2xbf16> {
+  // expected-error at +1 {{FTZ is required for vector<2xbf16>}}
+  %0 = nvvm.ex2 %arg0 : vector<2xbf16>
+  llvm.return %0 : vector<2xbf16>
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
index 5b7ecc9e1c87c..f4d27322e36c8 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 %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 %arg0 ftz = true : vector<2xbf16>
+  llvm.return %0 : vector<2xbf16>
+}