204 lines · plain
1// RUN: mlir-opt --split-input-file %s -convert-arith-to-amdgpu="chipset=gfx950" | FileCheck %s2// RUN: mlir-opt --split-input-file %s -convert-arith-to-amdgpu="chipset=gfx1100" | FileCheck %s --check-prefix=CHECK-GFX11003 4// CHECK-LABEL: @conversion_f8_fallback5// CHECK-DAG: %[[CST:.+]] = arith.constant dense<0.000000e+00> : vector<2x2xf8E5M2>6// CHECK-DAG: %[[SCALE_EXT:.+]] = arith.extf %arg1 : vector<2x2xf8E8M0FNU> to vector<2x2xf32>7// CHECK: %[[IN_SLICE_00:.+]] = vector.extract_strided_slice %arg0 {offsets = [0, 0], sizes = [1, 1], strides = [1, 1]}8// CHECK-NEXT: %[[IN_SCALAR_00:.+]] = vector.shape_cast %[[IN_SLICE_00]]9// CHECK-NEXT: %[[SCALE_SCALAR_00:.+]] = vector.extract %[[SCALE_EXT]][0, 0]10// CHECK-NEXT: %[[PACKED_00:.+]] = amdgpu.packed_scaled_trunc %[[IN_SCALAR_00]] into undef[0], %[[SCALE_SCALAR_00]]11// CHECK-NEXT: %[[OUT_SLICE_00:.+]] = vector.extract_strided_slice %[[PACKED_00]]12// CHECK-NEXT: %[[OUT_SCALAR_00:.+]] = vector.shape_cast %[[OUT_SLICE_00]]13// CHECK-NEXT: %[[ACC_A:.+]] = vector.insert_strided_slice %[[OUT_SCALAR_00]], %[[CST]]14// CHECK-NEXT: %[[IN_SLICE_01:.+]] = vector.extract_strided_slice %arg0 {offsets = [0, 1], sizes = [1, 1], strides = [1, 1]}15// CHECK-NEXT: %[[IN_SCALAR_01:.+]] = vector.shape_cast %[[IN_SLICE_01]]16// CHECK-NEXT: %[[SCALE_SCALAR_01:.+]] = vector.extract %[[SCALE_EXT]][0, 1]17// CHECK-NEXT: %[[PACKED_01:.+]] = amdgpu.packed_scaled_trunc %[[IN_SCALAR_01]] into undef[0], %[[SCALE_SCALAR_01]]18// CHECK-NEXT: %[[OUT_SLICE_01:.+]] = vector.extract_strided_slice %[[PACKED_01]]19// CHECK-NEXT: %[[OUT_SCALAR_01:.+]] = vector.shape_cast %[[OUT_SLICE_01]]20// CHECK-NEXT: %[[ACC_B:.+]] = vector.insert_strided_slice %[[OUT_SCALAR_01]], %[[ACC_A]]21// CHECK-NEXT: %[[IN_SLICE_10:.+]] = vector.extract_strided_slice %arg0 {offsets = [1, 0], sizes = [1, 1], strides = [1, 1]}22// CHECK-NEXT: %[[IN_SCALAR_10:.+]] = vector.shape_cast %[[IN_SLICE_10]]23// CHECK-NEXT: %[[SCALE_SCALAR_10:.+]] = vector.extract %[[SCALE_EXT]][1, 0]24// CHECK-NEXT: %[[PACKED_10:.+]] = amdgpu.packed_scaled_trunc %[[IN_SCALAR_10]] into undef[0], %[[SCALE_SCALAR_10]]25// CHECK-NEXT: %[[OUT_SLICE_10:.+]] = vector.extract_strided_slice %[[PACKED_10]]26// CHECK-NEXT: %[[OUT_SCALAR_10:.+]] = vector.shape_cast %[[OUT_SLICE_10]]27// CHECK-NEXT: %[[ACC_A:.+]] = vector.insert_strided_slice %[[OUT_SCALAR_10]], %[[ACC_B]]28// CHECK-NEXT: %[[IN_SLICE_11:.+]] = vector.extract_strided_slice %arg0 {offsets = [1, 1], sizes = [1, 1], strides = [1, 1]}29// CHECK-NEXT: %[[IN_SCALAR_11:.+]] = vector.shape_cast %[[IN_SLICE_11]]30// CHECK-NEXT: %[[SCALE_SCALAR_11:.+]] = vector.extract %[[SCALE_EXT]][1, 1]31// CHECK-NEXT: %[[PACKED_11:.+]] = amdgpu.packed_scaled_trunc %[[IN_SCALAR_11]] into undef[0], %[[SCALE_SCALAR_11]]32// CHECK-NEXT: %[[OUT_SLICE_11:.+]] = vector.extract_strided_slice %[[PACKED_11]]33// CHECK-NEXT: %[[OUT_SCALAR_11:.+]] = vector.shape_cast %[[OUT_SLICE_11]]34// CHECK-NEXT: %[[ACC_B:.+]] = vector.insert_strided_slice %[[OUT_SCALAR_11]], %[[ACC_A]]35// CHECK-NEXT: return %[[ACC_B]] : vector<2x2xf8E5M2>36func.func @conversion_f8_fallback(%in: vector<2x2xf32>, %scale: vector<2x2xf8E8M0FNU>) -> vector<2x2xf8E5M2> {37 %ext = arith.scaling_truncf %in, %scale : vector<2x2xf32>, vector<2x2xf8E8M0FNU> to vector<2x2xf8E5M2>38 return %ext : vector<2x2xf8E5M2>39}40 41// -----42 43// CHECK-LABEL: @conversion_f4_fallback44// CHECK-DAG: %[[CST:.+]] = arith.constant dense<0.000000e+00> : vector<2x2xf4E2M1FN>45// CHECK-DAG: %[[SCALE_EXT:.+]] = arith.extf %arg1 : vector<2x2xf8E8M0FNU> to vector<2x2xf32>46// CHECK: %[[IN_SLICE_00:.+]] = vector.extract_strided_slice %arg0 {offsets = [0, 0], sizes = [1, 1], strides = [1, 1]}47// CHECK-NEXT: %[[IN_SCALAR_00:.+]] = vector.shape_cast %[[IN_SLICE_00]]48// CHECK-NEXT: %[[SCALE_SCALAR_00:.+]] = vector.extract %[[SCALE_EXT]][0, 0]49// CHECK-NEXT: %[[PACKED_00:.+]] = amdgpu.packed_scaled_trunc %[[IN_SCALAR_00]] into undef[0], %[[SCALE_SCALAR_00]]50// CHECK-NEXT: %[[OUT_SLICE_00:.+]] = vector.extract_strided_slice %[[PACKED_00]]51// CHECK-NEXT: %[[OUT_SCALAR_00:.+]] = vector.shape_cast %[[OUT_SLICE_00]]52// CHECK-NEXT: %[[ACC_A:.+]] = vector.insert_strided_slice %[[OUT_SCALAR_00]], %[[CST]]53// CHECK-NEXT: %[[IN_SLICE_01:.+]] = vector.extract_strided_slice %arg0 {offsets = [0, 1], sizes = [1, 1], strides = [1, 1]}54// CHECK-NEXT: %[[IN_SCALAR_01:.+]] = vector.shape_cast %[[IN_SLICE_01]]55// CHECK-NEXT: %[[SCALE_SCALAR_01:.+]] = vector.extract %[[SCALE_EXT]][0, 1]56// CHECK-NEXT: %[[PACKED_01:.+]] = amdgpu.packed_scaled_trunc %[[IN_SCALAR_01]] into undef[0], %[[SCALE_SCALAR_01]]57// CHECK-NEXT: %[[OUT_SLICE_01:.+]] = vector.extract_strided_slice %[[PACKED_01]]58// CHECK-NEXT: %[[OUT_SCALAR_01:.+]] = vector.shape_cast %[[OUT_SLICE_01]]59// CHECK-NEXT: %[[ACC_B:.+]] = vector.insert_strided_slice %[[OUT_SCALAR_01]], %[[ACC_A]]60// CHECK-NEXT: %[[IN_SLICE_10:.+]] = vector.extract_strided_slice %arg0 {offsets = [1, 0], sizes = [1, 1], strides = [1, 1]}61// CHECK-NEXT: %[[IN_SCALAR_10:.+]] = vector.shape_cast %[[IN_SLICE_10]]62// CHECK-NEXT: %[[SCALE_SCALAR_10:.+]] = vector.extract %[[SCALE_EXT]][1, 0]63// CHECK-NEXT: %[[PACKED_10:.+]] = amdgpu.packed_scaled_trunc %[[IN_SCALAR_10]] into undef[0], %[[SCALE_SCALAR_10]]64// CHECK-NEXT: %[[OUT_SLICE_10:.+]] = vector.extract_strided_slice %[[PACKED_10]]65// CHECK-NEXT: %[[OUT_SCALAR_10:.+]] = vector.shape_cast %[[OUT_SLICE_10]]66// CHECK-NEXT: %[[ACC_A:.+]] = vector.insert_strided_slice %[[OUT_SCALAR_10]], %[[ACC_B]]67// CHECK-NEXT: %[[IN_SLICE_11:.+]] = vector.extract_strided_slice %arg0 {offsets = [1, 1], sizes = [1, 1], strides = [1, 1]}68// CHECK-NEXT: %[[IN_SCALAR_11:.+]] = vector.shape_cast %[[IN_SLICE_11]]69// CHECK-NEXT: %[[SCALE_SCALAR_11:.+]] = vector.extract %[[SCALE_EXT]][1, 1]70// CHECK-NEXT: %[[PACKED_11:.+]] = amdgpu.packed_scaled_trunc %[[IN_SCALAR_11]] into undef[0], %[[SCALE_SCALAR_11]]71// CHECK-NEXT: %[[OUT_SLICE_11:.+]] = vector.extract_strided_slice %[[PACKED_11]]72// CHECK-NEXT: %[[OUT_SCALAR_11:.+]] = vector.shape_cast %[[OUT_SLICE_11]]73// CHECK-NEXT: %[[ACC_B:.+]] = vector.insert_strided_slice %[[OUT_SCALAR_11]], %[[ACC_A]]74// CHECK-NEXT: return %[[ACC_B]] : vector<2x2xf4E2M1FN>75func.func @conversion_f4_fallback(%in: vector<2x2xf32>, %scale: vector<2x2xf8E8M0FNU>) -> vector<2x2xf4E2M1FN> {76 %ext = arith.scaling_truncf %in, %scale : vector<2x2xf32>, vector<2x2xf8E8M0FNU> to vector<2x2xf4E2M1FN>77 return %ext : vector<2x2xf4E2M1FN>78}79 80// -----81 82// CHECK-LABEL: @conversion_broadcast83// CHECK-DAG: %[[CST:.+]] = arith.constant dense<0.000000e+00> : vector<8x2x4xf8E5M2>84// CHECK-DAG: %[[BCAST:.+]] = vector.broadcast %arg185// CHECK-DAG: %[[IN_CAST:.+]] = vector.shape_cast %arg086// CHECK-DAG: %[[SCALE_CAST:.+]] = vector.shape_cast %[[BCAST]]87// CHECK-DAG: %[[SCALE_EXT:.+]] = arith.extf %[[SCALE_CAST]]88// CHECK-DAG: vector.extract_strided_slice %[[IN_CAST]] {offsets = [0, 0, 0], sizes = [1, 1, 4], strides = [1, 1, 1]}89// CHECK-NEXT: vector.shape_cast90// CHECK-NEXT: vector.extract %[[SCALE_EXT]][0, 0, 0]91// CHECK-NEXT: vector.extract_strided_slice %{{.+}} {offsets = [0], sizes = [2], strides = [1]}92// CHECK-NEXT: %[[P1:.+]] = amdgpu.packed_scaled_trunc93// CHECK-NEXT: vector.extract_strided_slice %{{.+}} {offsets = [2], sizes = [2], strides = [1]}94// CHECK-NEXT: %[[P2:.+]] = amdgpu.packed_scaled_trunc {{.*}} into %[[P1]][1]95// CHECK-NEXT: %[[P2_CAST:.+]] = vector.shape_cast %[[P2]] : vector<4xf8E5M2> to vector<1x1x4xf8E5M2>96// CHECK-NEXT: vector.insert_strided_slice %[[P2_CAST]], %{{.+}} {offsets = [0, 0, 0], strides = [1, 1, 1]}97// CHECK-NEXT: vector.extract_strided_slice %[[IN_CAST]] {offsets = [0, 1, 0], sizes = [1, 1, 4], strides = [1, 1, 1]}98// CHECK-NEXT: vector.shape_cast99// CHECK-NEXT: vector.extract %[[SCALE_EXT]][0, 1, 0]100// CHECK-NEXT: vector.extract_strided_slice %{{.+}} {offsets = [0], sizes = [2], strides = [1]}101// CHECK-NEXT: amdgpu.packed_scaled_trunc102// CHECK-NEXT: vector.extract_strided_slice %{{.+}} {offsets = [2], sizes = [2], strides = [1]}103// CHECK-NEXT: amdgpu.packed_scaled_trunc104// CHECK-NEXT: vector.shape_cast105// CHECK-NEXT: vector.insert_strided_slice %{{.+}}, %{{.+}} {offsets = [0, 1, 0], strides = [1, 1, 1]}106func.func @conversion_broadcast(%in: vector<8x8xf32>, %scale: vector<8x2xf8E8M0FNU>) -> vector<8x8xf8E5M2> {107 %bc = vector.broadcast %scale : vector<8x2xf8E8M0FNU> to vector<4x8x2xf8E8M0FNU>108 %cast1 = vector.shape_cast %in : vector<8x8xf32> to vector<8x2x4xf32>109 %cast2 = vector.shape_cast %bc : vector<4x8x2xf8E8M0FNU> to vector<8x2x4xf8E8M0FNU>110 %ext = arith.scaling_truncf %cast1, %cast2 : vector<8x2x4xf32>, vector<8x2x4xf8E8M0FNU> to vector<8x2x4xf8E5M2>111 %cast3 = vector.shape_cast %ext : vector<8x2x4xf8E5M2> to vector<8x8xf8E5M2>112 return %cast3 : vector<8x8xf8E5M2>113}114 115// -----116 117// CHECK-LABEL: @conversion_broadcast_odd118// CHECK-NEXT: %[[CST4:.+]] = arith.constant dense<0.000000e+00> : vector<4xf8E5M2>119// CHECK-NEXT: %[[CST6:.+]] = arith.constant dense<0.000000e+00> : vector<6xf8E5M2>120// CHECK-NEXT: %[[SCALE_BCAST:.+]] = vector.broadcast %arg1 : vector<2xf8E8M0FNU> to vector<3x2xf8E8M0FNU>121// CHECK-NEXT: %[[SCALE_FLAT:.+]] = vector.shape_cast %[[SCALE_BCAST]] : vector<3x2xf8E8M0FNU> to vector<6xf8E8M0FNU>122// CHECK-NEXT: %[[SCALE_EXTF:.+]] = arith.extf %[[SCALE_FLAT]] : vector<6xf8E8M0FNU> to vector<6xf32>123// CHECK-NEXT: %[[IN_CHUNK0:.+]] = vector.extract_strided_slice %arg0 {offsets = [0], sizes = [3], strides = [1]} : vector<6xf32> to vector<3xf32>124// CHECK-NEXT: %[[SCALE0:.+]] = vector.extract %[[SCALE_EXTF]][0] : f32 from vector<6xf32>125// CHECK-NEXT: %[[IN_CHUNK0_PART0:.+]] = vector.extract_strided_slice %[[IN_CHUNK0]] {offsets = [0], sizes = [2], strides = [1]} : vector<3xf32> to vector<2xf32>126// CHECK-NEXT: %[[PACKED0_PART0:.+]] = amdgpu.packed_scaled_trunc %[[IN_CHUNK0_PART0]] into %[[CST4]][0], %[[SCALE0]] : vector<2xf32> to vector<4xf8E5M2>127// CHECK-NEXT: %[[IN_CHUNK0_PART1:.+]] = vector.extract_strided_slice %[[IN_CHUNK0]] {offsets = [2], sizes = [1], strides = [1]} : vector<3xf32> to vector<1xf32>128// CHECK-NEXT: %[[PACKED0_PART1:.+]] = amdgpu.packed_scaled_trunc %[[IN_CHUNK0_PART1]] into %[[PACKED0_PART0]][1], %[[SCALE0]] : vector<1xf32> to vector<4xf8E5M2>129// CHECK-NEXT: %[[CHUNK0_RES:.+]] = vector.extract_strided_slice %[[PACKED0_PART1]] {offsets = [0], sizes = [3], strides = [1]} : vector<4xf8E5M2> to vector<3xf8E5M2>130// CHECK-NEXT: %[[FINAL_ACCUM_A:.+]] = vector.insert_strided_slice %[[CHUNK0_RES]], %[[CST6]] {offsets = [0], strides = [1]} : vector<3xf8E5M2> into vector<6xf8E5M2>131// CHECK-NEXT: %[[IN_CHUNK1:.+]] = vector.extract_strided_slice %arg0 {offsets = [3], sizes = [3], strides = [1]} : vector<6xf32> to vector<3xf32>132// CHECK-NEXT: %[[SCALE1:.+]] = vector.extract %[[SCALE_EXTF]][3] : f32 from vector<6xf32>133// CHECK-NEXT: %[[IN_CHUNK1_PART0:.+]] = vector.extract_strided_slice %[[IN_CHUNK1]] {offsets = [0], sizes = [2], strides = [1]} : vector<3xf32> to vector<2xf32>134// CHECK-NEXT: %[[PACKED1_PART0:.+]] = amdgpu.packed_scaled_trunc %[[IN_CHUNK1_PART0]] into %[[CST4]][0], %[[SCALE1]] : vector<2xf32> to vector<4xf8E5M2>135// CHECK-NEXT: %[[IN_CHUNK1_PART1:.+]] = vector.extract_strided_slice %[[IN_CHUNK1]] {offsets = [2], sizes = [1], strides = [1]} : vector<3xf32> to vector<1xf32>136// CHECK-NEXT: %[[PACKED1_PART1:.+]] = amdgpu.packed_scaled_trunc %[[IN_CHUNK1_PART1]] into %[[PACKED1_PART0]][1], %[[SCALE1]] : vector<1xf32> to vector<4xf8E5M2>137// CHECK-NEXT: %[[CHUNK1_RES:.+]] = vector.extract_strided_slice %[[PACKED1_PART1]] {offsets = [0], sizes = [3], strides = [1]} : vector<4xf8E5M2> to vector<3xf8E5M2>138// CHECK-NEXT: %[[FINAL_RESULT:.+]] = vector.insert_strided_slice %[[CHUNK1_RES]], %[[FINAL_ACCUM_A]] {offsets = [3], strides = [1]} : vector<3xf8E5M2> into vector<6xf8E5M2>139// CHECK-NEXT: return %[[FINAL_RESULT]] : vector<6xf8E5M2>140func.func @conversion_broadcast_odd(%in: vector<6xf32>, %scale: vector<2xf8E8M0FNU>) -> vector<6xf8E5M2> {141 %bc = vector.broadcast %scale : vector<2xf8E8M0FNU> to vector<3x2xf8E8M0FNU>142 %cast = vector.shape_cast %bc : vector<3x2xf8E8M0FNU> to vector<6xf8E8M0FNU>143 %ext = arith.scaling_truncf %in, %cast : vector<6xf32>, vector<6xf8E8M0FNU> to vector<6xf8E5M2>144 return %ext : vector<6xf8E5M2>145}146 147// -----148 149// CHECK-LABEL: @conversion_broadcast150// CHECK-DAG: %[[CST:.+]] = arith.constant dense<0.000000e+00> : vector<4xf8E5M2>151// CHECK-DAG: %[[SCALE_SPLAT:.+]] = vector.broadcast %arg1 : f8E8M0FNU to vector<4xf8E8M0FNU>152// CHECK-DAG: %[[SCALE_EXTF:.+]] = arith.extf %[[SCALE_SPLAT]] : vector<4xf8E8M0FNU> to vector<4xf32>153// CHECK-DAG: %[[SCALE_SCALAR:.+]] = vector.extract %[[SCALE_EXTF]][0] : f32 from vector<4xf32>154// CHECK: %[[IN_CHUNK0:.+]] = vector.extract_strided_slice %arg0 {offsets = [0], sizes = [2], strides = [1]} : vector<4xf32> to vector<2xf32>155// CHECK-NEXT: %[[PACKED0:.+]] = amdgpu.packed_scaled_trunc %[[IN_CHUNK0]] into %[[CST]][0], %[[SCALE_SCALAR]] : vector<2xf32> to vector<4xf8E5M2>156// CHECK-NEXT: %[[IN_CHUNK1:.+]] = vector.extract_strided_slice %arg0 {offsets = [2], sizes = [2], strides = [1]} : vector<4xf32> to vector<2xf32>157// CHECK-NEXT: %[[PACKED1:.+]] = amdgpu.packed_scaled_trunc %[[IN_CHUNK1]] into %[[PACKED0]][1], %[[SCALE_SCALAR]] : vector<2xf32> to vector<4xf8E5M2>158// CHECK-NEXT: return %[[PACKED1]] : vector<4xf8E5M2>159func.func @conversion_broadcast(%in: vector<4xf32>, %scale: f8E8M0FNU) -> vector<4xf8E5M2> {160 %splat = vector.broadcast %scale : f8E8M0FNU to vector<4xf8E8M0FNU>161 %ext = arith.scaling_truncf %in, %splat : vector<4xf32>, vector<4xf8E8M0FNU> to vector<4xf8E5M2>162 return %ext : vector<4xf8E5M2>163}164 165// -----166 167// CHECK-GFX1100-LABEL: @conversion_scalar168// CHECK-GFX1100: arith.scaling_truncf169 170// CHECK-LABEL: @conversion_scalar171// CHECK: %[[SCALE_F32:.+]] = arith.extf %arg1 : f8E8M0FNU to f32172// CHECK-NEXT: %[[SPLAT_IN:.+]] = vector.broadcast %arg0 : f32 to vector<1xf32>173// CHECK-NEXT: %[[PACKED_TRUNC:.+]] = amdgpu.packed_scaled_trunc %[[SPLAT_IN]] into undef[0], %[[SCALE_F32]]174// CHECK-NEXT: %[[RESULT:.+]] = vector.extract %[[PACKED_TRUNC]][0]175// CHECK-NEXT: return %[[RESULT]] : f8E5M2176func.func @conversion_scalar(%in: f32, %scale: f8E8M0FNU) -> f8E5M2 {177 %ext = arith.scaling_truncf %in, %scale : f32, f8E8M0FNU to f8E5M2178 return %ext : f8E5M2179}180 181// -----182 183// CHECK-LABEL: @long_fp4_broadcast184// CHECK-COUNT-4: amdgpu.packed_scaled_trunc %{{.*}} into %{{.+}}[3]185// CHECK-NOT: amdgpu.packed_scaled_trunc186// CHECK: return187func.func @long_fp4_broadcast(%in: vector<32xf32>, %scale: f32) -> vector<32xf4E2M1FN> {188 %splat = vector.broadcast %scale : f32 to vector<32xf32>189 %trunc = arith.scaling_truncf %in, %splat : vector<32xf32>, vector<32xf32> to vector<32xf4E2M1FN>190 return %trunc : vector<32xf4E2M1FN>191}192 193// -----194 195// CHECK-LABEL: @long_fp8_broadcast196// CHECK-COUNT-8: amdgpu.packed_scaled_trunc %{{.*}} into %{{.+}}[1]197// CHECK-NOT: amdgpu.packed_scaled_trunc198// CHECK: return199func.func @long_fp8_broadcast(%in: vector<32xf32>, %scale: f32) -> vector<32xf8E4M3FN> {200 %splat = vector.broadcast %scale : f32 to vector<32xf32>201 %trunc = arith.scaling_truncf %in, %splat : vector<32xf32>, vector<32xf32> to vector<32xf8E4M3FN>202 return %trunc : vector<32xf8E4M3FN>203}204