181 lines · plain
1// RUN: mlir-opt %s -test-scalar-vector-transfer-lowering -split-input-file | FileCheck %s2// RUN: mlir-opt %s -test-scalar-vector-transfer-lowering=allow-multiple-uses -split-input-file | FileCheck %s --check-prefix=MULTIUSE3 4// CHECK-LABEL: func @transfer_read_0d(5// CHECK-SAME: %[[m:.*]]: memref<?x?x?xf32>, %[[idx:.*]]: index6// CHECK: %[[r:.*]] = memref.load %[[m]][%[[idx]], %[[idx]], %[[idx]]]7// CHECK: return %[[r]]8func.func @transfer_read_0d(%m: memref<?x?x?xf32>, %idx: index) -> f32 {9 %cst = arith.constant 0.0 : f3210 %0 = vector.transfer_read %m[%idx, %idx, %idx], %cst : memref<?x?x?xf32>, vector<f32>11 %1 = vector.extract %0[] : f32 from vector<f32>12 return %1 : f3213}14 15// -----16 17// CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 + s1)>18// CHECK-LABEL: func @transfer_read_1d(19// CHECK-SAME: %[[m:.*]]: memref<?x?x?xf32>, %[[idx:.*]]: index, %[[idx2:.*]]: index20// CHECK: %[[added:.*]] = affine.apply #[[$map]]()[%[[idx]], %[[idx2]]]21// CHECK: %[[r:.*]] = memref.load %[[m]][%[[idx]], %[[idx]], %[[added]]]22// CHECK: return %[[r]]23func.func @transfer_read_1d(%m: memref<?x?x?xf32>, %idx: index, %idx2: index) -> f32 {24 %cst = arith.constant 0.0 : f3225 %c0 = arith.constant 0 : index26 %0 = vector.transfer_read %m[%idx, %idx, %idx], %cst {in_bounds = [true]} : memref<?x?x?xf32>, vector<5xf32>27 %1 = vector.extract %0[%idx2] : f32 from vector<5xf32>28 return %1 : f3229}30 31// -----32 33// CHECK-LABEL: func @tensor_transfer_read_0d(34// CHECK-SAME: %[[t:.*]]: tensor<?x?x?xf32>, %[[idx:.*]]: index35// CHECK: %[[r:.*]] = tensor.extract %[[t]][%[[idx]], %[[idx]], %[[idx]]]36// CHECK: return %[[r]]37func.func @tensor_transfer_read_0d(%t: tensor<?x?x?xf32>, %idx: index) -> f32 {38 %cst = arith.constant 0.0 : f3239 %0 = vector.transfer_read %t[%idx, %idx, %idx], %cst : tensor<?x?x?xf32>, vector<f32>40 %1 = vector.extract %0[] : f32 from vector<f32>41 return %1 : f3242}43 44// -----45 46// CHECK-LABEL: func @transfer_write_0d(47// CHECK-SAME: %[[m:.*]]: memref<?x?x?xf32>, %[[idx:.*]]: index, %[[f:.*]]: f3248// CHECK: memref.store %[[f]], %[[m]][%[[idx]], %[[idx]], %[[idx]]]49func.func @transfer_write_0d(%m: memref<?x?x?xf32>, %idx: index, %f: f32) {50 %0 = vector.broadcast %f : f32 to vector<f32>51 vector.transfer_write %0, %m[%idx, %idx, %idx] : vector<f32>, memref<?x?x?xf32>52 return53}54 55// -----56 57// CHECK-LABEL: func @transfer_write_1d(58// CHECK-SAME: %[[m:.*]]: memref<?x?x?xf32>, %[[idx:.*]]: index, %[[f:.*]]: f3259// CHECK: memref.store %[[f]], %[[m]][%[[idx]], %[[idx]], %[[idx]]]60func.func @transfer_write_1d(%m: memref<?x?x?xf32>, %idx: index, %f: f32) {61 %0 = vector.broadcast %f : f32 to vector<1xf32>62 vector.transfer_write %0, %m[%idx, %idx, %idx] : vector<1xf32>, memref<?x?x?xf32>63 return64}65 66// -----67 68// CHECK-LABEL: func @tensor_transfer_write_0d(69// CHECK-SAME: %[[t:.*]]: tensor<?x?x?xf32>, %[[idx:.*]]: index, %[[f:.*]]: f3270// CHECK: %[[r:.*]] = tensor.insert %[[f]] into %[[t]][%[[idx]], %[[idx]], %[[idx]]]71// CHECK: return %[[r]]72func.func @tensor_transfer_write_0d(%t: tensor<?x?x?xf32>, %idx: index, %f: f32) -> tensor<?x?x?xf32> {73 %0 = vector.broadcast %f : f32 to vector<f32>74 %1 = vector.transfer_write %0, %t[%idx, %idx, %idx] : vector<f32>, tensor<?x?x?xf32>75 return %1 : tensor<?x?x?xf32>76}77 78// -----79 80// CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 + 8)>81// CHECK: #[[$map1:.*]] = affine_map<()[s0] -> (s0 + 1)>82// CHECK-LABEL: func @transfer_read_2d_extract(83// CHECK-SAME: %[[m:.*]]: memref<?x?x?x?xf32>, %[[idx:.*]]: index, %[[idx2:.*]]: index84// CHECK: %[[added:.*]] = affine.apply #[[$map]]()[%[[idx]]]85// CHECK: %[[added1:.*]] = affine.apply #[[$map1]]()[%[[idx]]]86// CHECK: %[[r:.*]] = memref.load %[[m]][%[[idx]], %[[idx]], %[[added]], %[[added1]]]87// CHECK: return %[[r]]88func.func @transfer_read_2d_extract(%m: memref<?x?x?x?xf32>, %idx: index, %idx2: index) -> f32 {89 %cst = arith.constant 0.0 : f3290 %c0 = arith.constant 0 : index91 %0 = vector.transfer_read %m[%idx, %idx, %idx, %idx], %cst {in_bounds = [true, true]} : memref<?x?x?x?xf32>, vector<10x5xf32>92 %1 = vector.extract %0[8, 1] : f32 from vector<10x5xf32>93 return %1 : f3294}95 96// -----97 98// CHECK-LABEL: func @transfer_write_arith_constant(99// CHECK-SAME: %[[m:.*]]: memref<?x?x?xf32>, %[[idx:.*]]: index100// CHECK: %[[cst:.*]] = arith.constant 5.000000e+00 : f32101// CHECK: memref.store %[[cst]], %[[m]][%[[idx]], %[[idx]], %[[idx]]]102func.func @transfer_write_arith_constant(%m: memref<?x?x?xf32>, %idx: index) {103 %cst = arith.constant dense<5.000000e+00> : vector<1x1xf32>104 vector.transfer_write %cst, %m[%idx, %idx, %idx] : vector<1x1xf32>, memref<?x?x?xf32>105 return106}107 108// -----109 110// CHECK-LABEL: func @transfer_read_multi_use(111// CHECK-SAME: %[[m:.*]]: memref<?xf32>, %[[idx:.*]]: index112// CHECK-NOT: memref.load113// CHECK: %[[r:.*]] = vector.transfer_read %[[m]][%[[idx]]]114// CHECK: %[[e0:.*]] = vector.extract %[[r]][0]115// CHECK: %[[e1:.*]] = vector.extract %[[r]][1]116// CHECK: return %[[e0]], %[[e1]]117 118// MULTIUSE-LABEL: func @transfer_read_multi_use(119// MULTIUSE-SAME: %[[m:.*]]: memref<?xf32>, %[[idx0:.*]]: index120// MULTIUSE-NOT: vector.transfer_read121// MULTIUSE: %[[r0:.*]] = memref.load %[[m]][%[[idx0]]122// MULTIUSE: %[[idx1:.*]] = affine.apply123// MULTIUSE: %[[r1:.*]] = memref.load %[[m]][%[[idx1]]124// MULTIUSE: return %[[r0]], %[[r1]]125 126func.func @transfer_read_multi_use(%m: memref<?xf32>, %idx: index) -> (f32, f32) {127 %cst = arith.constant 0.0 : f32128 %0 = vector.transfer_read %m[%idx], %cst {in_bounds = [true]} : memref<?xf32>, vector<16xf32>129 %1 = vector.extract %0[0] : f32 from vector<16xf32>130 %2 = vector.extract %0[1] : f32 from vector<16xf32>131 return %1, %2 : f32, f32132}133 134// -----135 136// Check that patterns don't trigger for an sub-vector (not scalar) extraction.137// CHECK-LABEL: func @subvector_extract(138// CHECK-SAME: %[[m:.*]]: memref<?x?xf32>, %[[idx:.*]]: index139// CHECK-NOT: memref.load140// CHECK: %[[r:.*]] = vector.transfer_read %[[m]][%[[idx]], %[[idx]]]141// CHECK: %[[e0:.*]] = vector.extract %[[r]][0]142// CHECK: return %[[e0]]143 144func.func @subvector_extract(%m: memref<?x?xf32>, %idx: index) -> vector<16xf32> {145 %cst = arith.constant 0.0 : f32146 %0 = vector.transfer_read %m[%idx, %idx], %cst {in_bounds = [true, true]} : memref<?x?xf32>, vector<8x16xf32>147 %1 = vector.extract %0[0] : vector<16xf32> from vector<8x16xf32>148 return %1 : vector<16xf32>149}150 151// -----152 153// CHECK: #[[$MAP:.*]] = affine_map<()[s0, s1] -> (s0 + s1)>154// CHECK-LABEL: func @transfer_read_1d_extract_dynamic(155// CHECK-SAME: %[[MEMREF:.*]]: memref<?xf32>, %[[M_IDX:.*]]: index, %[[E_IDX:.*]]: index156// CHECK: %[[APPLY:.*]] = affine.apply #[[$MAP]]()[%[[M_IDX]], %[[E_IDX]]]157// CHECK: %[[RES:.*]] = memref.load %[[MEMREF]][%[[APPLY]]]158func.func @transfer_read_1d_extract_dynamic(%m: memref<?xf32>, %idx: index,159 %offset: index) -> f32 {160 %cst = arith.constant 0.0 : f32161 %vec = vector.transfer_read %m[%idx], %cst {in_bounds = [true]} : memref<?xf32>, vector<5xf32>162 %elem = vector.extract %vec[%offset] : f32 from vector<5xf32>163 return %elem : f32164}165 166// -----167 168// CHECK: #[[$MAP:.*]] = affine_map<()[s0, s1] -> (s0 + s1)>169// CHECK-LABEL: func @transfer_read_2d_extract_dynamic(170// CHECK-SAME: %[[MEMREF:.*]]: memref<?x?xf32>, %[[ROW_IDX:.*]]: index, %[[COL_IDX:.*]]: index, %[[ROW_OFFSET:.*]]: index, %[[COL_OFFSET:.*]]: index171// CHECK: %[[ROW_APPLY:.*]] = affine.apply #[[$MAP]]()[%[[ROW_IDX]], %[[ROW_OFFSET]]]172// CHECK: %[[COL_APPLY:.*]] = affine.apply #[[$MAP]]()[%[[COL_IDX]], %[[COL_OFFSET]]]173// CHECK: %[[RES:.*]] = memref.load %[[MEMREF]][%[[ROW_APPLY]], %[[COL_APPLY]]]174func.func @transfer_read_2d_extract_dynamic(%m: memref<?x?xf32>, %row_idx: index, %col_idx: index,175 %row_offset: index, %col_offset: index) -> f32 {176 %cst = arith.constant 0.0 : f32177 %vec = vector.transfer_read %m[%row_idx, %col_idx], %cst {in_bounds = [true, true]} : memref<?x?xf32>, vector<10x5xf32>178 %elem = vector.extract %vec[%row_offset, %col_offset] : f32 from vector<10x5xf32>179 return %elem : f32180}181