128 lines · plain
1// RUN: mlir-opt %s -convert-amdgpu-to-rocdl=chipset=gfx950 | FileCheck %s2// RUN: mlir-opt %s -convert-amdgpu-to-rocdl=chipset=gfx1200 | FileCheck %s3 4// CHECK-LABEL: func @ext_scalar5// CHECK: [[V:%.+]] = builtin.unrealized_conversion_cast %{{.+}} : f8E5M2 to i86// CHECK-DAG: [[UNDEF:%.+]] = llvm.mlir.undef : vector<4xi8>7// CHECK-DAG: [[C0_1:%.+]] = llvm.mlir.constant(0 : i32) : i328// CHECK: [[VEC:%.+]] = llvm.insertelement [[V]], [[UNDEF]]{{\[}}[[C0_1]] : i32] : vector<4xi8>9// CHECK: [[CAST:%.+]] = llvm.bitcast [[VEC]] : vector<4xi8> to i3210// CHECK: [[EXT:%.+]] = rocdl.cvt.f32.bf8 [[CAST]][0] : f3211// CHECK: return [[EXT]] : f3212func.func @ext_scalar(%v: f8E5M2) -> f32 {13 %ret = amdgpu.ext_packed_fp8 %v[0] : f8E5M2 to f3214 func.return %ret : f3215}16 17// CHECK-LABEL: func @ext_short_vec18// CHECK: [[V:%.+]] = builtin.unrealized_conversion_cast %{{.+}} : vector<2xf8E4M3FN> to vector<2xi8>19// CHECK-DAG: [[UNDEF:%.+]] = llvm.mlir.undef : vector<4xi8>20// CHECK-DAG: [[C0:%.+]] = llvm.mlir.constant(0 : i32) : i3221// CHECK: [[ELEM_0:%.+]] = llvm.extractelement [[V]]{{\[}}[[C0]] : i32] : vector<2xi8>22// CHECK: [[VEC_0:%.+]] = llvm.insertelement [[ELEM_0]], [[UNDEF]]{{\[}}[[C0]] : i32] : vector<4xi8>23// CHECK: [[C1_1:%.+]] = llvm.mlir.constant(1 : i32) : i3224// CHECK: [[ELEM_1:%.+]] = llvm.extractelement [[V]]{{\[}}[[C1_1]] : i32] : vector<2xi8>25// CHECK: [[VEC_1:%.+]] = llvm.insertelement [[ELEM_1]], [[VEC_0]]{{\[}}[[C1_1]] : i32] : vector<4xi8>26// CHECK: [[CAST:%.+]] = llvm.bitcast [[VEC_1]] : vector<4xi8> to i3227// CHECK: [[EXT:%.+]] = rocdl.cvt.f32.fp8 [[CAST]][1] : f3228// CHECK: return [[EXT]] : f3229func.func @ext_short_vec(%v: vector<2xf8E4M3FN>) -> f32 {30 %ret = amdgpu.ext_packed_fp8 %v[1] : vector<2xf8E4M3FN> to f3231 func.return %ret : f3232}33 34// CHECK-LABEL: func @ext_full_vec(35// CHECK: [[V:%.+]] = builtin.unrealized_conversion_cast %{{.+}} : vector<4xf8E4M3FN> to vector<4xi8>36// CHECK: [[CAST:%.+]] = llvm.bitcast [[V]] : vector<4xi8> to i3237// CHECK: [[EXT:%.+]] = rocdl.cvt.f32.fp8 [[CAST]][3] : f3238// CHECK: return [[EXT]] : f3239func.func @ext_full_vec(%v: vector<4xf8E4M3FN>) -> f32 {40 %ret = amdgpu.ext_packed_fp8 %v[3] : vector<4xf8E4M3FN> to f3241 func.return %ret : f3242}43 44// CHECK-LABEL: func @ext_packed_2xfp845// CHECK: [[V:%.+]] = builtin.unrealized_conversion_cast %{{.+}} : vector<2xf8E4M3FN> to vector<2xi8>46// CHECK-DAG: [[UNDEF:%.+]] = llvm.mlir.undef : vector<4xi8>47// CHECK-DAG: [[C0:%.+]] = llvm.mlir.constant(0 : i32) : i3248// CHECK: [[ELEM_0:%.+]] = llvm.extractelement [[V]]{{\[}}[[C0]] : i32] : vector<2xi8>49// CHECK: [[VEC_0:%.+]] = llvm.insertelement [[ELEM_0]], [[UNDEF]]{{\[}}[[C0]] : i32] : vector<4xi8>50// CHECK: [[C1_1:%.+]] = llvm.mlir.constant(1 : i32) : i3251// CHECK: [[ELEM_1:%.+]] = llvm.extractelement [[V]]{{\[}}[[C1_1]] : i32] : vector<2xi8>52// CHECK: [[VEC_1:%.+]] = llvm.insertelement [[ELEM_1]], [[VEC_0]]{{\[}}[[C1_1]] : i32] : vector<4xi8>53// CHECK: [[CAST:%.+]] = llvm.bitcast [[VEC_1]] : vector<4xi8> to i3254// CHECK: [[EXT:%.+]] = rocdl.cvt.pk.f32.fp8 [[CAST]][false] : vector<2xf32>55// CHECK: return [[EXT]]56func.func @ext_packed_2xfp8(%v: vector<2xf8E4M3FN>) -> vector<2xf32> {57 %ret = amdgpu.ext_packed_fp8 %v[0] : vector<2xf8E4M3FN> to vector<2xf32>58 func.return %ret : vector<2xf32>59}60 61// CHECK-LABEL: func @ext_packed_4xfp862// CHECK: [[V:%.+]] = builtin.unrealized_conversion_cast %{{.+}} : vector<4xf8E4M3FN> to vector<4xi8>63// CHECK: [[CAST:%.+]] = llvm.bitcast [[V]] : vector<4xi8> to i3264// CHECK: [[EXT:%.+]] = rocdl.cvt.pk.f32.fp8 [[CAST]][true] : vector<2xf32>65// CHECK: return [[EXT]] : vector<2xf32>66func.func @ext_packed_4xfp8(%v: vector<4xf8E4M3FN>) -> vector<2xf32> {67 %ret = amdgpu.ext_packed_fp8 %v[1] : vector<4xf8E4M3FN> to vector<2xf32>68 func.return %ret : vector<2xf32>69}70 71// CHECK-LABEL: func @packed_trunc72// CHECK-SAME: ([[V:%.+]]: f32)73// CHECK: [[V2:%.+]] = llvm.mlir.undef : f3274// CHECK: [[EXISTING:%.+]] = llvm.mlir.undef : i3275// CHECK: [[PACKED:%.+]] = rocdl.cvt.pk.fp8.f32 [[V]], [[V2]] -> [[EXISTING]][false] : i3276// CHECK: [[CAST:%.+]] = llvm.bitcast [[PACKED]] : i32 to vector<4xi8>77// CHECK: builtin.unrealized_conversion_cast [[CAST]] : vector<4xi8> to vector<4xf8E4M3FN>78func.func @packed_trunc(%v: f32) -> vector<4xf8E4M3FN> {79 %ret = amdgpu.packed_trunc_2xfp8 %v, undef into undef[word 0] : f32 to vector<4xf8E4M3FN>80 func.return %ret : vector<4xf8E4M3FN>81}82 83// CHECK-LABEL: func @packed_truncx284// CHECK-SAME: ([[V:%.+]]: f32, [[W:%.+]]: f32)85// CHECK: [[EXISTING:%.+]] = llvm.mlir.undef : i3286// CHECK: [[PACKED:%.+]] = rocdl.cvt.pk.fp8.f32 [[V]], [[W]] -> [[EXISTING]][false] : i3287// CHECK: [[CAST:%.+]] = llvm.bitcast [[PACKED]] : i32 to vector<4xi8>88// CHECK: builtin.unrealized_conversion_cast [[CAST]] : vector<4xi8> to vector<4xf8E4M3FN>89func.func @packed_truncx2(%v: f32, %w: f32) -> vector<4xf8E4M3FN> {90 %ret = amdgpu.packed_trunc_2xfp8 %v, %w into undef[word 0] : f32 to vector<4xf8E4M3FN>91 func.return %ret : vector<4xf8E4M3FN>92}93 94// CHECK-LABEL: func @packed_truncx2_into95// CHECK-SAME: ([[V:%.+]]: f32, [[W:%.+]]: f32, [[EXISTING:%.+]]: vector<4xf8E5M2>)96// CHECK: [[EXISTING_BYTES:%.+]] = builtin.unrealized_conversion_cast [[EXISTING]] : vector<4xf8E5M2> to vector<4xi8>97// CHECK: [[EXISTING_INT:%.+]] = llvm.bitcast [[EXISTING_BYTES]] : vector<4xi8> to i3298// CHECK: [[PACKED:%.+]] = rocdl.cvt.pk.bf8.f32 [[V]], [[W]] -> [[EXISTING_INT]][true] : i3299// CHECK: [[CAST:%.+]] = llvm.bitcast [[PACKED]] : i32 to vector<4xi8>100// CHECK: builtin.unrealized_conversion_cast [[CAST]] : vector<4xi8> to vector<4xf8E5M2>101func.func @packed_truncx2_into(%v: f32, %w: f32, %existing: vector<4xf8E5M2>) -> vector<4xf8E5M2> {102 %ret = amdgpu.packed_trunc_2xfp8 %v, %w into %existing[word 1] : f32 to vector<4xf8E5M2> into vector<4xf8E5M2>103 func.return %ret : vector<4xf8E5M2>104}105 106// CHECK-LABEL: func @packed_stoch_round107// CHECK-SAME: ([[V:%.+]]: f32, [[S:%.+]]: i32)108// CHECK: [[EXISTING:%.+]] = llvm.mlir.undef : i32109// CHECK: [[PACKED:%.+]] = rocdl.cvt.sr.fp8.f32 [[V]], [[S]] -> [[EXISTING]][0] : i32110// CHECK: [[CAST:%.+]] = llvm.bitcast [[PACKED]] : i32 to vector<4xi8>111// CHECK: builtin.unrealized_conversion_cast [[CAST]] : vector<4xi8> to vector<4xf8E4M3FN>112func.func @packed_stoch_round(%v: f32, %s: i32) -> vector<4xf8E4M3FN> {113 %ret = amdgpu.packed_stoch_round_fp8 %v + %s into undef[0] : f32 to vector<4xf8E4M3FN>114 func.return %ret : vector<4xf8E4M3FN>115}116 117// CHECK-LABEL: func @packed_stoch_round_into118// CHECK-SAME: ([[V:%.+]]: f32, [[S:%.+]]: i32, [[EXISTING:%.+]]: vector<4xf8E5M2>)119// CHECK: [[EXISTING_BYTES:%.+]] = builtin.unrealized_conversion_cast [[EXISTING]] : vector<4xf8E5M2> to vector<4xi8>120// CHECK: [[EXISTING_INT:%.+]] = llvm.bitcast [[EXISTING_BYTES]] : vector<4xi8> to i32121// CHECK: [[PACKED:%.+]] = rocdl.cvt.sr.bf8.f32 [[V]], [[S]] -> [[EXISTING_INT]][1] : i32122// CHECK: [[CAST:%.+]] = llvm.bitcast [[PACKED]] : i32 to vector<4xi8>123// CHECK: builtin.unrealized_conversion_cast [[CAST]] : vector<4xi8> to vector<4xf8E5M2>124func.func @packed_stoch_round_into(%v: f32, %s: i32, %existing: vector<4xf8E5M2>) -> vector<4xf8E5M2> {125 %ret = amdgpu.packed_stoch_round_fp8 %v + %s into %existing[1] : f32 to vector<4xf8E5M2> into vector<4xf8E5M2>126 func.return %ret : vector<4xf8E5M2>127}128