[Mlir-commits] [mlir] [MLIR][NVGPU] Add convert.fpext and convert.fptrunc Ops (PR #199700)
Srinivasa Ravi
llvmlistbot at llvm.org
Wed May 27 07:07:25 PDT 2026
================
@@ -0,0 +1,327 @@
+// RUN: mlir-opt %s -convert-nvgpu-to-nvvm | FileCheck %s
+// RUN: mlir-opt %s -convert-nvgpu-to-nvvm -convert-vector-to-llvm | FileCheck %s --check-prefix=CHECK-E2E
+
+// Basic aligned vector inputs.
+
+// CHECK-LABEL: @cvt_float_f32_to_f16(
+// CHECK-SAME: %[[IN:.+]]: vector<4xf32>
+func.func @cvt_float_f32_to_f16(%in : vector<4xf32>) {
+ // CHECK: llvm.bitcast %[[IN]] : vector<4xf32> to vector<4xi32>
+ // CHECK: nvvm.convert.f32x2.to.f16x2
+ // CHECK-SAME: rnd = #nvvm.fp_rnd_mode<rn>
+ // CHECK-SAME: : vector<2xf16>
+ // CHECK: nvvm.convert.f32x2.to.f16x2
+ // CHECK-SAME: rnd = #nvvm.fp_rnd_mode<rn>
+ // CHECK-SAME: : vector<2xf16>
+ // CHECK: llvm.bitcast {{.*}} : vector<2xi32> to vector<4xf16>
+ %out = nvgpu.convert.fptrunc %in : vector<4xf32> to vector<4xf16>
+ return
+}
+
+// CHECK-LABEL: @cvt_float_f32_to_f16_v8(
+// CHECK-SAME: %[[IN:.+]]: vector<8xf32>
+func.func @cvt_float_f32_to_f16_v8(%in : vector<8xf32>) {
+ // CHECK: llvm.bitcast %[[IN]] : vector<8xf32> to vector<8xi32>
+ // CHECK: nvvm.convert.f32x2.to.f16x2
+ // CHECK: nvvm.convert.f32x2.to.f16x2
+ // CHECK: nvvm.convert.f32x2.to.f16x2
+ // CHECK: nvvm.convert.f32x2.to.f16x2
+ // CHECK: llvm.bitcast {{.*}} : vector<4xi32> to vector<8xf16>
+ %out = nvgpu.convert.fptrunc %in : vector<8xf32> to vector<8xf16>
+ return
+}
+
+// CHECK-LABEL: @cvt_float_f32_to_bf16(
+// CHECK-SAME: %[[IN:.+]]: vector<4xf32>
+func.func @cvt_float_f32_to_bf16(%in : vector<4xf32>) {
+ // CHECK: llvm.bitcast %[[IN]] : vector<4xf32> to vector<4xi32>
+ // CHECK: nvvm.convert.f32x2.to.bf16x2
+ // CHECK-SAME: rnd = #nvvm.fp_rnd_mode<rn>
+ // CHECK-SAME: : vector<2xbf16>
+ // CHECK: nvvm.convert.f32x2.to.bf16x2
+ // CHECK: llvm.bitcast {{.*}} : vector<2xi32> to vector<4xbf16>
+ %out = nvgpu.convert.fptrunc %in : vector<4xf32> to vector<4xbf16>
+ return
+}
+
+// CHECK-LABEL: @cvt_float_f32_to_e4m3(
+// CHECK-SAME: %[[IN:.+]]: vector<8xf32>
+func.func @cvt_float_f32_to_e4m3(%in : vector<8xf32>) {
+ // CHECK: %[[IN_I32:.+]] = llvm.bitcast %[[IN]] : vector<8xf32> to vector<8xi32>
+ // CHECK: %[[OUT_I32:.+]] = llvm.mlir.undef : vector<2xi32>
+ // CHECK: %[[IDX_0:.+]] = llvm.mlir.constant(0 : i64) : i64
+ // CHECK: %[[SUB_VEC_0:.+]] = llvm.mlir.undef : vector<2xi16>
+ // CHECK: nvvm.convert.f32x2.to.f8x2
+ // CHECK-SAME: sat = #nvvm.sat_mode<satfinite>
+ // CHECK-SAME: : i16(f8E4M3FN)
+ // CHECK: llvm.insertelement {{.*}} : vector<2xi16>
+ // CHECK: nvvm.convert.f32x2.to.f8x2
+ // CHECK-SAME: : i16(f8E4M3FN)
+ // CHECK: llvm.insertelement {{.*}} : vector<2xi16>
+ // CHECK: llvm.bitcast {{.*}} : vector<2xi16> to i32
+ // CHECK: llvm.insertelement {{.*}} : vector<2xi32>
+ // CHECK: llvm.mlir.undef : vector<2xi16>
+ // CHECK: nvvm.convert.f32x2.to.f8x2
+ // CHECK: nvvm.convert.f32x2.to.f8x2
+ // CHECK: llvm.bitcast {{.*}} : vector<2xi16> to i32
+ // CHECK: llvm.insertelement {{.*}} : vector<2xi32>
+ // CHECK: llvm.bitcast {{.*}} : vector<2xi32> to vector<8xi8>
+ %out = nvgpu.convert.fptrunc %in : vector<8xf32> to vector<8xf8E4M3FN>
+ return
+}
+
+// CHECK-LABEL: @cvt_float_f32_to_e2m3(
+// CHECK-SAME: %[[IN:.+]]: vector<8xf32>
+func.func @cvt_float_f32_to_e2m3(%in : vector<8xf32>) {
+ // CHECK: llvm.bitcast %[[IN]] : vector<8xf32> to vector<8xi32>
+ // CHECK: llvm.mlir.undef : vector<2xi32>
+ // CHECK: llvm.mlir.undef : vector<2xi16>
+ // CHECK: nvvm.convert.f32x2.to.f6x2
+ // CHECK-SAME: : i16(f6E2M3FN)
+ // CHECK: llvm.insertelement {{.*}} : vector<2xi16>
+ // CHECK: nvvm.convert.f32x2.to.f6x2
+ // CHECK-SAME: : i16(f6E2M3FN)
+ // CHECK: llvm.insertelement {{.*}} : vector<2xi16>
+ // CHECK: llvm.bitcast {{.*}} : vector<2xi16> to i32
+ // CHECK: llvm.insertelement {{.*}} : vector<2xi32>
+ // CHECK: llvm.mlir.undef : vector<2xi16>
+ // CHECK: nvvm.convert.f32x2.to.f6x2
+ // CHECK: nvvm.convert.f32x2.to.f6x2
+ // CHECK: llvm.bitcast {{.*}} : vector<2xi16> to i32
+ // CHECK: llvm.insertelement {{.*}} : vector<2xi32>
+ // CHECK: llvm.bitcast {{.*}} : vector<2xi32> to vector<8xi8>
+ %out = nvgpu.convert.fptrunc %in {packed_kind = #nvgpu.subbytes_packedkind<u6_unpack_u8_e2m3>}: vector<8xf32> to vector<8xi8>
+ return
+}
+
+// Scalar inputs (canonicalize: broadcast + pad + extract).
+
+// CHECK-LABEL: @fptrunc_scalar_f32_to_f16
+// CHECK-SAME: %[[IN:.+]]: f32
+func.func @fptrunc_scalar_f32_to_f16(%in : f32) -> f16 {
+ // CHECK: vector.broadcast %[[IN]] : f32 to vector<1xf32>
+ // CHECK: vector.insert_strided_slice
+ // CHECK: nvvm.convert.f32x2.to.f16x2
+ // CHECK: vector.extract_strided_slice
+ // CHECK: vector.extract {{.*}}[0] : f16 from vector<1xf16>
+ %out = nvgpu.convert.fptrunc %in : f32 to f16
+ return %out : f16
+}
+
+// CHECK-LABEL: @fptrunc_scalar_f32_to_bf16
+// CHECK-SAME: %[[IN:.+]]: f32
+func.func @fptrunc_scalar_f32_to_bf16(%in : f32) -> bf16 {
+ // CHECK: vector.broadcast %[[IN]]
+ // CHECK: nvvm.convert.f32x2.to.bf16x2
+ // CHECK: vector.extract
+ %out = nvgpu.convert.fptrunc %in : f32 to bf16
+ return %out : bf16
+}
+
+// CHECK-LABEL: @fptrunc_scalar_f64_to_f32
+func.func @fptrunc_scalar_f64_to_f32(%arg0: f64) -> f32 {
+ // CHECK: vector.broadcast
+ // CHECK: llvm.fptrunc
+ // CHECK: vector.extract
+ %out = nvgpu.convert.fptrunc %arg0 : f64 to f32
+ return %out : f32
+}
+
+// Multi-rank vectors (canonicalize: shape_cast flatten).
+
+// CHECK-LABEL: @fptrunc_v2x4_f32_to_f16
+// CHECK-SAME: %[[IN:.+]]: vector<2x4xf32>
+func.func @fptrunc_v2x4_f32_to_f16(%in : vector<2x4xf32>) -> vector<2x4xf16> {
+ // CHECK: vector.shape_cast %[[IN]] : vector<2x4xf32> to vector<8xf32>
+ // CHECK: nvvm.convert.f32x2.to.f16x2
+ // CHECK: vector.shape_cast {{.*}} to vector<2x4xf16>
+ %out = nvgpu.convert.fptrunc %in : vector<2x4xf32> to vector<2x4xf16>
+ return %out : vector<2x4xf16>
+}
+
+// CHECK-LABEL: @fptrunc_v4x2_f32_to_f8
+// CHECK-SAME: %[[IN:.+]]: vector<4x2xf32>
+func.func @fptrunc_v4x2_f32_to_f8(%in : vector<4x2xf32>) -> vector<4x2xf8E4M3FN> {
+ // CHECK: vector.shape_cast %[[IN]] : vector<4x2xf32> to vector<8xf32>
+ // CHECK: nvvm.convert.f32x2.to.f8x2
+ // CHECK: vector.shape_cast {{.*}} to vector<4x2xf8E4M3FN>
+ %out = nvgpu.convert.fptrunc %in : vector<4x2xf32> to vector<4x2xf8E4M3FN>
+ return %out : vector<4x2xf8E4M3FN>
+}
+
+// Non-aligned 1-D vectors (canonicalize: pad via insert/extract_strided_slice).
+
+// CHECK-LABEL: @fptrunc_v1f32_to_v1f16
+// CHECK-SAME: %[[IN:.+]]: vector<1xf32>
+func.func @fptrunc_v1f32_to_v1f16(%in : vector<1xf32>) -> vector<1xf16> {
+ // CHECK: vector.insert_strided_slice %[[IN]]
+ // CHECK: nvvm.convert.f32x2.to.f16x2
+ // CHECK: vector.extract_strided_slice
+ %out = nvgpu.convert.fptrunc %in : vector<1xf32> to vector<1xf16>
+ return %out : vector<1xf16>
+}
+
+// CHECK-LABEL: @fptrunc_v3f16_to_v3f8
+// CHECK-SAME: %[[IN:.+]]: vector<3xf16>
+func.func @fptrunc_v3f16_to_v3f8(%in : vector<3xf16>) -> vector<3xf8E4M3FN> {
+ // CHECK: vector.insert_strided_slice %[[IN]]
+ // CHECK: nvvm.convert.f16x2.to.f8x2
+ // CHECK: vector.extract_strided_slice
+ %out = nvgpu.convert.fptrunc %in : vector<3xf16> to vector<3xf8E4M3FN>
+ return %out : vector<3xf8E4M3FN>
+}
+
+// Multi-rank + padding combined.
+
+// CHECK-LABEL: @fptrunc_v3x1_f32_to_f16
+// CHECK-SAME: %[[IN:.+]]: vector<3x1xf32>
+func.func @fptrunc_v3x1_f32_to_f16(%in : vector<3x1xf32>) -> vector<3x1xf16> {
+ // CHECK: vector.shape_cast %[[IN]] : vector<3x1xf32> to vector<3xf32>
+ // CHECK: vector.insert_strided_slice
+ // CHECK: nvvm.convert.f32x2.to.f16x2
+ // CHECK: vector.extract_strided_slice
+ // CHECK: vector.shape_cast {{.*}} to vector<3x1xf16>
+ %out = nvgpu.convert.fptrunc %in : vector<3x1xf32> to vector<3x1xf16>
+ return %out : vector<3x1xf16>
+}
+
+// f64 source truncation.
+
+// CHECK-LABEL: @fptrunc_f64_to_f32
+func.func @fptrunc_f64_to_f32(%arg0: vector<4xf64>) -> vector<4xf32> {
+ // CHECK: llvm.fptrunc %{{.*}} : vector<4xf64> to vector<4xf32>
+ // CHECK-NOT: nvvm
+ %out = nvgpu.convert.fptrunc %arg0 : vector<4xf64> to vector<4xf32>
+ return %out : vector<4xf32>
+}
+
+// CHECK-LABEL: @fptrunc_f64_to_f16
+func.func @fptrunc_f64_to_f16(%arg0: vector<2xf64>) -> vector<2xf16> {
+ // CHECK: vector.insert_strided_slice
+ // CHECK: llvm.fptrunc %{{.*}} : vector<4xf64> to vector<4xf32>
+ // CHECK: nvvm.convert.f32x2.to.f16x2
+ // CHECK: vector.extract_strided_slice
+ %out = nvgpu.convert.fptrunc %arg0 : vector<2xf64> to vector<2xf16>
+ return %out : vector<2xf16>
+}
+
+// CHECK-LABEL: @fptrunc_f64_to_bf16
+func.func @fptrunc_f64_to_bf16(%arg0: vector<2xf64>) -> vector<2xbf16> {
+ // CHECK: vector.insert_strided_slice
+ // CHECK: llvm.fptrunc %{{.*}} : vector<4xf64> to vector<4xf32>
+ // CHECK: nvvm.convert.f32x2.to.bf16x2
+ // CHECK: vector.extract_strided_slice
+ %out = nvgpu.convert.fptrunc %arg0 : vector<2xf64> to vector<2xbf16>
+ return %out : vector<2xbf16>
+}
+
+// CHECK-LABEL: @fptrunc_f64_to_f8
+func.func @fptrunc_f64_to_f8(%arg0: vector<4xf64>) -> vector<4xf8E4M3FN> {
+ // CHECK: vector.insert_strided_slice
+ // CHECK: llvm.fptrunc %{{.*}} : vector<8xf64> to vector<8xf32>
+ // CHECK: nvvm.convert.f32x2.to.f8x2
+ // CHECK: vector.extract_strided_slice
+ %out = nvgpu.convert.fptrunc %arg0 : vector<4xf64> to vector<4xf8E4M3FN>
+ return %out : vector<4xf8E4M3FN>
+}
+
+// Saturation and relu attributes.
+
+// CHECK-LABEL: @fptrunc_f32_to_f8_satfinite
+func.func @fptrunc_f32_to_f8_satfinite(%in : vector<8xf32>) {
+ // CHECK: nvvm.convert.f32x2.to.f8x2
+ // CHECK-SAME: sat = #nvvm.sat_mode<satfinite>
+ %out = nvgpu.convert.fptrunc %in {sat = #nvvm.sat_mode<satfinite>}
+ : vector<8xf32> to vector<8xf8E4M3FN>
+ return
+}
+
+// CHECK-LABEL: @fptrunc_f32_to_f16_relu
+func.func @fptrunc_f32_to_f16_relu(%in : vector<4xf32>) {
+ // CHECK: nvvm.convert.f32x2.to.f16x2
+ // CHECK-SAME: relu = true
+ %out = nvgpu.convert.fptrunc %in {relu = true}
+ : vector<4xf32> to vector<4xf16>
+ return
+}
+
+// Default SATFINITE behavior.
+
+// CHECK-LABEL: @fptrunc_f32_to_f8_default_sat
+func.func @fptrunc_f32_to_f8_default_sat(%in : vector<8xf32>) {
+ // CHECK: nvvm.convert.f32x2.to.f8x2
+ // CHECK-SAME: sat = #nvvm.sat_mode<satfinite>
+ %out = nvgpu.convert.fptrunc %in : vector<8xf32> to vector<8xf8E4M3FN>
+ return
+}
+
+// CHECK-LABEL: @fptrunc_f32_to_f16_default_sat
+func.func @fptrunc_f32_to_f16_default_sat(%in : vector<4xf32>) {
+ // CHECK: nvvm.convert.f32x2.to.f16x2
+ // CHECK-SAME: sat = #nvvm.sat_mode<satfinite>
+ %out = nvgpu.convert.fptrunc %in : vector<4xf32> to vector<4xf16>
+ return
+}
+
+// CHECK-LABEL: @fptrunc_f32_to_f16_explicit_none
+func.func @fptrunc_f32_to_f16_explicit_none(%in : vector<4xf32>) {
+ // CHECK: nvvm.convert.f32x2.to.f16x2
+ // CHECK-NOT: satfinite
+ %out = nvgpu.convert.fptrunc %in {sat = #nvvm.sat_mode<none>}
+ : vector<4xf32> to vector<4xf16>
+ return
+}
+
+// Stochastic rounding (RS + random_bits).
+
+// CHECK-LABEL: @fptrunc_f32_to_f16_rs
+func.func @fptrunc_f32_to_f16_rs(%in : vector<4xf32>, %rbits : i32) {
+ // CHECK: nvvm.convert.f32x2.to.f16x2
+ // CHECK-SAME: rnd = #nvvm.fp_rnd_mode<rs>
+ %out = nvgpu.convert.fptrunc %in, %rbits {rnd = #nvvm.fp_rnd_mode<rs>}
+ : vector<4xf32> to vector<4xf16>
+ return
+}
+
+// CHECK-LABEL: @fptrunc_f32_to_bf16_rs
+func.func @fptrunc_f32_to_bf16_rs(%in : vector<4xf32>, %rbits : i32) {
+ // CHECK: nvvm.convert.f32x2.to.bf16x2
+ // CHECK-SAME: rnd = #nvvm.fp_rnd_mode<rs>
+ %out = nvgpu.convert.fptrunc %in, %rbits {rnd = #nvvm.fp_rnd_mode<rs>}
+ : vector<4xf32> to vector<4xbf16>
----------------
Wolfram70 wrote:
Added large vector tests for both Ops in new files. Thanks!
https://github.com/llvm/llvm-project/pull/199700
More information about the Mlir-commits
mailing list