[Mlir-commits] [mlir] [mlir][NVVM] Reject vector<1xf32> in nvvm.ex2 (PR #231517)

[Bastian Hagedorn]

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

>From 9e7a5f7e8daffc285d0ea5be24e8b21ca7101d02 Mon Sep 17 00:00:00 2001
From: Bastian Hagedorn <bhagedorn at nvidia.com>
Date: Tue, 13 Oct 2026 08:03:46 +0000
Subject: [PATCH 1/2] [mlir][NVVM] Reject vector<1xf32> in nvvm.ex2

The vector<1xf32> form added in #220501 lowers to
llvm.nvvm.ex2.approx(.ftz).v1f32, which NVPTX cannot select. A one-lane
vector is not a distinct PTX form; producers should emit f32 instead.
Restrict nvvm.ex2 to f32, f16x2 and bf16x2 and add a negative test.

Supersedes #223622.
---
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td       | 11 ++++-------
 mlir/test/Target/LLVMIR/nvvm/ex2_invalid.mlir     |  8 ++++++++
 mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir | 14 --------------
 3 files changed, 12 insertions(+), 21 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index b15ee81291e5c..ead5c00bbbd34 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -499,18 +499,15 @@ def NVVM_Ex2Op
   let summary = "Base-2 exponential (fast approximation)";
   let description = [{
     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.
+    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<[1], [F32]>,
-               VectorOfLengthAndType<[2], [F16]>,
+    AnyTypeOf<[F32, VectorOfLengthAndType<[2], [F16]>,
                VectorOfLengthAndType<[2], [BF16]>]>:$src,
     DefaultValuedAttr<BoolAttr, "false">:$ftz);
   let results = (outs
-    AnyTypeOf<[F32, VectorOfLengthAndType<[1], [F32]>,
-               VectorOfLengthAndType<[2], [F16]>,
+    AnyTypeOf<[F32, VectorOfLengthAndType<[2], [F16]>,
                VectorOfLengthAndType<[2], [BF16]>]>:$res);
   let assemblyFormat = [{
     $src oilist(`ftz` `=` $ftz) attr-dict `:` type($src)
diff --git a/mlir/test/Target/LLVMIR/nvvm/ex2_invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/ex2_invalid.mlir
index 0ff00a258b961..6f129b690b8aa 100644
--- a/mlir/test/Target/LLVMIR/nvvm/ex2_invalid.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/ex2_invalid.mlir
@@ -15,3 +15,11 @@ llvm.func @ex2_invalid_no_ftz_bf16x2(%arg0: vector<2xbf16>) -> vector<2xbf16> {
   %0 = nvvm.ex2 %arg0 : vector<2xbf16>
   llvm.return %0 : vector<2xbf16>
 }
+
+// -----
+
+llvm.func @ex2_invalid_v1f32(%arg0: vector<1xf32>) -> vector<1xf32> {
+  // expected-error at +1 {{op operand #0 must be 32-bit float or vector of}}
+  %0 = nvvm.ex2 %arg0 : vector<1xf32>
+  llvm.return %0 : vector<1xf32>
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
index 06f5dafaf8b6a..3c86787430071 100644
--- a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
@@ -56,20 +56,6 @@ llvm.func @nvvm_ex2_ftz(%arg0: f32) -> 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> %{{.*}})

>From 80adda56fd70f566f0ecb394fed254c5351cdf19 Mon Sep 17 00:00:00 2001
From: Bastian Hagedorn <bhagedorn at nvidia.com>
Date: Wed, 14 Oct 2026 08:08:53 +0200
Subject: [PATCH 2/2] [mlir][NVVM] Spell out the allowed ex2 types in the
 negative test

---
 mlir/test/Target/LLVMIR/nvvm/ex2_invalid.mlir | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/test/Target/LLVMIR/nvvm/ex2_invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/ex2_invalid.mlir
index 6f129b690b8aa..504aac07ea9a4 100644
--- a/mlir/test/Target/LLVMIR/nvvm/ex2_invalid.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/ex2_invalid.mlir
@@ -19,7 +19,7 @@ llvm.func @ex2_invalid_no_ftz_bf16x2(%arg0: vector<2xbf16>) -> vector<2xbf16> {
 // -----
 
 llvm.func @ex2_invalid_v1f32(%arg0: vector<1xf32>) -> vector<1xf32> {
-  // expected-error at +1 {{op operand #0 must be 32-bit float or vector of}}
+  // expected-error at +1 {{op operand #0 must be 32-bit float or vector of 16-bit float values of length 2 or vector of bfloat16 type values of length 2, but got 'vector<1xf32>'}}
   %0 = nvvm.ex2 %arg0 : vector<1xf32>
   llvm.return %0 : vector<1xf32>
 }