98 lines · plain
1// RUN: mlir-opt %s -convert-vector-to-xegpu -split-input-file | FileCheck %s2 3func.func @load_1D_vector(%source: memref<8x16x32xf32>, %offset: index) -> vector<8xf32> {4 %0 = vector.load %source[%offset, %offset, %offset]5 : memref<8x16x32xf32>, vector<8xf32>6 return %0 : vector<8xf32>7}8 9// CHECK-LABEL: @load_1D_vector(10// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,11// CHECK-SAME: %[[OFFSET:.+]]: index12// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], %[[OFFSET]], 0]13// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc14// CHECK-SAME: %[[COLLAPSED]]15// CHECK-SAME: memref<32xf32, strided<[1], offset: ?>> -> !xegpu.tensor_desc<8xf32,16// CHECK-SAME: boundary_check = false17// CHECK: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFFSET]]]{{.*}}-> vector<8xf32>18// CHECK: return %[[VEC]]19 20// -----21 22func.func @load_2D_vector(%source: memref<8x16x32xf32>,23 %offset: index) -> vector<8x16xf32> {24 %0 = vector.load %source[%offset, %offset, %offset]25 : memref<8x16x32xf32>, vector<8x16xf32>26 return %0 : vector<8x16xf32>27}28 29// CHECK-LABEL: @load_2D_vector(30// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,31// CHECK-SAME: %[[OFFSET:.+]]: index32// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], 0, 0]33// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc34// CHECK-SAME: %[[COLLAPSED]]35// CHECK-SAME: memref<16x32xf32, strided<[32, 1], offset: ?>> -> !xegpu.tensor_desc<8x16xf32>36// CHECK: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFFSET]], %[[OFFSET]]]{{.*}}-> vector<8x16xf32>37// CHECK: return %[[VEC]]38 39// -----40 41func.func @load_dynamic_source(%source: memref<?x?x?xf32>,42 %i: index, %j: index, %k: index) -> vector<8x16xf32> {43 %0 = vector.load %source[%i, %j, %k]44 : memref<?x?x?xf32>, vector<8x16xf32>45 return %0 : vector<8x16xf32>46}47 48// CHECK-LABEL: @load_dynamic_source(49// CHECK-SAME: %[[SRC:.+]]: memref<?x?x?xf32>,50// CHECK-SAME: %[[OFF0:.+]]: index, %[[OFF1:.+]]: index, %[[OFF2:.+]]: index51// CHECK: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFF0]], 0, 0]52// CHECK: {{.*}} %[[SIZES:.+]]:2, %[[STRIDES:.+]]:2 = memref.extract_strided_metadata %[[COLLAPSED]]53// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[COLLAPSED]]54// CHECK: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFF1]], %[[OFF2]]]{{.*}}-> vector<8x16xf32>55// CHECK: return %[[VEC]]56 57// -----58 59func.func @load_out_of_bounds(%source: memref<7x15xf32>,60 %offset: index) -> vector<8x16xf32> {61 %0 = vector.load %source[%offset, %offset]62 : memref<7x15xf32>, vector<8x16xf32>63 return %0 : vector<8x16xf32>64}65 66// CHECK-LABEL: @load_out_of_bounds(67// CHECK-SAME: %[[SRC:.+]]: memref<7x15xf32>,68// CHECK-SAME: %[[OFFSET:.+]]: index69// CHECK: %[[DESC:.+]] = xegpu.create_nd_tdesc70// CHECK-SAME: %[[SRC]]71// CHECK-SAME: memref<7x15xf32> -> !xegpu.tensor_desc<8x16xf32>72// CHECK: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFFSET]], %[[OFFSET]]]{{.*}}-> vector<8x16xf32>73// CHECK: return %[[VEC]]74 75// -----76 77func.func @no_load_high_dim_vector(%source: memref<16x32x64xf32>,78 %offset: index) -> vector<8x16x32xf32> {79 %0 = vector.load %source[%offset, %offset, %offset]80 : memref<16x32x64xf32>, vector<8x16x32xf32>81 return %0 : vector<8x16x32xf32>82}83 84// CHECK-LABEL: @no_load_high_dim_vector(85// CHECK: vector.load86 87// -----88 89func.func @no_load_zero_dim_vector(%source: memref<64xf32>,90 %offset: index) -> vector<f32> {91 %0 = vector.load %source[%offset]92 : memref<64xf32>, vector<f32>93 return %0 : vector<f32>94}95 96// CHECK-LABEL: @no_load_zero_dim_vector(97// CHECK: vector.load98