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

[llvmlistbot at llvm.org]

Author: Bastian Hagedorn
Date: 2026-09-14T14:19:34+02:00
New Revision: bc5c36113551b13e41606573da58334ea17db1ab

URL: https://github.com/llvm/llvm-project/commit/bc5c36113551b13e41606573da58334ea17db1ab
DIFF: https://github.com/llvm/llvm-project/commit/bc5c36113551b13e41606573da58334ea17db1ab.diff

LOG: [mlir][NVVM] Add packed EX2 operations (#220501)

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.

Added: 
    mlir/test/Target/LLVMIR/nvvm/ex2_invalid.mlir

Modified: 
    mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
    mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
    mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index c5e99d999920d..3226063edeb88 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -494,13 +494,28 @@ 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, vector<1xf32>, and packed f16x2 and bf16x2 values. The `ftz`
+    attribute is optional for f32 and vector<1xf32>, unavailable for f16x2,
+    and required for bf16x2.
   }];
+  let arguments = (ins
+    AnyTypeOf<[F32, VectorOfLengthAndType<[1], [F32]>,
+               VectorOfLengthAndType<[2], [F16]>,
+               VectorOfLengthAndType<[2], [BF16]>]>:$src,
+    DefaultValuedAttr<BoolAttr, "false">:$ftz);
+  let results = (outs
+    AnyTypeOf<[F32, VectorOfLengthAndType<[1], [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 688b3b0cd468a..2c0eb4b694c12 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3930,6 +3930,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
 RsqrtOp::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 78d1acc190193..06f5dafaf8b6a 100644
--- a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
@@ -55,3 +55,31 @@ llvm.func @nvvm_ex2_ftz(%arg0: f32) -> f32 {
   %0 = nvvm.ex2 %arg0 ftz = true : f32
   llvm.return %0 : f32
 }
+
+// CHECK-LABEL: @nvvm_ex2_v1f32
+llvm.func @nvvm_ex2_v1f32(%arg0: vector<1xf32>) -> vector<1xf32> {
+  // CHECK: call <1 x float> @llvm.nvvm.ex2.approx.v1f32(<1 x float> %{{.*}})
+  %0 = nvvm.ex2 %arg0 : vector<1xf32>
+  llvm.return %0 : vector<1xf32>
+}
+
+// CHECK-LABEL: @nvvm_ex2_v1f32_ftz
+llvm.func @nvvm_ex2_v1f32_ftz(%arg0: vector<1xf32>) -> vector<1xf32> {
+  // CHECK: call <1 x float> @llvm.nvvm.ex2.approx.ftz.v1f32(<1 x float> %{{.*}})
+  %0 = nvvm.ex2 %arg0 ftz = true : vector<1xf32>
+  llvm.return %0 : vector<1xf32>
+}
+
+// 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>
+}