207 lines · plain
1// RUN: mlir-opt %s -convert-vector-to-xegpu -split-input-file | FileCheck %s2 3gpu.module @xevm_module {4gpu.func @store_1D_vector(%vec: vector<8xf32>, %source: memref<8x16x32xf32>,5 %off1: index, %off2: index, %off3: index,6 %indices: vector<8xindex>, %mask: vector<8xi1>) {7 vector.scatter %source[%off1, %off2, %off3][%indices], %mask, %vec8 : memref<8x16x32xf32>, vector<8xindex>, vector<8xi1>, vector<8xf32>9 gpu.return10}11// CHECK-LABEL: @store_1D_vector(12// CHECK-SAME: %[[VAL:.+]]: vector<8xf32>, %[[SRC:.+]]: memref<8x16x32xf32>,13// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,14// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>, %[[MASK:.+]]: vector<8xi1>) {15// CHECK-COUNT2: arith.muli {{.*}} : index16// CHECK-COUNT2: arith.addi {{.*}} : index17// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>18// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8xindex>19// CHECK: %[[BASE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index20// CHECK: %[[BASE_I64:.+]] = arith.index_cast %[[BASE]] : index to i6421// CHECK: xegpu.store %[[VAL]], %[[BASE_I64]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : vector<8xf32>, i64, vector<8xindex>, vector<8xi1>22// CHECK: gpu.return23}24 25// -----26gpu.module @xevm_module {27gpu.func @store_2D_memref(%vec: vector<8xf32>, %source: memref<8x32xf32>,28 %off1: index, %off2: index,29 %indices: vector<8xindex>, %mask: vector<8xi1>) {30 vector.scatter %source[%off1, %off2][%indices], %mask, %vec31 : memref<8x32xf32>, vector<8xindex>, vector<8xi1>, vector<8xf32>32 gpu.return33}34// CHECK-LABEL: @store_2D_memref(35// CHECK-SAME: %[[VAL:.+]]: vector<8xf32>, %[[SRC:.+]]: memref<8x32xf32>,36// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index37// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>, %[[MASK:.+]]: vector<8xi1>) {38// CHECK-COUNT1: arith.muli {{.*}} : index39// CHECK-COUNT1: arith.addi {{.*}} : index40// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>41// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8xindex>42// CHECK: %[[BASE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x32xf32> -> index43// CHECK: %[[BASE_I64:.+]] = arith.index_cast %[[BASE]] : index to i6444// CHECK: xegpu.store %[[VAL]], %[[BASE_I64]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : vector<8xf32>, i64, vector<8xindex>, vector<8xi1>45// CHECK: gpu.return46}47 48// -----49gpu.module @xevm_module {50gpu.func @store_2D_vector(%vec: vector<8x16xf32>, %source: memref<8x16x32xf32>,51 %off0: index, %off1: index, %off2: index,52 %indices: vector<8x16xindex>, %mask: vector<8x16xi1>) {53 vector.scatter %source[%off0, %off1, %off2][%indices], %mask, %vec54 : memref<8x16x32xf32>, vector<8x16xindex>, vector<8x16xi1>, vector<8x16xf32>55 gpu.return56}57// CHECK-LABEL: @store_2D_vector(58// CHECK-SAME: %[[VAL:.+]]: vector<8x16xf32>, %[[SRC:.+]]: memref<8x16x32xf32>,59// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,60// CHECK-SAME: %[[INDICES:.+]]: vector<8x16xindex>, %[[MASK:.+]]: vector<8x16xi1>) {61// CHECK-COUNT2: arith.muli {{.*}} : index62// CHECK-COUNT2: arith.addi {{.*}} : index63// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8x16xindex>64// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8x16xindex>65// CHECK: %[[BASE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index66// CHECK: %[[BASE_I64:.+]] = arith.index_cast %[[BASE]] : index to i6467// CHECK: xegpu.store %[[VAL]], %[[BASE_I64]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : vector<8x16xf32>, i64, vector<8x16xindex>, vector<8x16xi1>68// CHECK: gpu.return69}70 71// -----72gpu.module @xevm_module {73gpu.func @store_dynamic_source(%vec: vector<8x16xf32>, %source: memref<?x?x?xf32>,74 %off0: index, %off1: index, %off2: index,75 %indices: vector<8x16xindex>, %mask: vector<8x16xi1>) {76 vector.scatter %source[%off0, %off1, %off2][%indices], %mask, %vec77 : memref<?x?x?xf32>, vector<8x16xindex>, vector<8x16xi1>, vector<8x16xf32>78 gpu.return79}80// CHECK-LABEL: @store_dynamic_source(81// CHECK-SAME: %[[VAL:.+]]: vector<8x16xf32>, %[[SRC:.+]]: memref<?x?x?xf32>,82// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,83// CHECK-SAME: %[[INDICES:.+]]: vector<8x16xindex>, %[[MASK:.+]]: vector<8x16xi1>) {84// CHECK: memref.extract_strided_metadata %[[SRC]]85// CHECK-COUNT2: arith.muli {{.*}} : index86// CHECK-COUNT2: arith.addi {{.*}} : index87// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8x16xindex>88// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8x16xindex>89// CHECK: %[[BASE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x?x?xf32> -> index90// CHECK: %[[BASE_I64:.+]] = arith.index_cast %[[BASE]] : index to i6491// CHECK: xegpu.store %[[VAL]], %[[BASE_I64]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : vector<8x16xf32>, i64, vector<8x16xindex>, vector<8x16xi1>92// CHECK: gpu.return93}94 95// -----96gpu.module @xevm_module {97gpu.func @store_dynamic_source2(%vec: vector<8x16xf32>, %source: memref<?x8x16xf32>,98 %off0: index, %off1: index, %off2: index,99 %indices: vector<8x16xindex>, %mask: vector<8x16xi1>) {100 vector.scatter %source[%off0, %off1, %off2][%indices], %mask, %vec101 : memref<?x8x16xf32>, vector<8x16xindex>, vector<8x16xi1>, vector<8x16xf32>102 gpu.return103}104// CHECK-LABEL: @store_dynamic_source2(105// CHECK-SAME: %[[VAL:.+]]: vector<8x16xf32>, %[[SRC:.+]]: memref<?x8x16xf32>,106// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,107// CHECK-SAME: %[[INDICES:.+]]: vector<8x16xindex>, %[[MASK:.+]]: vector<8x16xi1>) {108// CHECK-NOT: memref.extract_strided_metadata %[[SRC]]109// CHECK-COUNT2: arith.muli {{.*}} : index110// CHECK-COUNT2: arith.addi {{.*}} : index111// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8x16xindex>112// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8x16xindex>113// CHECK: %[[BASE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x8x16xf32> -> index114// CHECK: %[[BASE_I64:.+]] = arith.index_cast %[[BASE]] : index to i64115// CHECK: xegpu.store %[[VAL]], %[[BASE_I64]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : vector<8x16xf32>, i64, vector<8x16xindex>, vector<8x16xi1>116// CHECK: gpu.return117}118 119// -----120gpu.module @xevm_module {121gpu.func @non_unit_inner_stride_1D(122 %vec: vector<8xf32>, %source: memref<32xf32, strided<[?], offset: ?>>,123 %off: index, %indices: vector<8xindex>, %mask: vector<8xi1>) {124 vector.scatter %source[%off][%indices], %mask, %vec125 : memref<32xf32, strided<[?], offset: ?>>, vector<8xindex>, vector<8xi1>, vector<8xf32>126 gpu.return127}128// CHECK-LABEL: @non_unit_inner_stride_1D(129// CHECK-SAME: %[[VAL:.+]]: vector<8xf32>, %[[SRC:.+]]: memref<32xf32, strided<[?], offset: ?>>,130// CHECK-SAME: %[[OFF1:.+]]: index,131// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>, %[[MASK:.+]]: vector<8xi1>) {132// CHECK: %[[BB:.+]], %[[M_OFF:.+]], %[[SZ:.+]], %[[STRIDE:.+]] = memref.extract_strided_metadata %[[SRC]]133// CHECK: arith.muli %[[OFF1]], %[[STRIDE]] : index134// CHECK: arith.addi {{.*}} : index135// CHECK: %[[STRD_VEC:.+]] = vector.broadcast %[[STRIDE]] : index to vector<8xindex>136// CHECK: %[[STRD_INDICES:.+]] = arith.muli %[[STRD_VEC:.+]], %[[INDICES]] : vector<8xindex>137// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>138// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[STRD_INDICES]] : vector<8xindex>139// CHECK: %[[BASE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<32xf32, strided<[?], offset: ?>> -> index140// CHECK: %[[BASE_I64:.+]] = arith.index_cast %[[BASE]] : index to i64141// CHECK: xegpu.store %[[VAL]], %[[BASE_I64]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : vector<8xf32>, i64, vector<8xindex>, vector<8xi1>142// CHECK: gpu.return143}144 145// -----146gpu.module @xevm_module {147gpu.func @non_unit_inner_stride_3D(148 %vec: vector<8xf32>,149 %source: memref<4x8x32xf32, strided<[?, 128, 2], offset: ?>>,150 %off0: index, %off1: index, %off2: index,151 %indices: vector<8xindex>, %mask: vector<8xi1>) {152 vector.scatter %source[%off0, %off1, %off2][%indices], %mask, %vec153 : memref<4x8x32xf32, strided<[?, 128, 2], offset: ?>>,154 vector<8xindex>, vector<8xi1>, vector<8xf32>155 gpu.return156}157// CHECK-LABEL: @non_unit_inner_stride_3D(158// CHECK-SAME: %[[VAL:.+]]: vector<8xf32>, %[[SRC:.+]]: memref<4x8x32xf32, strided<[?, 128, 2], offset: ?>>,159// CHECK-SAME: %[[OFF0:.+]]: index, %[[OFF1:.+]]: index, %[[OFF2:.+]]: index,160// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>, %[[MASK:.+]]: vector<8xi1>) {161// CHECK: %[[BB:.+]], %[[M_OFF:.+]], %[[SIZES:.+]]:3, %[[STRIDES:.+]]:3 = memref.extract_strided_metadata %[[SRC]]162// CHECK: arith.muli %[[OFF0]], %[[STRIDES]]#0 : index163// CHECK: arith.addi {{.*}} : index164// CHECK-COUNT2: arith.muli {{.*}} : index165// CHECK-COUNT2: arith.addi {{.*}} : index166// CHECK: %[[STRD_INDICES:.+]] = arith.muli {{.*}}%[[INDICES]]{{.*}} : vector<8xindex>167// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<8xindex>168// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[STRD_INDICES]] : vector<8xindex>169// CHECK: %[[BASE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<4x8x32xf32, strided<[?, 128, 2], offset: ?>> -> index170// CHECK: %[[BASE_I64:.+]] = arith.index_cast %[[BASE]] : index to i64171// CHECK: xegpu.store %[[VAL]], %[[BASE_I64]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : vector<8xf32>, i64, vector<8xindex>, vector<8xi1>172// CHECK: gpu.return173}174 175// -----176gpu.module @xevm_module {177gpu.func @scatter_into_subview(%vals: vector<8xf16>,178 %source: memref<4096x4096xf16>,179 %memref_off: index, %off1: index, %off2: index,180 %indices: vector<8xindex>,181 %mask: vector<8xi1>) {182 %subview = memref.subview %source[%memref_off, %memref_off] [256, 256] [1, 1]183 : memref<4096x4096xf16>184 to memref<256x256xf16, strided<[4096, 1], offset: ?>>185 vector.scatter %subview[%off1, %off2][%indices], %mask, %vals186 : memref<256x256xf16, strided<[4096, 1], offset: ?>>,187 vector<8xindex>, vector<8xi1>, vector<8xf16>188 gpu.return189}190// CHECK-LABEL: @scatter_into_subview(191// CHECK-SAME: %[[VALS:.+]]: vector<8xf16>,192// CHECK-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,193// CHECK-SAME: %[[MEMREF_OFF:.+]]: index, %[[OFF1:.+]]: index, %[[OFF2:.+]]: index,194// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>, %[[MASK:.+]]: vector<8xi1>) {195// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[MEMREF_OFF]], %[[MEMREF_OFF]]] [256, 256] [1, 1]196// CHECK: %[[BB:.+]], %[[OFFSET:.+]],{{.*}},{{.*}} = memref.extract_strided_metadata %[[SUBVIEW]] : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> memref<f16>, index, index, index, index, index197// CHECK: arith.muli {{.*}}%[[OFF1]]{{.*}} : index198// CHECK: arith.addi %[[OFFSET]]{{.*}} : index199// CHECK: %[[BASE_OFF:.+]] = arith.addi {{.*}}%[[OFF2]]{{.*}} : index200// CHECK: %[[SPLAT:.+]] = vector.broadcast %[[BASE_OFF]] : index to vector<8xindex>201// CHECK: %[[LIN:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8xindex>202// CHECK: %[[BASE_IDX:.+]] = memref.extract_aligned_pointer_as_index %[[SUBVIEW]] : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> index203// CHECK: %[[BASE_I64:.+]] = arith.index_cast %[[BASE_IDX]] : index to i64204// CHECK: xegpu.store %[[VALS]], %[[BASE_I64]]{{\[}}%[[LIN]]{{\]}}, %[[MASK]] : vector<8xf16>, i64, vector<8xindex>, vector<8xi1>205// CHECK: gpu.return206}207