99 lines · plain
1// RUN: mlir-opt %s -convert-vector-to-xegpu -split-input-file | FileCheck %s2 3func.func @store_1D_vector(%vec: vector<8xf32>,4 %source: memref<8x16x32xf32>, %offset: index) {5 vector.store %vec, %source[%offset, %offset, %offset]6 : memref<8x16x32xf32>, vector<8xf32>7 return8}9 10// CHECK-LABEL: @store_1D_vector(11// CHECK-SAME: %[[VEC:.+]]: vector<8xf32>,12// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,13// CHECK-SAME: %[[OFFSET:.+]]: index14// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], %[[OFFSET]], 0]15// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc16// CHECK-SAME: %[[COLLAPSED]]17// CHECK-SAME: memref<32xf32, strided<[1], offset: ?>> -> !xegpu.tensor_desc<8xf32,18// CHECK-SAME: boundary_check = false19// CHECK: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFFSET]]] : vector<8xf32>20 21// -----22 23func.func @store_2D_vector(%vec: vector<8x16xf32>,24 %source: memref<8x16x32xf32>, %offset: index) {25 vector.store %vec, %source[%offset, %offset, %offset]26 : memref<8x16x32xf32>, vector<8x16xf32>27 return28}29 30// CHECK-LABEL: @store_2D_vector(31// CHECK-SAME: %[[VEC:.+]]: vector<8x16xf32>,32// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,33// CHECK-SAME: %[[OFFSET:.+]]: index34// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], 0, 0]35// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc36// CHECK-SAME: %[[COLLAPSED]]37// CHECK-SAME: memref<16x32xf32, strided<[32, 1], offset: ?>> -> !xegpu.tensor_desc<8x16xf32>38// CHECK: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFFSET]], %[[OFFSET]]] : vector<8x16xf32>39 40// -----41 42func.func @store_dynamic_source(%vec: vector<8x16xf32>,43 %source: memref<?x?x?xf32>, %i: index, %j: index, %k: index) {44 vector.store %vec, %source[%i, %j, %k]45 : memref<?x?x?xf32>, vector<8x16xf32>46 return47}48 49// CHECK-LABEL: @store_dynamic_source(50// CHECK-SAME: %[[VEC:.+]]: vector<8x16xf32>,51// CHECK-SAME: %[[SRC:.+]]: memref<?x?x?xf32>,52// CHECK-SAME: %[[OFF0:.+]]: index, %[[OFF1:.+]]: index, %[[OFF2:.+]]: index53// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFF0]], 0, 0]54// CHECK: {{.*}} %[[SIZES:.+]]:2, %[[STRIDES:.+]]:2 = memref.extract_strided_metadata %[[COLLAPSED]]55// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]]56// CHECK: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFF1]], %[[OFF2]]] : vector<8x16xf32>57 58// -----59 60func.func @store_out_of_bounds(%vec: vector<8x16xf32>,61 %source: memref<7x64xf32>, %offset: index) {62 vector.store %vec, %source[%offset, %offset]63 : memref<7x64xf32>, vector<8x16xf32>64 return65}66 67// CHECK-LABEL: @store_out_of_bounds(68// CHECK-SAME: %[[VEC:.+]]: vector<8x16xf32>,69// CHECK-SAME: %[[SRC:.+]]: memref<7x64xf32>,70// CHECK-SAME: %[[OFFSET:.+]]: index71// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc72// CHECK-SAME: %[[SRC]]73// CHECK-SAME: memref<7x64xf32> -> !xegpu.tensor_desc<8x16xf32>74// CHECK: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFFSET]], %[[OFFSET]]] : vector<8x16xf32>75 76// -----77 78func.func @no_store_high_dim_vector(%vec: vector<8x16x32xf32>,79 %source: memref<16x32x64xf32>, %offset: index) {80 vector.store %vec, %source[%offset, %offset, %offset]81 : memref<16x32x64xf32>, vector<8x16x32xf32>82 return83}84 85// CHECK-LABEL: @no_store_high_dim_vector(86// CHECK: vector.store87 88// -----89 90func.func @no_store_zero_dim_vector(%vec: vector<f32>,91 %source: memref<64xf32>, %offset: index) {92 vector.store %vec, %source[%offset]93 : memref<64xf32>, vector<f32>94 return95}96 97// CHECK-LABEL: @no_store_zero_dim_vector(98// CHECK: vector.store99