324 lines · plain
1// RUN: mlir-opt %s --xevm-attach-target='module=xevm_* O=3 chip=pvc' -convert-vector-to-xegpu -split-input-file | FileCheck %s --check-prefix=STORE-ND2// RUN: mlir-opt %s -convert-vector-to-xegpu -split-input-file | FileCheck %s --check-prefix=STORE-SCATTER3 4 5gpu.module @xevm_module {6gpu.func @store_1D_vector(%vec: vector<8xf32>,7 %source: memref<8x16x32xf32>, %offset: index) {8 vector.transfer_write %vec, %source[%offset, %offset, %offset]9 {in_bounds = [true]}10 : vector<8xf32>, memref<8x16x32xf32>11 gpu.return12}13 14// STORE-ND-LABEL: @store_1D_vector(15// STORE-ND-SAME: %[[VEC:.+]]: vector<8xf32>,16// STORE-ND-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,17// STORE-ND-SAME: %[[OFFSET:.+]]: index18// STORE-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], %[[OFFSET]], 0]19// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc20// STORE-ND-SAME: %[[COLLAPSED]]21// STORE-ND-SAME: memref<32xf32, strided<[1], offset: ?>> -> !xegpu.tensor_desc<8xf32,22// STORE-ND-SAME: boundary_check = false23// STORE-ND: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFFSET]]] : vector<8xf32>24 25// STORE-SCATTER-LABEL: @store_1D_vector(26// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<8xf32>,27// STORE-SCATTER-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,28// STORE-SCATTER-DAG: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>29// STORE-SCATTER-DAG: %[[STEP:.+]] = vector.step30// STORE-SCATTER-COUNT2: arith.muli {{.*}} : index31// STORE-SCATTER-COUNT2: arith.addi {{.*}} : index32// STORE-SCATTER-DAG: %[[BCAST:.+]] = vector.broadcast {{.*}} : index to vector<8xindex>33// STORE-SCATTER-DAG: %[[IDX:.+]] = arith.addi %[[BCAST]], %{{.*}} : vector<8xindex>34// STORE-SCATTER-DAG: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index35// STORE-SCATTER-DAG: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i6436// STORE-SCATTER: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8xf32>, i64, vector<8xindex>, vector<8xi1>37}38 39// -----40gpu.module @xevm_module {41gpu.func @store_2D_vector(%vec: vector<8x16xf32>,42 %source: memref<8x16x32xf32>, %offset: index) {43 vector.transfer_write %vec, %source[%offset, %offset, %offset]44 {in_bounds = [true, true]}45 : vector<8x16xf32>, memref<8x16x32xf32>46 gpu.return47}48 49// STORE-ND-LABEL: @store_2D_vector(50// STORE-ND-SAME: %[[VEC:.+]]: vector<8x16xf32>,51// STORE-ND-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,52// STORE-ND-SAME: %[[OFFSET:.+]]: index53// STORE-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], 0, 0]54// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc55// STORE-ND-SAME: %[[COLLAPSED]]56// STORE-ND-SAME: memref<16x32xf32, strided<[32, 1], offset: ?>> -> !xegpu.tensor_desc<8x16xf32,57// STORE-ND-SAME: boundary_check = false58// STORE-ND: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFFSET]], %[[OFFSET]]] : vector<8x16xf32>59 60// STORE-SCATTER-LABEL: @store_2D_vector(61// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<8x16xf32>,62// STORE-SCATTER-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,63// STORE-SCATTER-SAME: %[[OFFSET:.+]]: index64// STORE-SCATTER: %[[CST:.+]] = arith.constant dense<true> : vector<8x16xi1>65// STORE-SCATTER-COUNT2: %[[STEP:.+]] = vector.step66// STORE-SCATTER-COUNT2: vector.shape_cast {{.*}}67// STORE-SCATTER-COUNT2: vector.broadcast {{.*}} : vector<8x16xindex>68// STORE-SCATTER-DAG: %[[BCAST2:.+]] = vector.broadcast {{.*}} : index to vector<8x16xindex>69// STORE-SCATTER-DAG: %[[IDX:.+]] = arith.addi %[[BCAST2]], {{.*}} : vector<8x16xindex>70// STORE-SCATTER-DAG: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index71// STORE-SCATTER-DAG: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i6472// STORE-SCATTER: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8x16xf32>, i64, vector<8x16xindex>, vector<8x16xi1>73}74 75// -----76gpu.module @xevm_module {77gpu.func @store_dynamic_source(%vec: vector<8x16xf32>,78 %source: memref<?x?x?xf32>, %i: index, %j: index, %k: index) {79 vector.transfer_write %vec, %source[%i, %j, %k]80 {in_bounds = [true, true]}81 : vector<8x16xf32>, memref<?x?x?xf32>82 gpu.return83}84 85// STORE-ND-LABEL: @store_dynamic_source(86// STORE-ND-SAME: %[[VEC:.+]]: vector<8x16xf32>,87// STORE-ND-SAME: %[[SRC:.+]]: memref<?x?x?xf32>,88// STORE-ND-SAME: %[[OFF0:.+]]: index, %[[OFF1:.+]]: index, %[[OFF2:.+]]: index89// STORE-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFF0]], 0, 0]90// STORE-ND: {{.*}} %[[SIZES:.+]]:2, %[[STRIDES:.+]]:2 = memref.extract_strided_metadata %[[COLLAPSED]]91// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]]92// STORE-ND: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFF1]], %[[OFF2]]] : vector<8x16xf32>93 94// STORE-SCATTER-LABEL: @store_dynamic_source(95// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<8x16xf32>,96// STORE-SCATTER-SAME: %[[SRC:.+]]: memref<?x?x?xf32>,97// STORE-SCATTER-DAG: %[[CST:.+]] = arith.constant dense<true> : vector<8x16xi1>98// STORE-SCATTER-DAG: memref.extract_strided_metadata %[[SRC]] : memref<?x?x?xf32> -> memref<f32>, index, index, index, index, index, index, index99// STORE-SCATTER-COUNT2: %[[STEP:.+]] = vector.step100// STORE-SCATTER-COUNT2: vector.shape_cast {{.*}}101// STORE-SCATTER-COUNT2: vector.broadcast {{.*}} : vector<8x16xindex>102// STORE-SCATTER-DAG: %[[BCAST2:.+]] = vector.broadcast {{.*}} : index to vector<8x16xindex>103// STORE-SCATTER-DAG: %[[IDX:.+]] = arith.addi %[[BCAST2]], {{.*}} : vector<8x16xindex>104// STORE-SCATTER-DAG: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x?x?xf32> -> index105// STORE-SCATTER-DAG: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64106// STORE-SCATTER: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8x16xf32>, i64, vector<8x16xindex>, vector<8x16xi1>107}108 109// -----110gpu.module @xevm_module {111gpu.func @store_out_of_bounds(%vec: vector<8x16xf32>,112 %source: memref<7x64xf32>, %offset: index) {113 vector.transfer_write %vec, %source[%offset, %offset]114 {in_bounds = [false, true]}115 : vector<8x16xf32>, memref<7x64xf32>116 gpu.return117}118 119// STORE-ND-LABEL: @store_out_of_bounds(120// STORE-ND-SAME: %[[VEC:.+]]: vector<8x16xf32>,121// STORE-ND-SAME: %[[SRC:.+]]: memref<7x64xf32>,122// STORE-ND-SAME: %[[OFFSET:.+]]: index123// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc124// STORE-ND-SAME: %[[SRC]]125// STORE-ND-SAME: memref<7x64xf32> -> !xegpu.tensor_desc<8x16xf32>126// STORE-ND: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFFSET]], %[[OFFSET]]] : vector<8x16xf32>127 128// STORE-SCATTER-LABEL: @store_out_of_bounds(129// STORE-SCATTER: vector.transfer_write130}131 132// -----133gpu.module @xevm_module {134gpu.func @no_store_transposed(%vec: vector<8x16xf32>,135 %source: memref<32x64xf32>, %offset: index) {136 vector.transfer_write %vec, %source[%offset, %offset]137 {permutation_map = affine_map<(d0, d1) -> (d1, d0)>,138 in_bounds = [true, true]}139 : vector<8x16xf32>, memref<32x64xf32>140 gpu.return141}142 143// STORE-ND-LABEL: @no_store_transposed(144// STORE-ND: vector.transfer_write145 146// STORE-SCATTER-LABEL: @no_store_transposed(147// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<8x16xf32>,148// STORE-SCATTER-SAME: %[[SRC:.+]]: memref<32x64xf32>,149// STORE-SCATTER-SAME: %[[OFFSET:.+]]: index150// STORE-SCATTER: %[[CST:.+]] = arith.constant dense<true> : vector<8x16xi1>151// STORE-SCATTER-COUNT2: %[[STEP:.+]] = vector.step152// STORE-SCATTER-COUNT2: vector.shape_cast {{.*}}153// STORE-SCATTER-COUNT2: vector.broadcast {{.*}} : vector<8x16xindex>154// STORE-SCATTER-DAG: %[[BCAST2:.+]] = vector.broadcast {{.*}} : index to vector<8x16xindex>155// STORE-SCATTER-DAG: %[[IDX:.+]] = arith.addi %[[BCAST2]], {{.*}} : vector<8x16xindex>156// STORE-SCATTER-DAG: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<32x64xf32> -> index157// STORE-SCATTER-DAG: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64158// STORE-SCATTER: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8x16xf32>, i64, vector<8x16xindex>, vector<8x16xi1>159}160 161// -----162gpu.module @xevm_module {163gpu.func @store_high_dim_vector(%vec: vector<8x16x32xf32>,164 %source: memref<16x32x64xf32>, %offset: index) {165 vector.transfer_write %vec, %source[%offset, %offset, %offset]166 {in_bounds = [true, true, true]}167 : vector<8x16x32xf32>, memref<16x32x64xf32>168 gpu.return169}170 171// STORE-ND-LABEL: @store_high_dim_vector(172// STORE-ND: vector.transfer_write173 174// STORE-SCATTER-LABEL: @store_high_dim_vector(175// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<8x16x32xf32>,176// STORE-SCATTER-SAME: %[[SRC:.+]]: memref<16x32x64xf32>177// STORE-SCATTER: %[[CST:.+]] = arith.constant dense<true> : vector<8x16x32xi1>178// STORE-SCATTER: %[[CST_0:.+]] = arith.constant dense<64> : vector<16xindex>179// STORE-SCATTER: %[[CST_1:.+]] = arith.constant dense<2048> : vector<8xindex>180// STORE-SCATTER: %[[C2048:.+]] = arith.constant 2048 : index181// STORE-SCATTER: %[[C64:.+]] = arith.constant 64 : index182// STORE-SCATTER-COUNT3: vector.step183// STORE-SCATTER-COUNT3: vector.shape_cast184// STORE-SCATTER-COUNT3: vector.broadcast {{.*}} : vector<8x16x32xindex>185// STORE-SCATTER-COUNT2: arith.addi {{.*}} : vector<8x16x32xindex>186// STORE-SCATTER: %[[BCASTOFF:.+]] = vector.broadcast {{.*}} : index to vector<8x16x32xindex>187// STORE-SCATTER: %[[IDX:.+]] = arith.addi %[[BCASTOFF]], {{.*}} : vector<8x16x32xindex>188// STORE-SCATTER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<16x32x64xf32> -> index189// STORE-SCATTER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64190// STORE-SCATTER: xegpu.store %[[VEC]], %[[COLLAPSE_I]][%[[IDX]]], %[[CST]] : vector<8x16x32xf32>, i64, vector<8x16x32xindex>, vector<8x16x32xi1>191}192 193// -----194gpu.module @xevm_module {195gpu.func @no_store_masked(%vec: vector<4xf32>,196 %source: memref<4xf32>, %offset: index) {197 %mask = arith.constant dense<[0, 1, 0, 1]> : vector<4xi1>198 vector.transfer_write %vec, %source[%offset], %mask199 {in_bounds = [true]}200 : vector<4xf32>, memref<4xf32>201 gpu.return202}203 204// STORE-ND-LABEL: @no_store_masked(205// STORE-ND: vector.transfer_write206 207// STORE-SCATTER-LABEL: @no_store_masked(208// STORE-SCATTER: vector.transfer_write209}210 211// -----212gpu.module @xevm_module {213gpu.func @no_store_tensor(%vec: vector<8x16xf32>,214 %source: tensor<32x64xf32>, %offset: index) -> tensor<32x64xf32> {215 %0 = vector.transfer_write %vec, %source[%offset, %offset]216 {in_bounds = [true, true]}217 : vector<8x16xf32>, tensor<32x64xf32>218 gpu.return %0 : tensor<32x64xf32>219}220 221// STORE-ND-LABEL: @no_store_tensor(222// STORE-ND: vector.transfer_write223 224// STORE-SCATTER-LABEL: @no_store_tensor(225// STORE-SCATTER: vector.transfer_write226}227 228// -----229gpu.module @xevm_module {230gpu.func @no_store_non_unit_inner_stride(%vec: vector<8xf32>,231 %source: memref<32xf32, strided<[?], offset: ?>>, %offset: index) {232 vector.transfer_write %vec, %source[%offset]233 {in_bounds = [true]}234 : vector<8xf32>, memref<32xf32, strided<[?], offset: ?>>235 gpu.return236}237 238// STORE-ND-LABEL: @no_store_non_unit_inner_stride(239// STORE-ND: vector.transfer_write240 241// STORE-SCATTER-LABEL: @no_store_non_unit_inner_stride(242// STORE-SCATTER: vector.transfer_write243}244 245// -----246gpu.module @xevm_module {247gpu.func @no_store_unsupported_map(%vec: vector<8x16xf32>,248 %source: memref<16x32x64xf32>, %offset: index) {249 vector.transfer_write %vec, %source[%offset, %offset, %offset]250 {permutation_map = affine_map<(d0, d1, d2) -> (d0, d2)>,251 in_bounds = [true, true]}252 : vector<8x16xf32>, memref<16x32x64xf32>253 gpu.return254}255 256// STORE-ND-LABEL: @no_store_unsupported_map(257// STORE-ND: vector.transfer_write258 259// STORE-SCATTER-LABEL: @no_store_unsupported_map(260// STORE-SCATTER: vector.transfer_write261}262 263// -----264gpu.module @xevm_module {265gpu.func @no_store_out_of_bounds_1D_vector(%vec: vector<8xf32>,266 %source: memref<8x16x32xf32>, %offset: index) {267 vector.transfer_write %vec, %source[%offset, %offset, %offset]268 {in_bounds = [false]}269 : vector<8xf32>, memref<8x16x32xf32>270 gpu.return271}272 273// STORE-ND-LABEL: @no_store_out_of_bounds_1D_vector(274// STORE-ND: vector.transfer_write275 276// STORE-SCATTER-LABEL: @no_store_out_of_bounds_1D_vector(277// STORE-SCATTER: vector.transfer_write278}279 280// -----281gpu.module @xevm_module {282gpu.func @store_to_subview(%vec: vector<8xf16>,283 %source: memref<4096x4096xf16>, %off1: index, %off2: index) {284 %subview = memref.subview %source[%off1, %off2] [256, 256] [1, 1]285 : memref<4096x4096xf16>286 to memref<256x256xf16, strided<[4096, 1], offset: ?>>287 vector.transfer_write %vec, %subview[%off2, %off2]288 {in_bounds = [true]}289 : vector<8xf16>, memref<256x256xf16, strided<[4096, 1], offset: ?>>290 gpu.return291}292// STORE-ND-LABEL: @store_to_subview(293// STORE-ND-SAME: %[[VEC:.+]]: vector<8xf16>,294// STORE-ND-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,295// STORE-ND-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index296// STORE-ND: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[OFF1]], %[[OFF2]]] [256, 256] [1, 1] : memref<4096x4096xf16> to memref<256x256xf16, strided<[4096, 1], offset: ?>>297// STORE-ND: %[[COLLAPSED:.+]] = memref.subview %[[SUBVIEW]][%[[OFF2]], 0]298// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc299// STORE-ND-SAME: %[[COLLAPSED]]300// STORE-ND-SAME: memref<256xf16, strided<[1], offset: ?>> -> !xegpu.tensor_desc<8xf16,301// STORE-ND-SAME: boundary_check = false302// STORE-ND: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFF2]]] : vector<8xf16>303 304// STORE-SCATTER-LABEL: @store_to_subview(305// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<8xf16>,306// STORE-SCATTER-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,307// STORE-SCATTER-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index308// STORE-SCATTER: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>309// STORE-SCATTER: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[OFF1]], %[[OFF2]]] [256, 256] [1, 1]310// STORE-SCATTER-SAME: : memref<4096x4096xf16> to memref<256x256xf16, strided<[4096, 1], offset: ?>>311// STORE-SCATTER: %[[BB:.+]], %[[OFFSET:.+]], {{.*}}, {{.*}} = memref.extract_strided_metadata %[[SUBVIEW]]312// STORE-SCATTER-SAME: : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> memref<f16>, index, index, index, index, index313// STORE-SCATTER: %[[STEP:.+]] = vector.step : vector<8xindex>314// STORE-SCATTER: arith.muli {{.*}} : index315// STORE-SCATTER: arith.addi %[[OFFSET]]{{.*}} : index316// STORE-SCATTER: arith.addi {{.*}} : index317// STORE-SCATTER: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<8xindex>318// STORE-SCATTER: %[[IDX:.+]] = arith.addi %[[SPLAT]], %[[STEP]] : vector<8xindex>319// STORE-SCATTER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SUBVIEW]]320// STORE-SCATTER-SAME: : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> index321// STORE-SCATTER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64322// STORE-SCATTER: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8xf16>, i64, vector<8xindex>, vector<8xi1>323}324