252 lines · plain
1// RUN: mlir-opt %s -convert-vector-to-xegpu -split-input-file | FileCheck %s2 3gpu.module @xevm_module {4gpu.func @load_1D_vector(%source: memref<8x16x32xf32>,5 %off1: index, %off2: index, %off3: index,6 %indices: vector<8xindex>, %mask: vector<8xi1>,7 %pass_thru: vector<8xf32>) -> vector<8xf32> {8 %0 = vector.gather %source[%off1, %off2, %off3][%indices], %mask,9 %pass_thru : memref<8x16x32xf32>, vector<8xindex>, vector<8xi1>, vector<8xf32> into vector<8xf32>10 gpu.return %0 : vector<8xf32>11}12// CHECK-LABEL: @load_1D_vector(13// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,14// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,15// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>16// CHECK-SAME: %[[MASK:.+]]: vector<8xi1>17// CHECK-SAME: %[[PASS_THRU:.+]]: vector<8xf32>) -> vector<8xf32> {18// CHECK-COUNT2: arith.muli {{.*}} : index19// CHECK-COUNT2: arith.addi {{.*}} : index20// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>21// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8xindex>22// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index23// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i6424// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf32>25// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PASS_THRU]] : vector<8xi1>, vector<8xf32>26// CHECK: gpu.return %[[RES]] : vector<8xf32>27}28 29// -----30gpu.module @xevm_module {31gpu.func @load_2D_memref(%source: memref<8x32xf32>,32 %off1: index, %off2: index,33 %indices: vector<8xindex>, %mask: vector<8xi1>,34 %pass_thru: vector<8xf32>) -> vector<8xf32> {35 %0 = vector.gather %source[%off1, %off2][%indices], %mask,36 %pass_thru : memref<8x32xf32>, vector<8xindex>, vector<8xi1>, vector<8xf32> into vector<8xf32>37 gpu.return %0 : vector<8xf32>38}39// CHECK-LABEL: @load_2D_memref(40// CHECK-SAME: %[[SRC:.+]]: memref<8x32xf32>,41// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index42// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>43// CHECK-SAME: %[[MASK:.+]]: vector<8xi1>44// CHECK-SAME: %[[PASS_THRU:.+]]: vector<8xf32>) -> vector<8xf32> {45// CHECK-COUNT1: arith.muli {{.*}} : index46// CHECK-COUNT1: arith.addi {{.*}} : index47// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>48// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8xindex>49// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x32xf32> -> index50// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i6451// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf32>52// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PASS_THRU]] : vector<8xi1>, vector<8xf32>53// CHECK: gpu.return %[[RES]] : vector<8xf32>54}55 56// -----57gpu.module @xevm_module {58gpu.func @load_2D_vector(%source: memref<8x16x32xf32>,59 %off0: index, %off1: index, %off2: index,60 %indices: vector<8x16xindex>, %mask: vector<8x16xi1>,61 %pass_thru: vector<8x16xf32>) -> vector<8x16xf32> {62 %0 = vector.gather %source[%off0, %off1, %off2][%indices], %mask,63 %pass_thru : memref<8x16x32xf32>, vector<8x16xindex>, vector<8x16xi1>, vector<8x16xf32> into vector<8x16xf32>64 gpu.return %0 : vector<8x16xf32>65}66// CHECK-LABEL: @load_2D_vector(67// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,68// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,69// CHECK-SAME: %[[INDICES:.+]]: vector<8x16xindex>70// CHECK-SAME: %[[MASK:.+]]: vector<8x16xi1>71// CHECK-SAME: %[[PASS_THRU:.+]]: vector<8x16xf32>) -> vector<8x16xf32> {72// CHECK-COUNT2: arith.muli {{.*}} : index73// CHECK-COUNT2: arith.addi {{.*}} : index74// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8x16xindex>75// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8x16xindex>76// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index77// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i6478// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : i64, vector<8x16xindex>, vector<8x16xi1> -> vector<8x16xf32>79// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PASS_THRU]] : vector<8x16xi1>, vector<8x16xf32>80// CHECK: gpu.return %[[RES]] : vector<8x16xf32>81}82 83// -----84gpu.module @xevm_module {85gpu.func @load_dynamic_source(%source: memref<?x?x?xf32>,86 %off0: index, %off1: index, %off2: index,87 %indices: vector<8x16xindex>, %mask: vector<8x16xi1>,88 %pass_thru: vector<8x16xf32>) -> vector<8x16xf32> {89 %0 = vector.gather %source[%off0, %off1, %off2][%indices], %mask,90 %pass_thru : memref<?x?x?xf32>, vector<8x16xindex>, vector<8x16xi1>, vector<8x16xf32> into vector<8x16xf32>91 gpu.return %0 : vector<8x16xf32>92}93// CHECK-LABEL: @load_dynamic_source(94// CHECK-SAME: %[[SRC:.+]]: memref<?x?x?xf32>,95// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,96// CHECK-SAME: %[[INDICES:.+]]: vector<8x16xindex>97// CHECK-SAME: %[[MASK:.+]]: vector<8x16xi1>98// CHECK-SAME: %[[PASS_THRU:.+]]: vector<8x16xf32>) -> vector<8x16xf32> {99// CHECK: memref.extract_strided_metadata %[[SRC]]100// CHECK-COUNT2: arith.muli {{.*}} : index101// CHECK-COUNT2: arith.addi {{.*}} : index102// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8x16xindex>103// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8x16xindex>104// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x?x?xf32> -> index105// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64106// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : i64, vector<8x16xindex>, vector<8x16xi1> -> vector<8x16xf32>107// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PASS_THRU]] : vector<8x16xi1>, vector<8x16xf32>108// CHECK: gpu.return %[[RES]] : vector<8x16xf32>109}110 111// -----112gpu.module @xevm_module {113gpu.func @load_dynamic_source2(%source: memref<?x8x16xf32>,114 %off0: index, %off1: index, %off2: index,115 %indices: vector<8x16xindex>, %mask: vector<8x16xi1>,116 %pass_thru: vector<8x16xf32>) -> vector<8x16xf32> {117 %0 = vector.gather %source[%off0, %off1, %off2][%indices], %mask,118 %pass_thru : memref<?x8x16xf32>, vector<8x16xindex>, vector<8x16xi1>, vector<8x16xf32> into vector<8x16xf32>119 gpu.return %0 : vector<8x16xf32>120}121// CHECK-LABEL: @load_dynamic_source2(122// CHECK-SAME: %[[SRC:.+]]: memref<?x8x16xf32>,123// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,124// CHECK-SAME: %[[INDICES:.+]]: vector<8x16xindex>125// CHECK-SAME: %[[MASK:.+]]: vector<8x16xi1>126// CHECK-SAME: %[[PASS_THRU:.+]]: vector<8x16xf32>) -> vector<8x16xf32> {127// CHECK-NOT: memref.extract_strided_metadata %[[SRC]]128// CHECK-COUNT2: arith.muli {{.*}} : index129// CHECK-COUNT2: arith.addi {{.*}} : index130// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8x16xindex>131// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8x16xindex>132// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x8x16xf32> -> index133// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64134// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : i64, vector<8x16xindex>, vector<8x16xi1> -> vector<8x16xf32>135// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PASS_THRU]] : vector<8x16xi1>, vector<8x16xf32>136// CHECK: gpu.return %[[RES]] : vector<8x16xf32>137}138 139// -----140gpu.module @xevm_module {141gpu.func @no_load_tensor(%source: tensor<32x64xf32>,142 %off: index, %indices: vector<8x16xindex>,143 %mask: vector<8x16xi1>, %pass_thru: vector<8x16xf32>) -> vector<8x16xf32> {144 %0 = vector.gather %source[%off, %off][%indices], %mask,145 %pass_thru : tensor<32x64xf32>, vector<8x16xindex>, vector<8x16xi1>, vector<8x16xf32> into vector<8x16xf32>146 gpu.return %0 : vector<8x16xf32>147}148// CHECK-LABEL: @no_load_tensor(149// CHECK: vector.gather150}151 152// -----153gpu.module @xevm_module {154gpu.func @gather_from_subview(%source: memref<4096x4096xf16>,155 %memref_off: index, %off1: index, %off2: index,156 %indices: vector<8xindex>,157 %mask: vector<8xi1>,158 %pass_thru: vector<8xf16>) -> vector<8xf16> {159 %subview = memref.subview %source[%memref_off, %memref_off] [256, 256] [1, 1]160 : memref<4096x4096xf16>161 to memref<256x256xf16, strided<[4096, 1], offset: ?>>162 %0 = vector.gather %subview[%off1, %off2][%indices], %mask, %pass_thru163 : memref<256x256xf16, strided<[4096, 1], offset: ?>>,164 vector<8xindex>, vector<8xi1>, vector<8xf16>165 into vector<8xf16>166 gpu.return %0 : vector<8xf16>167}168// CHECK-LABEL: @gather_from_subview(169// CHECK-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,170// CHECK-SAME: %[[MEMREF_OFF:.+]]: index, %[[OFF1:.+]]: index, %[[OFF2:.+]]: index,171// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>,172// CHECK-SAME: %[[MASK:.+]]: vector<8xi1>,173// CHECK-SAME: %[[PASS:.+]]: vector<8xf16>) -> vector<8xf16> {174// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[MEMREF_OFF]], %[[MEMREF_OFF]]] [256, 256] [1, 1]175// CHECK: %[[BB:.+]], %[[OFFSET:.+]],{{.*}},{{.*}} = memref.extract_strided_metadata %[[SUBVIEW]] : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> memref<f16>, index, index, index, index, index176// CHECK: arith.muli {{.*}}%[[OFF1]]{{.*}} : index177// CHECK: arith.addi %[[OFFSET]]{{.*}} : index178// CHECK: %[[BASE_OFF:.+]] = arith.addi {{.*}}%[[OFF2]]{{.*}} : index179// CHECK: %[[SPLAT:.+]] = vector.broadcast %[[BASE_OFF]] : index to vector<8xindex>180// CHECK: %[[LIN:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8xindex>181// CHECK: %[[BASE_IDX:.+]] = memref.extract_aligned_pointer_as_index %[[SUBVIEW]] : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> index182// CHECK: %[[BASE_I64:.+]] = arith.index_cast %[[BASE_IDX]] : index to i64183// CHECK: %[[VEC:.+]] = xegpu.load %[[BASE_I64]]{{\[}}%[[LIN]]{{\]}}, %[[MASK]]184// CHECK-SAME: : i64, vector<8xindex>, vector<8xi1> -> vector<8xf16>185// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PASS]] : vector<8xi1>, vector<8xf16>186// CHECK: gpu.return %[[RES]] : vector<8xf16>187}188 189// -----190gpu.module @xevm_module {191gpu.func @non_unit_inner_stride_1D(192 %source: memref<32xf32, strided<[?], offset: ?>>,193 %off: index, %indices: vector<8xindex>, %mask: vector<8xi1>,194 %pass_thru: vector<8xf32>) -> vector<8xf32> {195 %0 = vector.gather %source[%off][%indices], %mask, %pass_thru196 : memref<32xf32, strided<[?], offset: ?>>,197 vector<8xindex>, vector<8xi1>, vector<8xf32>198 into vector<8xf32>199 gpu.return %0 : vector<8xf32>200}201// CHECK-LABEL: @non_unit_inner_stride_1D(202// CHECK-SAME: %[[SRC:.+]]: memref<32xf32, strided<[?], offset: ?>>,203// CHECK-SAME: %[[OFF1:.+]]: index,204// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>,205// CHECK-SAME: %[[MASK:.+]]: vector<8xi1>, %[[PASS:.+]]: vector<8xf32>) -> vector<8xf32> {206// CHECK: %[[BB:.+]], %[[M_OFF:.+]], %[[SZ:.+]], %[[STRIDE:.+]] = memref.extract_strided_metadata %[[SRC]]207// CHECK: arith.muli %[[OFF1]], %[[STRIDE]] : index208// CHECK: arith.addi {{.*}} : index209// CHECK: %[[STRD_VEC:.+]] = vector.broadcast %[[STRIDE]] : index to vector<8xindex>210// CHECK: %[[STRD_INDICES:.+]] = arith.muli %[[STRD_VEC:.+]], %[[INDICES]] : vector<8xindex>211// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>212// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[STRD_INDICES]] : vector<8xindex>213// CHECK: %[[BASE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<32xf32, strided<[?], offset: ?>> -> index214// CHECK: %[[BASE_I64:.+]] = arith.index_cast %[[BASE]] : index to i64215// CHECK: %[[V:.+]] = xegpu.load %[[BASE_I64]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf32>216// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[V]], %[[PASS]] : vector<8xi1>, vector<8xf32>217// CHECK: gpu.return %[[RES]] : vector<8xf32>218}219 220// -----221gpu.module @xevm_module {222gpu.func @non_unit_inner_stride_3D(223 %source: memref<4x8x32xf32, strided<[?, 128, 2], offset: ?>>,224 %off0: index, %off1: index, %off2: index,225 %indices: vector<8xindex>, %mask: vector<8xi1>,226 %pass_thru: vector<8xf32>) -> vector<8xf32> {227 %0 = vector.gather %source[%off0, %off1, %off2][%indices], %mask, %pass_thru228 : memref<4x8x32xf32, strided<[?, 128, 2], offset: ?>>,229 vector<8xindex>, vector<8xi1>, vector<8xf32>230 into vector<8xf32>231 gpu.return %0 : vector<8xf32>232}233// CHECK-LABEL: @non_unit_inner_stride_3D(234// CHECK-SAME: %[[SRC:.+]]: memref<4x8x32xf32, strided<[?, 128, 2], offset: ?>>,235// CHECK-SAME: %[[OFF0:.+]]: index, %[[OFF1:.+]]: index, %[[OFF2:.+]]: index,236// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>, %[[MASK:.+]]: vector<8xi1>,237// CHECK-SAME: %[[PASS:.+]]: vector<8xf32>) -> vector<8xf32> {238// CHECK: %[[BB:.+]], %[[M_OFF:.+]], %[[SIZES:.+]]:3, %[[STRIDES:.+]]:3 = memref.extract_strided_metadata %[[SRC]]239// CHECK: arith.muli %[[OFF0]], %[[STRIDES]]#0 : index240// CHECK: arith.addi {{.*}} : index241// CHECK-COUNT2: arith.muli {{.*}} : index242// CHECK-COUNT2: arith.addi {{.*}} : index243// CHECK: %[[STRD_INDICES:.+]] = arith.muli {{.*}}%[[INDICES]]{{.*}} : vector<8xindex>244// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<8xindex>245// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[STRD_INDICES]] : vector<8xindex>246// CHECK: %[[BASE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<4x8x32xf32, strided<[?, 128, 2], offset: ?>> -> index247// CHECK: %[[BASE_I64:.+]] = arith.index_cast %[[BASE]] : index to i64248// CHECK: %[[V:.+]] = xegpu.load %[[BASE_I64]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf32>249// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[V]], %[[PASS]] : vector<8xi1>, vector<8xf32>250// CHECK: gpu.return %[[RES]] : vector<8xf32>251}252