[Mlir-commits] [mlir] [mlir][NVVM] Lower vector<1xf32> nvvm.ex2 through the scalar intrinsic (PR #231517)
Bastian Hagedorn
llvmlistbot at llvm.org
Wed Oct 14 02:38:25 PDT 2026
https://github.com/bastianhagedorn updated https://github.com/llvm/llvm-project/pull/231517
>From 81f56f3d909461c09ce2822bdace240031313d1c Mon Sep 17 00:00:00 2001
From: Bastian Hagedorn <bhagedorn at nvidia.com>
Date: Wed, 14 Oct 2026 08:24:56 +0200
Subject: [PATCH 1/2] [mlir][NVVM] Lower vector<1xf32> nvvm.ex2 through the
scalar intrinsic
The vector<1xf32> form added in #220501 translated to
llvm.nvvm.ex2.approx(.ftz).v1f32, which NVPTX cannot select. A one-lane
vector is not a distinct PTX form, so lower it on the MLIR side: extract
the lane, call the f32 intrinsic and re-insert the result. The op keeps
accepting vector<1xf32>.
Supersedes #223622.
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 20 ++++++++++++++++++-
.../Target/LLVMIR/nvvm/transcendentals.mlir | 8 ++++++--
2 files changed, 25 insertions(+), 3 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index b15ee81291e5c..9ac2cf6156643 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -501,7 +501,8 @@ def NVVM_Ex2Op
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.
+ and required for bf16x2. A vector<1xf32> operand is not a distinct PTX
+ form: it is lowered through the scalar f32 intrinsic.
}];
let arguments = (ins
AnyTypeOf<[F32, VectorOfLengthAndType<[1], [F32]>,
@@ -516,6 +517,23 @@ def NVVM_Ex2Op
$src oilist(`ftz` `=` $ftz) attr-dict `:` type($src)
}];
let hasVerifier = 1;
+ let llvmBuilder = [{
+ auto [id, args] =
+ NVVM::Ex2Op::getIntrinsicIDAndArgs(*op, moduleTranslation, builder);
+ llvm::Value *src = args[0];
+ auto *vecTy = llvm::dyn_cast<llvm::FixedVectorType>(src->getType());
+ if (vecTy && vecTy->getNumElements() == 1) {
+ // There is no one-lane PTX form; use the scalar f32 intrinsic.
+ llvm::Value *lane = builder.CreateExtractElement(src, builder.getInt32(0));
+ llvm::Value *res = createIntrinsicCall(builder, id, lane->getType(), {lane});
+ $res = builder.CreateInsertElement(llvm::PoisonValue::get(vecTy), res,
+ builder.getInt32(0));
+ } else {
+ $res = createIntrinsicCall(
+ builder, id, moduleTranslation.convertType(op.getRes().getType()),
+ args);
+ }
+ }];
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
index 06f5dafaf8b6a..1f3aff350aa9c 100644
--- a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
@@ -58,14 +58,18 @@ llvm.func @nvvm_ex2_ftz(%arg0: f32) -> 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> %{{.*}})
+ // CHECK: %[[LANE:.*]] = extractelement <1 x float> %{{.*}}, i32 0
+ // CHECK: %[[RES:.*]] = call float @llvm.nvvm.ex2.approx.f32(float %[[LANE]])
+ // CHECK: insertelement <1 x float> poison, float %[[RES]], i32 0
%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> %{{.*}})
+ // CHECK: %[[LANE:.*]] = extractelement <1 x float> %{{.*}}, i32 0
+ // CHECK: %[[RES:.*]] = call float @llvm.nvvm.ex2.approx.ftz.f32(float %[[LANE]])
+ // CHECK: insertelement <1 x float> poison, float %[[RES]], i32 0
%0 = nvvm.ex2 %arg0 ftz = true : vector<1xf32>
llvm.return %0 : vector<1xf32>
}
>From 6ab28abd319b3f6bec8f3c9fb65d30dbc663f36f Mon Sep 17 00:00:00 2001
From: Bastian Hagedorn <bhagedorn at nvidia.com>
Date: Wed, 14 Oct 2026 11:38:11 +0200
Subject: [PATCH 2/2] [mlir][NVVM] Use bitcasts for the vector<1xf32> ex2
lowering
Address review: a bitcast between <1 x float> and float is simpler than
extract/insert.
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 8 ++++----
mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir | 12 ++++++------
2 files changed, 10 insertions(+), 10 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 9ac2cf6156643..47675b6f61c18 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -524,10 +524,10 @@ def NVVM_Ex2Op
auto *vecTy = llvm::dyn_cast<llvm::FixedVectorType>(src->getType());
if (vecTy && vecTy->getNumElements() == 1) {
// There is no one-lane PTX form; use the scalar f32 intrinsic.
- llvm::Value *lane = builder.CreateExtractElement(src, builder.getInt32(0));
- llvm::Value *res = createIntrinsicCall(builder, id, lane->getType(), {lane});
- $res = builder.CreateInsertElement(llvm::PoisonValue::get(vecTy), res,
- builder.getInt32(0));
+ llvm::Type *f32Ty = vecTy->getElementType();
+ llvm::Value *scalar = builder.CreateBitCast(src, f32Ty);
+ llvm::Value *res = createIntrinsicCall(builder, id, f32Ty, {scalar});
+ $res = builder.CreateBitCast(res, vecTy);
} else {
$res = createIntrinsicCall(
builder, id, moduleTranslation.convertType(op.getRes().getType()),
diff --git a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
index 1f3aff350aa9c..d5106857ed6ba 100644
--- a/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/transcendentals.mlir
@@ -58,18 +58,18 @@ llvm.func @nvvm_ex2_ftz(%arg0: f32) -> f32 {
// CHECK-LABEL: @nvvm_ex2_v1f32
llvm.func @nvvm_ex2_v1f32(%arg0: vector<1xf32>) -> vector<1xf32> {
- // CHECK: %[[LANE:.*]] = extractelement <1 x float> %{{.*}}, i32 0
- // CHECK: %[[RES:.*]] = call float @llvm.nvvm.ex2.approx.f32(float %[[LANE]])
- // CHECK: insertelement <1 x float> poison, float %[[RES]], i32 0
+ // CHECK: %[[SRC:.*]] = bitcast <1 x float> %{{.*}} to float
+ // CHECK: %[[RES:.*]] = call float @llvm.nvvm.ex2.approx.f32(float %[[SRC]])
+ // CHECK: bitcast float %[[RES]] to <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: %[[LANE:.*]] = extractelement <1 x float> %{{.*}}, i32 0
- // CHECK: %[[RES:.*]] = call float @llvm.nvvm.ex2.approx.ftz.f32(float %[[LANE]])
- // CHECK: insertelement <1 x float> poison, float %[[RES]], i32 0
+ // CHECK: %[[SRC:.*]] = bitcast <1 x float> %{{.*}} to float
+ // CHECK: %[[RES:.*]] = call float @llvm.nvvm.ex2.approx.ftz.f32(float %[[SRC]])
+ // CHECK: bitcast float %[[RES]] to <1 x float>
%0 = nvvm.ex2 %arg0 ftz = true : vector<1xf32>
llvm.return %0 : vector<1xf32>
}
More information about the Mlir-commits
mailing list