[Mlir-commits] [mlir] [mlir][amdgpu] Add conversion from arith.scaling_extf to amdgpu (PR #146372)

Tim Gymnich llvmlistbot at llvm.org
Tue Jul 8 06:01:33 PDT 2025


================
@@ -0,0 +1,635 @@
+// RUN: mlir-opt --split-input-file %s -convert-arith-to-amdgpu="chipset=gfx950" | FileCheck %s
+
+// CHECK-LABEL: @conversion_f8_f32_fallback
+// CHECK:         [[CST:%.+]] = arith.constant dense<0.000000e+00> : vector<2x2xf32>
+// CHECK-NEXT:    [[SCALE_EXT:%.+]] = arith.extf %arg1 : vector<2x2xf8E8M0FNU> to vector<2x2xf32>
+// CHECK-NEXT:    [[IN_SLICE_00:%.+]] = vector.extract_strided_slice %arg0 {offsets = [0, 0], sizes = [1, 1], strides = [1, 1]} : vector<2x2xf8E5M2> to vector<1x1xf8E5M2>
+// CHECK-NEXT:    [[IN_VEC_00:%.+]] = vector.shape_cast [[IN_SLICE_00]] : vector<1x1xf8E5M2> to vector<1xf8E5M2>
+// CHECK-NEXT:    [[SCALE_SCALAR_00:%.+]] = vector.extract [[SCALE_EXT]][0, 0] : f32 from vector<2x2xf32>
+// CHECK-NEXT:    [[PACKED_00:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_00]][0], [[SCALE_SCALAR_00]] : vector<1xf8E5M2> to vector<2xf32>
+// CHECK-NEXT:    [[OUT_SLICE_00:%.+]] = vector.extract_strided_slice [[PACKED_00]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf32> to vector<1xf32>
+// CHECK-NEXT:    [[OUT_VEC_00:%.+]] = vector.shape_cast [[OUT_SLICE_00]] : vector<1xf32> to vector<1x1xf32>
+// CHECK-NEXT:    [[ACC_A:%.+]] = vector.insert_strided_slice [[OUT_VEC_00]], [[CST]] {offsets = [0, 0], strides = [1, 1]} : vector<1x1xf32> into vector<2x2xf32>
+// CHECK-NEXT:    [[IN_SLICE_01:%.+]] = vector.extract_strided_slice %arg0 {offsets = [0, 1], sizes = [1, 1], strides = [1, 1]} : vector<2x2xf8E5M2> to vector<1x1xf8E5M2>
+// CHECK-NEXT:    [[IN_VEC_01:%.+]] = vector.shape_cast [[IN_SLICE_01]] : vector<1x1xf8E5M2> to vector<1xf8E5M2>
+// CHECK-NEXT:    [[SCALE_SCALAR_01:%.+]] = vector.extract [[SCALE_EXT]][0, 1] : f32 from vector<2x2xf32>
+// CHECK-NEXT:    [[PACKED_01:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_01]][0], [[SCALE_SCALAR_01]] : vector<1xf8E5M2> to vector<2xf32>
+// CHECK-NEXT:    [[OUT_SLICE_01:%.+]] = vector.extract_strided_slice [[PACKED_01]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf32> to vector<1xf32>
+// CHECK-NEXT:    [[OUT_VEC_01:%.+]] = vector.shape_cast [[OUT_SLICE_01]] : vector<1xf32> to vector<1x1xf32>
+// CHECK-NEXT:    [[ACC_B:%.+]] = vector.insert_strided_slice [[OUT_VEC_01]], [[ACC_A]] {offsets = [0, 1], strides = [1, 1]} : vector<1x1xf32> into vector<2x2xf32>
+// CHECK-NEXT:    [[IN_SLICE_10:%.+]] = vector.extract_strided_slice %arg0 {offsets = [1, 0], sizes = [1, 1], strides = [1, 1]} : vector<2x2xf8E5M2> to vector<1x1xf8E5M2>
+// CHECK-NEXT:    [[IN_VEC_10:%.+]] = vector.shape_cast [[IN_SLICE_10]] : vector<1x1xf8E5M2> to vector<1xf8E5M2>
+// CHECK-NEXT:    [[SCALE_SCALAR_10:%.+]] = vector.extract [[SCALE_EXT]][1, 0] : f32 from vector<2x2xf32>
+// CHECK-NEXT:    [[PACKED_10:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_10]][0], [[SCALE_SCALAR_10]] : vector<1xf8E5M2> to vector<2xf32>
+// CHECK-NEXT:    [[OUT_SLICE_10:%.+]] = vector.extract_strided_slice [[PACKED_10]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf32> to vector<1xf32>
+// CHECK-NEXT:    [[OUT_VEC_10:%.+]] = vector.shape_cast [[OUT_SLICE_10]] : vector<1xf32> to vector<1x1xf32>
+// CHECK-NEXT:    [[ACC_A:%.+]] = vector.insert_strided_slice [[OUT_VEC_10]], [[ACC_B]] {offsets = [1, 0], strides = [1, 1]} : vector<1x1xf32> into vector<2x2xf32>
+// CHECK-NEXT:    [[IN_SLICE_11:%.+]] = vector.extract_strided_slice %arg0 {offsets = [1, 1], sizes = [1, 1], strides = [1, 1]} : vector<2x2xf8E5M2> to vector<1x1xf8E5M2>
+// CHECK-NEXT:    [[IN_VEC_11:%.+]] = vector.shape_cast [[IN_SLICE_11]] : vector<1x1xf8E5M2> to vector<1xf8E5M2>
+// CHECK-NEXT:    [[SCALE_SCALAR_11:%.+]] = vector.extract [[SCALE_EXT]][1, 1] : f32 from vector<2x2xf32>
+// CHECK-NEXT:    [[PACKED_11:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_11]][0], [[SCALE_SCALAR_11]] : vector<1xf8E5M2> to vector<2xf32>
+// CHECK-NEXT:    [[OUT_SLICE_11:%.+]] = vector.extract_strided_slice [[PACKED_11]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf32> to vector<1xf32>
+// CHECK-NEXT:    [[OUT_VEC_11:%.+]] = vector.shape_cast [[OUT_SLICE_11]] : vector<1xf32> to vector<1x1xf32>
+// CHECK-NEXT:    [[ACC_B:%.+]] = vector.insert_strided_slice [[OUT_VEC_11]], [[ACC_A]] {offsets = [1, 1], strides = [1, 1]} : vector<1x1xf32> into vector<2x2xf32>
+// CHECK-NEXT:    return [[ACC_B]] : vector<2x2xf32>
+func.func @conversion_f8_f32_fallback(%in: vector<2x2xf8E5M2>, %scale: vector<2x2xf8E8M0FNU>) -> vector<2x2xf32> {
+    %ext = arith.scaling_extf %in, %scale : vector<2x2xf8E5M2>, vector<2x2xf8E8M0FNU> to vector<2x2xf32>
+    return %ext : vector<2x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @conversion_f4_f32_fallback
+// CHECK:         [[CST:%.+]] = arith.constant dense<0.000000e+00> : vector<2x2xf32>
+// CHECK-NEXT:    [[SCALE_EXT:%.+]] = arith.extf %arg1 : vector<2x2xf8E8M0FNU> to vector<2x2xf32>
+// CHECK-NEXT:    [[IN_SLICE_00:%.+]] = vector.extract_strided_slice %arg0 {offsets = [0, 0], sizes = [1, 1], strides = [1, 1]} : vector<2x2xf4E2M1FN> to vector<1x1xf4E2M1FN>
+// CHECK-NEXT:    [[IN_VEC_00:%.+]] = vector.shape_cast [[IN_SLICE_00]] : vector<1x1xf4E2M1FN> to vector<1xf4E2M1FN>
+// CHECK-NEXT:    [[SCALE_SCALAR_00:%.+]] = vector.extract [[SCALE_EXT]][0, 0] : f32 from vector<2x2xf32>
+// CHECK-NEXT:    [[PACKED_00:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_00]][0], [[SCALE_SCALAR_00]] : vector<1xf4E2M1FN> to vector<2xf32>
+// CHECK-NEXT:    [[OUT_SLICE_00:%.+]] = vector.extract_strided_slice [[PACKED_00]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf32> to vector<1xf32>
+// CHECK-NEXT:    [[OUT_VEC_00:%.+]] = vector.shape_cast [[OUT_SLICE_00]] : vector<1xf32> to vector<1x1xf32>
+// CHECK-NEXT:    [[ACC_A:%.+]] = vector.insert_strided_slice [[OUT_VEC_00]], [[CST]] {offsets = [0, 0], strides = [1, 1]} : vector<1x1xf32> into vector<2x2xf32>
+// CHECK-NEXT:    [[IN_SLICE_01:%.+]] = vector.extract_strided_slice %arg0 {offsets = [0, 1], sizes = [1, 1], strides = [1, 1]} : vector<2x2xf4E2M1FN> to vector<1x1xf4E2M1FN>
+// CHECK-NEXT:    [[IN_VEC_01:%.+]] = vector.shape_cast [[IN_SLICE_01]] : vector<1x1xf4E2M1FN> to vector<1xf4E2M1FN>
+// CHECK-NEXT:    [[SCALE_SCALAR_01:%.+]] = vector.extract [[SCALE_EXT]][0, 1] : f32 from vector<2x2xf32>
+// CHECK-NEXT:    [[PACKED_01:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_01]][0], [[SCALE_SCALAR_01]] : vector<1xf4E2M1FN> to vector<2xf32>
+// CHECK-NEXT:    [[OUT_SLICE_01:%.+]] = vector.extract_strided_slice [[PACKED_01]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf32> to vector<1xf32>
+// CHECK-NEXT:    [[OUT_VEC_01:%.+]] = vector.shape_cast [[OUT_SLICE_01]] : vector<1xf32> to vector<1x1xf32>
+// CHECK-NEXT:    [[ACC_B:%.+]] = vector.insert_strided_slice [[OUT_VEC_01]], [[ACC_A]] {offsets = [0, 1], strides = [1, 1]} : vector<1x1xf32> into vector<2x2xf32>
+// CHECK-NEXT:    [[IN_SLICE_10:%.+]] = vector.extract_strided_slice %arg0 {offsets = [1, 0], sizes = [1, 1], strides = [1, 1]} : vector<2x2xf4E2M1FN> to vector<1x1xf4E2M1FN>
+// CHECK-NEXT:    [[IN_VEC_10:%.+]] = vector.shape_cast [[IN_SLICE_10]] : vector<1x1xf4E2M1FN> to vector<1xf4E2M1FN>
+// CHECK-NEXT:    [[SCALE_SCALAR_10:%.+]] = vector.extract [[SCALE_EXT]][1, 0] : f32 from vector<2x2xf32>
+// CHECK-NEXT:    [[PACKED_10:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_10]][0], [[SCALE_SCALAR_10]] : vector<1xf4E2M1FN> to vector<2xf32>
+// CHECK-NEXT:    [[OUT_SLICE_10:%.+]] = vector.extract_strided_slice [[PACKED_10]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf32> to vector<1xf32>
+// CHECK-NEXT:    [[OUT_VEC_10:%.+]] = vector.shape_cast [[OUT_SLICE_10]] : vector<1xf32> to vector<1x1xf32>
+// CHECK-NEXT:    [[ACC_A:%.+]] = vector.insert_strided_slice [[OUT_VEC_10]], [[ACC_B]] {offsets = [1, 0], strides = [1, 1]} : vector<1x1xf32> into vector<2x2xf32>
+// CHECK-NEXT:    [[IN_SLICE_11:%.+]] = vector.extract_strided_slice %arg0 {offsets = [1, 1], sizes = [1, 1], strides = [1, 1]} : vector<2x2xf4E2M1FN> to vector<1x1xf4E2M1FN>
+// CHECK-NEXT:    [[IN_VEC_11:%.+]] = vector.shape_cast [[IN_SLICE_11]] : vector<1x1xf4E2M1FN> to vector<1xf4E2M1FN>
+// CHECK-NEXT:    [[SCALE_SCALAR_11:%.+]] = vector.extract [[SCALE_EXT]][1, 1] : f32 from vector<2x2xf32>
+// CHECK-NEXT:    [[PACKED_11:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_11]][0], [[SCALE_SCALAR_11]] : vector<1xf4E2M1FN> to vector<2xf32>
+// CHECK-NEXT:    [[OUT_SLICE_11:%.+]] = vector.extract_strided_slice [[PACKED_11]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf32> to vector<1xf32>
+// CHECK-NEXT:    [[OUT_VEC_11:%.+]] = vector.shape_cast [[OUT_SLICE_11]] : vector<1xf32> to vector<1x1xf32>
+// CHECK-NEXT:    [[ACC_B:%.+]] = vector.insert_strided_slice [[OUT_VEC_11]], [[ACC_A]] {offsets = [1, 1], strides = [1, 1]} : vector<1x1xf32> into vector<2x2xf32>
+// CHECK-NEXT:    return [[ACC_B]] : vector<2x2xf32>
+func.func @conversion_f4_f32_fallback(%in: vector<2x2xf4E2M1FN>, %scale: vector<2x2xf8E8M0FNU>) -> vector<2x2xf32> {
+    %ext = arith.scaling_extf %in, %scale : vector<2x2xf4E2M1FN>, vector<2x2xf8E8M0FNU> to vector<2x2xf32>
+    return %ext : vector<2x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @conversion_f8_f16_fallback
+// CHECK:         [[CST:%.+]] = arith.constant dense<0.000000e+00> : vector<2x2xf16>
+// CHECK-NEXT:    [[SCALE_EXT:%.+]] = arith.extf %arg1 : vector<2x2xf8E8M0FNU> to vector<2x2xf32>
+// CHECK-NEXT:    [[IN_SLICE_00:%.+]] = vector.extract_strided_slice %arg0 {offsets = [0, 0], sizes = [1, 1], strides = [1, 1]} : vector<2x2xf8E5M2> to vector<1x1xf8E5M2>
+// CHECK-NEXT:    [[IN_VEC_00:%.+]] = vector.shape_cast [[IN_SLICE_00]] : vector<1x1xf8E5M2> to vector<1xf8E5M2>
+// CHECK-NEXT:    [[SCALE_SCALAR_00:%.+]] = vector.extract [[SCALE_EXT]][0, 0] : f32 from vector<2x2xf32>
+// CHECK-NEXT:    [[PACKED_00:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_00]][0], [[SCALE_SCALAR_00]] : vector<1xf8E5M2> to vector<2xf16>
+// CHECK-NEXT:    [[OUT_SLICE_00:%.+]] = vector.extract_strided_slice [[PACKED_00]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf16> to vector<1xf16>
+// CHECK-NEXT:    [[OUT_VEC_00:%.+]] = vector.shape_cast [[OUT_SLICE_00]] : vector<1xf16> to vector<1x1xf16>
+// CHECK-NEXT:    [[ACC_A:%.+]] = vector.insert_strided_slice [[OUT_VEC_00]], [[CST]] {offsets = [0, 0], strides = [1, 1]} : vector<1x1xf16> into vector<2x2xf16>
+// CHECK-NEXT:    [[IN_SLICE_01:%.+]] = vector.extract_strided_slice %arg0 {offsets = [0, 1], sizes = [1, 1], strides = [1, 1]} : vector<2x2xf8E5M2> to vector<1x1xf8E5M2>
+// CHECK-NEXT:    [[IN_VEC_01:%.+]] = vector.shape_cast [[IN_SLICE_01]] : vector<1x1xf8E5M2> to vector<1xf8E5M2>
+// CHECK-NEXT:    [[SCALE_SCALAR_01:%.+]] = vector.extract [[SCALE_EXT]][0, 1] : f32 from vector<2x2xf32>
+// CHECK-NEXT:    [[PACKED_01:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_01]][0], [[SCALE_SCALAR_01]] : vector<1xf8E5M2> to vector<2xf16>
+// CHECK-NEXT:    [[OUT_SLICE_01:%.+]] = vector.extract_strided_slice [[PACKED_01]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf16> to vector<1xf16>
+// CHECK-NEXT:    [[OUT_VEC_01:%.+]] = vector.shape_cast [[OUT_SLICE_01]] : vector<1xf16> to vector<1x1xf16>
+// CHECK-NEXT:    [[ACC_B:%.+]] = vector.insert_strided_slice [[OUT_VEC_01]], [[ACC_A]] {offsets = [0, 1], strides = [1, 1]} : vector<1x1xf16> into vector<2x2xf16>
+// CHECK-NEXT:    [[IN_SLICE_10:%.+]] = vector.extract_strided_slice %arg0 {offsets = [1, 0], sizes = [1, 1], strides = [1, 1]} : vector<2x2xf8E5M2> to vector<1x1xf8E5M2>
+// CHECK-NEXT:    [[IN_VEC_10:%.+]] = vector.shape_cast [[IN_SLICE_10]] : vector<1x1xf8E5M2> to vector<1xf8E5M2>
+// CHECK-NEXT:    [[SCALE_SCALAR_10:%.+]] = vector.extract [[SCALE_EXT]][1, 0] : f32 from vector<2x2xf32>
+// CHECK-NEXT:    [[PACKED_10:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_10]][0], [[SCALE_SCALAR_10]] : vector<1xf8E5M2> to vector<2xf16>
+// CHECK-NEXT:    [[OUT_SLICE_10:%.+]] = vector.extract_strided_slice [[PACKED_10]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf16> to vector<1xf16>
+// CHECK-NEXT:    [[OUT_VEC_10:%.+]] = vector.shape_cast [[OUT_SLICE_10]] : vector<1xf16> to vector<1x1xf16>
+// CHECK-NEXT:    [[ACC_A:%.+]] = vector.insert_strided_slice [[OUT_VEC_10]], [[ACC_B]] {offsets = [1, 0], strides = [1, 1]} : vector<1x1xf16> into vector<2x2xf16>
+// CHECK-NEXT:    [[IN_SLICE_11:%.+]] = vector.extract_strided_slice %arg0 {offsets = [1, 1], sizes = [1, 1], strides = [1, 1]} : vector<2x2xf8E5M2> to vector<1x1xf8E5M2>
+// CHECK-NEXT:    [[IN_VEC_11:%.+]] = vector.shape_cast [[IN_SLICE_11]] : vector<1x1xf8E5M2> to vector<1xf8E5M2>
+// CHECK-NEXT:    [[SCALE_SCALAR_11:%.+]] = vector.extract [[SCALE_EXT]][1, 1] : f32 from vector<2x2xf32>
+// CHECK-NEXT:    [[PACKED_11:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_11]][0], [[SCALE_SCALAR_11]] : vector<1xf8E5M2> to vector<2xf16>
+// CHECK-NEXT:    [[OUT_SLICE_11:%.+]] = vector.extract_strided_slice [[PACKED_11]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf16> to vector<1xf16>
+// CHECK-NEXT:    [[OUT_VEC_11:%.+]] = vector.shape_cast [[OUT_SLICE_11]] : vector<1xf16> to vector<1x1xf16>
+// CHECK-NEXT:    [[ACC_B:%.+]] = vector.insert_strided_slice [[OUT_VEC_11]], [[ACC_A]] {offsets = [1, 1], strides = [1, 1]} : vector<1x1xf16> into vector<2x2xf16>
+// CHECK-NEXT:    return [[ACC_B]] : vector<2x2xf16>
+func.func @conversion_f8_f16_fallback(%in: vector<2x2xf8E5M2>, %scale: vector<2x2xf8E8M0FNU>) -> vector<2x2xf16> {
+    %ext = arith.scaling_extf %in, %scale : vector<2x2xf8E5M2>, vector<2x2xf8E8M0FNU> to vector<2x2xf16>
+    return %ext : vector<2x2xf16>
+}
+
+// -----
+
+// CHECK-LABEL: @conversion_f4_f16_fallback
+// CHECK:         [[CST:%.+]] = arith.constant dense<0.000000e+00> : vector<2x2xf16>
+// CHECK-NEXT:    [[SCALE_EXT:%.+]] = arith.extf %arg1 : vector<2x2xf8E8M0FNU> to vector<2x2xf32>
+// CHECK-NEXT:    [[IN_SLICE_00:%.+]] = vector.extract_strided_slice %arg0 {offsets = [0, 0], sizes = [1, 1], strides = [1, 1]} : vector<2x2xf4E2M1FN> to vector<1x1xf4E2M1FN>
+// CHECK-NEXT:    [[IN_VEC_00:%.+]] = vector.shape_cast [[IN_SLICE_00]] : vector<1x1xf4E2M1FN> to vector<1xf4E2M1FN>
+// CHECK-NEXT:    [[SCALE_SCALAR_00:%.+]] = vector.extract [[SCALE_EXT]][0, 0] : f32 from vector<2x2xf32>
+// CHECK-NEXT:    [[PACKED_00:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_00]][0], [[SCALE_SCALAR_00]] : vector<1xf4E2M1FN> to vector<2xf16>
+// CHECK-NEXT:    [[OUT_SLICE_00:%.+]] = vector.extract_strided_slice [[PACKED_00]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf16> to vector<1xf16>
+// CHECK-NEXT:    [[OUT_VEC_00:%.+]] = vector.shape_cast [[OUT_SLICE_00]] : vector<1xf16> to vector<1x1xf16>
+// CHECK-NEXT:    [[ACC_A:%.+]] = vector.insert_strided_slice [[OUT_VEC_00]], [[CST]] {offsets = [0, 0], strides = [1, 1]} : vector<1x1xf16> into vector<2x2xf16>
+// CHECK-NEXT:    [[IN_SLICE_01:%.+]] = vector.extract_strided_slice %arg0 {offsets = [0, 1], sizes = [1, 1], strides = [1, 1]} : vector<2x2xf4E2M1FN> to vector<1x1xf4E2M1FN>
+// CHECK-NEXT:    [[IN_VEC_01:%.+]] = vector.shape_cast [[IN_SLICE_01]] : vector<1x1xf4E2M1FN> to vector<1xf4E2M1FN>
+// CHECK-NEXT:    [[SCALE_SCALAR_01:%.+]] = vector.extract [[SCALE_EXT]][0, 1] : f32 from vector<2x2xf32>
+// CHECK-NEXT:    [[PACKED_01:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_01]][0], [[SCALE_SCALAR_01]] : vector<1xf4E2M1FN> to vector<2xf16>
+// CHECK-NEXT:    [[OUT_SLICE_01:%.+]] = vector.extract_strided_slice [[PACKED_01]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf16> to vector<1xf16>
+// CHECK-NEXT:    [[OUT_VEC_01:%.+]] = vector.shape_cast [[OUT_SLICE_01]] : vector<1xf16> to vector<1x1xf16>
+// CHECK-NEXT:    [[ACC_B:%.+]] = vector.insert_strided_slice [[OUT_VEC_01]], [[ACC_A]] {offsets = [0, 1], strides = [1, 1]} : vector<1x1xf16> into vector<2x2xf16>
+// CHECK-NEXT:    [[IN_SLICE_10:%.+]] = vector.extract_strided_slice %arg0 {offsets = [1, 0], sizes = [1, 1], strides = [1, 1]} : vector<2x2xf4E2M1FN> to vector<1x1xf4E2M1FN>
+// CHECK-NEXT:    [[IN_VEC_10:%.+]] = vector.shape_cast [[IN_SLICE_10]] : vector<1x1xf4E2M1FN> to vector<1xf4E2M1FN>
+// CHECK-NEXT:    [[SCALE_SCALAR_10:%.+]] = vector.extract [[SCALE_EXT]][1, 0] : f32 from vector<2x2xf32>
+// CHECK-NEXT:    [[PACKED_10:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_10]][0], [[SCALE_SCALAR_10]] : vector<1xf4E2M1FN> to vector<2xf16>
+// CHECK-NEXT:    [[OUT_SLICE_10:%.+]] = vector.extract_strided_slice [[PACKED_10]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf16> to vector<1xf16>
+// CHECK-NEXT:    [[OUT_VEC_10:%.+]] = vector.shape_cast [[OUT_SLICE_10]] : vector<1xf16> to vector<1x1xf16>
+// CHECK-NEXT:    [[ACC_A:%.+]] = vector.insert_strided_slice [[OUT_VEC_10]], [[ACC_B]] {offsets = [1, 0], strides = [1, 1]} : vector<1x1xf16> into vector<2x2xf16>
+// CHECK-NEXT:    [[IN_SLICE_11:%.+]] = vector.extract_strided_slice %arg0 {offsets = [1, 1], sizes = [1, 1], strides = [1, 1]} : vector<2x2xf4E2M1FN> to vector<1x1xf4E2M1FN>
+// CHECK-NEXT:    [[IN_VEC_11:%.+]] = vector.shape_cast [[IN_SLICE_11]] : vector<1x1xf4E2M1FN> to vector<1xf4E2M1FN>
+// CHECK-NEXT:    [[SCALE_SCALAR_11:%.+]] = vector.extract [[SCALE_EXT]][1, 1] : f32 from vector<2x2xf32>
+// CHECK-NEXT:    [[PACKED_11:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_11]][0], [[SCALE_SCALAR_11]] : vector<1xf4E2M1FN> to vector<2xf16>
+// CHECK-NEXT:    [[OUT_SLICE_11:%.+]] = vector.extract_strided_slice [[PACKED_11]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf16> to vector<1xf16>
+// CHECK-NEXT:    [[OUT_VEC_11:%.+]] = vector.shape_cast [[OUT_SLICE_11]] : vector<1xf16> to vector<1x1xf16>
+// CHECK-NEXT:    [[ACC_B:%.+]] = vector.insert_strided_slice [[OUT_VEC_11]], [[ACC_A]] {offsets = [1, 1], strides = [1, 1]} : vector<1x1xf16> into vector<2x2xf16>
+// CHECK-NEXT:    return [[ACC_B]] : vector<2x2xf16>
+func.func @conversion_f4_f16_fallback(%in: vector<2x2xf4E2M1FN>, %scale: vector<2x2xf8E8M0FNU>) -> vector<2x2xf16> {
+    %ext = arith.scaling_extf %in, %scale : vector<2x2xf4E2M1FN>, vector<2x2xf8E8M0FNU> to vector<2x2xf16>
+    return %ext : vector<2x2xf16>
+}
+
+// -----
+
+// CHECK-LABEL: @conversion_broadcast
+// CHECK:         [[CST:%.+]] = arith.constant dense<0.000000e+00> : vector<8x2x4xf32>
+// CHECK-NEXT:    [[BCAST:%.+]] = vector.broadcast %arg1 : vector<8x2xf8E8M0FNU> to vector<4x8x2xf8E8M0FNU>
+// CHECK-NEXT:    [[IN_CAST:%.+]] = vector.shape_cast %arg0 : vector<8x8xf8E5M2> to vector<8x2x4xf8E5M2>
+// CHECK-NEXT:    [[SCALE_CAST:%.+]] = vector.shape_cast [[BCAST]] : vector<4x8x2xf8E8M0FNU> to vector<8x2x4xf8E8M0FNU>
+// CHECK-NEXT:    [[SCALE_EXT:%.+]] = arith.extf [[SCALE_CAST]] : vector<8x2x4xf8E8M0FNU> to vector<8x2x4xf32>
+// CHECK-NEXT:    [[IN_SLICE_0:%.+]] = vector.extract_strided_slice [[IN_CAST]] {offsets = [0, 0, 0], sizes = [1, 1, 1], strides = [1, 1, 1]} : vector<8x2x4xf8E5M2> to vector<1x1x1xf8E5M2>
+// CHECK-NEXT:    [[IN_VEC_0:%.+]] = vector.shape_cast [[IN_SLICE_0]] : vector<1x1x1xf8E5M2> to vector<1xf8E5M2>
+// CHECK-NEXT:    [[SCALE_SCALAR_0:%.+]] = vector.extract [[SCALE_EXT]][0, 0, 0] : f32 from vector<8x2x4xf32>
+// CHECK-NEXT:    [[PACKED_0:%.+]] = amdgpu.scaled_ext_packed [[IN_VEC_0]][0], [[SCALE_SCALAR_0]] : vector<1xf8E5M2> to vector<2xf32>
+// CHECK-NEXT:    [[OUT_SLICE_0:%.+]] = vector.extract_strided_slice [[PACKED_0]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xf32> to vector<1xf32>
+// CHECK-NEXT:    [[OUT_VEC_0:%.+]] = vector.shape_cast [[OUT_SLICE_0]] : vector<1xf32> to vector<1x1x1xf32>
----------------
tgymnich wrote:

reduced

https://github.com/llvm/llvm-project/pull/146372


More information about the Mlir-commits mailing list