301 lines · plain
1// RUN: mlir-opt --flatten-memref %s --split-input-file --verify-diagnostics | FileCheck %s2 3func.func @load_scalar_from_memref(%input: memref<4x8xf32, strided<[8, 1], offset: 100>>) -> f32 {4 %c1 = arith.constant 1 : index5 %c2 = arith.constant 2 : index6 %value = memref.load %input[%c1, %c2] : memref<4x8xf32, strided<[8, 1], offset: 100>>7 return %value : f328}9// CHECK-LABEL: func @load_scalar_from_memref10// CHECK-NEXT: %[[C10:.*]] = arith.constant 10 : index11// CHECK-NEXT: %[[REINT:.*]] = memref.reinterpret_cast %arg0 to offset: [100], sizes: [32], strides: [1]12// CHECK-SAME: memref<4x8xf32, strided<[8, 1], offset: 100>> to memref<32xf32, strided<[1], offset: 100>>13// CHECK-NEXT: memref.load %[[REINT]][%[[C10]]] : memref<32xf32, strided<[1], offset: 100>>14 15 16// -----17 18func.func @load_scalar_from_memref_dynamic_dim(%input: memref<?x?xf32, strided<[?, ?], offset: ?>>, %row: index, %col: index) -> f32 {19 %value = memref.load %input[%col, %row] : memref<?x?xf32, strided<[?, ?], offset: ?>>20 return %value : f3221}22 23// CHECK: #[[MAP:.*]] = affine_map<()[s0, s1, s2, s3] -> (s0 * s1 + s2 * s3)>24// CHECK: #[[MAP1:.*]] = affine_map<()[s0, s1, s2, s3] -> (s0 * s1, s2 * s3)>25// CHECK: func @load_scalar_from_memref_dynamic_dim26// CHECK-SAME: (%[[ARG0:.*]]: memref<?x?xf32, strided<[?, ?], offset: ?>>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)27// CHECK: %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[ARG0]]28// CHECK: %[[IDX:.*]] = affine.apply #[[MAP]]()[%[[ARG2]], %[[STRIDES]]#0, %[[ARG1]], %[[STRIDES]]#1]29// CHECK: %[[SIZE:.*]] = affine.max #[[MAP1]]()[%[[STRIDES]]#0, %[[SIZES]]#0, %[[STRIDES]]#1, %[[SIZES]]#1]30// CHECK: %[[REINT:.*]] = memref.reinterpret_cast %arg0 to offset: [%[[OFFSET]]], sizes: [%[[SIZE]]], strides: [1] : memref<?x?xf32, strided<[?, ?], offset: ?>> to memref<?xf32, strided<[1], offset: ?>> 31// CHECK: memref.load %[[REINT]][%[[IDX]]]32 33// -----34 35func.func @load_scalar_from_memref_static_dim(%input: memref<8x12xf32, strided<[24, 2], offset: 100>>) -> f32 {36 %c7 = arith.constant 7 : index37 %c10 = arith.constant 10 : index38 %value = memref.load %input[%c7, %c10] : memref<8x12xf32, strided<[24, 2], offset: 100>>39 return %value : f3240}41 42// CHECK-LABEL: func @load_scalar_from_memref_static_dim43// CHECK-SAME: (%[[ARG0:.*]]: memref<8x12xf32, strided<[24, 2], offset: 100>>)44// CHECK: %[[C188:.*]] = arith.constant 188 : index45// CHECK: %[[REINT:.*]] = memref.reinterpret_cast %[[ARG0]] to offset: [100], sizes: [192], strides: [1] : memref<8x12xf32, strided<[24, 2], offset: 100>> to memref<192xf32, strided<[1], offset: 100>>46// CHECK: memref.load %[[REINT]][%[[C188]]] : memref<192xf32, strided<[1], offset: 100>>47 48// -----49 50func.func @store_scalar_from_memref_padded(%input: memref<4x8xf32, strided<[18, 2], offset: 100>>, %row: index, %col: index, %value: f32) {51 memref.store %value, %input[%col, %row] : memref<4x8xf32, strided<[18, 2], offset: 100>>52 return53}54// CHECK: #[[MAP:.*]] = affine_map<()[s0, s1] -> (s0 * 18 + s1 * 2)>55// CHECK: func @store_scalar_from_memref_padded56// CHECK-SAME: (%[[ARG0:.*]]: memref<4x8xf32, strided<[18, 2], offset: 100>>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: f32)57// CHECK: %[[IDX:.*]] = affine.apply #[[MAP]]()[%[[ARG2]], %[[ARG1]]]58// CHECK: %[[REINT:.*]] = memref.reinterpret_cast %[[ARG0]]59// CHECK: memref.store %[[ARG3]], %[[REINT]][%[[IDX]]] : memref<72xf32, strided<[1], offset: 100>>60 61// -----62 63func.func @store_scalar_from_memref_dynamic_dim(%input: memref<?x?xf32, strided<[?, ?], offset: ?>>, %row: index, %col: index, %value: f32) {64 memref.store %value, %input[%col, %row] : memref<?x?xf32, strided<[?, ?], offset: ?>>65 return66}67// CHECK: #[[MAP:.*]] = affine_map<()[s0, s1, s2, s3] -> (s0 * s1 + s2 * s3)>68// CHECK: #[[MAP1:.*]] = affine_map<()[s0, s1, s2, s3] -> (s0 * s1, s2 * s3)>69// CHECK: func @store_scalar_from_memref_dynamic_dim70// CHECK-SAME: (%[[ARG0:.*]]: memref<?x?xf32, strided<[?, ?], offset: ?>>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: f32)71// CHECK: %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[ARG0]]72// CHECK: %[[IDX:.*]] = affine.apply #[[MAP]]()[%[[ARG2]], %[[STRIDES]]#0, %[[ARG1]], %[[STRIDES]]#1]73// CHECK: %[[SIZE:.*]] = affine.max #[[MAP1]]()[%[[STRIDES]]#0, %[[SIZES]]#0, %[[STRIDES]]#1, %[[SIZES]]#1]74// CHECK: %[[REINT:.*]] = memref.reinterpret_cast %[[ARG0]] to offset: [%[[OFFSET]]], sizes: [%[[SIZE]]], strides: [1]75// CHECK: memref.store %[[ARG3]], %[[REINT]][%[[IDX]]]76 77// -----78 79func.func @load_vector_from_memref(%input: memref<4x8xf32>) -> vector<8xf32> {80 %c3 = arith.constant 3 : index81 %c6 = arith.constant 6 : index82 %value = vector.load %input[%c3, %c6] : memref<4x8xf32>, vector<8xf32>83 return %value : vector<8xf32>84}85// CHECK-LABEL: func @load_vector_from_memref86// CHECK: %[[C30:.*]] = arith.constant 3087// CHECK-NEXT: %[[REINT:.*]] = memref.reinterpret_cast %arg0 to offset: [0], sizes: [32], strides: [1]88// CHECK-NEXT: vector.load %[[REINT]][%[[C30]]]89 90// -----91 92func.func @load_vector_from_memref_odd(%input: memref<3x7xi2>) -> vector<3xi2> {93 %c1 = arith.constant 1 : index94 %c3 = arith.constant 3 : index95 %value = vector.load %input[%c1, %c3] : memref<3x7xi2>, vector<3xi2>96 return %value : vector<3xi2>97}98// CHECK-LABEL: func @load_vector_from_memref_odd99// CHECK: %[[C10:.*]] = arith.constant 10 : index100// CHECK-NEXT: %[[REINT:.*]] = memref.reinterpret_cast101// CHECK-NEXT: vector.load %[[REINT]][%[[C10]]]102 103// -----104 105func.func @load_vector_from_memref_dynamic(%input: memref<3x7xi2>, %row: index, %col: index) -> vector<3xi2> {106 %value = vector.load %input[%col, %row] : memref<3x7xi2>, vector<3xi2>107 return %value : vector<3xi2>108}109// CHECK: #[[MAP:.*]] = affine_map<()[s0, s1] -> (s0 * 7 + s1)>110// CHECK: func @load_vector_from_memref_dynamic111// CHECK: %[[IDX:.*]] = affine.apply #[[MAP]]()112// CHECK: %[[REINT:.*]] = memref.reinterpret_cast113// CHECK: vector.load %[[REINT]][%[[IDX]]] : memref<21xi2, strided<[1]>>, vector<3xi2>114 115// -----116 117func.func @store_vector_to_memref_odd(%input: memref<3x7xi2>, %value: vector<3xi2>) {118 %c1 = arith.constant 1 : index119 %c3 = arith.constant 3 : index120 vector.store %value, %input[%c1, %c3] : memref<3x7xi2>, vector<3xi2>121 return122}123// CHECK-LABEL: func @store_vector_to_memref_odd124// CHECK-SAME: (%[[ARG0:.*]]: memref<3x7xi2>, %[[ARG1:.*]]: vector<3xi2>)125// CHECK: %[[C10:.*]] = arith.constant 10 : index126// CHECK-NEXT: %[[REINT:.*]] = memref.reinterpret_cast127// CHECK-NEXT: vector.store %[[ARG1]], %[[REINT]][%[[C10]]] : memref<21xi2, strided<[1]>128 129// -----130 131func.func @store_vector_to_memref_dynamic(%input: memref<3x7xi2>, %value: vector<3xi2>, %row: index, %col: index) {132 vector.store %value, %input[%col, %row] : memref<3x7xi2>, vector<3xi2>133 return134}135// CHECK: #[[MAP:.*]] = affine_map<()[s0, s1] -> (s0 * 7 + s1)>136// CHECK: func @store_vector_to_memref_dynamic137// CHECK-SAME: (%[[ARG0:.*]]: memref<3x7xi2>, %[[ARG1:.*]]: vector<3xi2>, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index)138// CHECK: %[[IDX:.*]] = affine.apply #[[MAP]]()[%[[ARG3]], %[[ARG2]]]139// CHECK: %[[REINT:.*]] = memref.reinterpret_cast %[[ARG0]] to offset: [0], sizes: [21], strides: [1]140// CHECK: vector.store %[[ARG1]], %[[REINT]][%[[IDX]]]141 142// -----143 144func.func @mask_store_vector_to_memref_odd(%input: memref<3x7xi2>, %value: vector<3xi2>, %mask: vector<3xi1>) {145 %c1 = arith.constant 1 : index146 %c3 = arith.constant 3 : index147 vector.maskedstore %input[%c1, %c3], %mask, %value : memref<3x7xi2>, vector<3xi1>, vector<3xi2>148 return149}150// CHECK-LABEL: func @mask_store_vector_to_memref_odd151// CHECK-SAME: (%[[ARG0:.*]]: memref<3x7xi2>, %[[ARG1:.*]]: vector<3xi2>, %[[ARG2:.*]]: vector<3xi1>)152// CHECK: %[[C10:.*]] = arith.constant 10 : index153// CHECK-NEXT: %[[REINT:.*]] = memref.reinterpret_cast154// CHECK: vector.maskedstore %[[REINT]][%[[C10]]], %[[ARG2]], %[[ARG1]]155 156// -----157 158func.func @mask_store_vector_to_memref_dynamic(%input: memref<3x7xi2>, %value: vector<3xi2>, %row: index, %col: index, %mask: vector<3xi1>) {159 vector.maskedstore %input[%col, %row], %mask, %value : memref<3x7xi2>, vector<3xi1>, vector<3xi2>160 return161}162// CHECK: #map = affine_map<()[s0, s1] -> (s0 * 7 + s1)>163// CHECK: func @mask_store_vector_to_memref_dynamic164// CHECK-SAME: (%[[ARG0:.*]]: memref<3x7xi2>, %[[ARG1:.*]]: vector<3xi2>, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: vector<3xi1>)165// CHECK: %[[IDX:.*]] = affine.apply #map()[%[[ARG3]], %[[ARG2]]]166// CHECK: %[[REINT:.*]] = memref.reinterpret_cast %[[ARG0]]167// CHECK: vector.maskedstore %[[REINT]][%[[IDX]]], %[[ARG4]], %[[ARG1]]168 169// -----170func.func @mask_load_vector_from_memref_odd(%input: memref<3x7xi2>, %mask: vector<3xi1>, %passthru: vector<3xi2>) -> vector<3xi2> {171 %c1 = arith.constant 1 : index172 %c3 = arith.constant 3 : index173 %result = vector.maskedload %input[%c1, %c3], %mask, %passthru : memref<3x7xi2>, vector<3xi1>, vector<3xi2> into vector<3xi2>174 return %result : vector<3xi2>175}176// CHECK-LABEL: func @mask_load_vector_from_memref_odd177// CHECK-SAME: (%[[ARG0:.*]]: memref<3x7xi2>, %[[MASK:.*]]: vector<3xi1>, %[[PASSTHRU:.*]]: vector<3xi2>)178// CHECK: %[[C10:.*]] = arith.constant 10 : index179// CHECK: %[[REINT:.*]] = memref.reinterpret_cast %[[ARG0]] to offset: [0], sizes: [21], strides: [1]180// CHECK: vector.maskedload %[[REINT]][%[[C10]]], %[[MASK]], %[[PASSTHRU]]181 182// -----183 184func.func @mask_load_vector_from_memref_dynamic(%input: memref<3x7xi2>, %row: index, %col: index, %mask: vector<3xi1>, %passthru: vector<3xi2>) -> vector<3xi2> {185 %result = vector.maskedload %input[%col, %row], %mask, %passthru : memref<3x7xi2>, vector<3xi1>, vector<3xi2> into vector<3xi2>186 return %result : vector<3xi2>187}188// CHECK: #[[MAP:.*]] = affine_map<()[s0, s1] -> (s0 * 7 + s1)>189// CHECK: func @mask_load_vector_from_memref_dynamic190// CHECK-SAME: (%[[ARG0:.*]]: memref<3x7xi2>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: vector<3xi1>, %[[ARG4:.*]]: vector<3xi2>)191// CHECK: %[[IDX:.*]] = affine.apply #[[MAP]]()[%[[ARG2]], %[[ARG1]]]192// CHECK: %[[REINT:.*]] = memref.reinterpret_cast %[[ARG0]]193// CHECK: vector.maskedload %[[REINT]][%[[IDX]]], %[[ARG3]]194 195// -----196 197func.func @transfer_read_memref(%input: memref<4x8xi2>, %value: vector<8xi2>, %row: index, %col: index) -> vector<8xi2> {198 %c0 = arith.constant 0 : i2199 %0 = vector.transfer_read %input[%col, %row], %c0 {in_bounds = [true]} : memref<4x8xi2>, vector<8xi2>200 return %0 : vector<8xi2>201}202 203// CHECK: #[[MAP:.*]] = affine_map<()[s0, s1] -> (s0 * 8 + s1)>204// CHECK: func @transfer_read_memref205// CHECK-SAME: (%[[ARG0:.*]]: memref<4x8xi2>, %[[ARG1:.*]]: vector<8xi2>, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index)206// CHECK: %[[C0:.*]] = arith.constant 0 : i2207// CHECK: %[[IDX:.*]] = affine.apply #[[MAP]]()[%[[ARG3]], %[[ARG2]]]208// CHECK-NEXT: %[[REINT:.*]] = memref.reinterpret_cast %[[ARG0]]209// CHECK-NEXT: vector.transfer_read %[[REINT]][%[[IDX]]], %[[C0]]210 211// -----212 213func.func @transfer_read_memref_not_inbound(%input: memref<4x8xi2>, %value: vector<8xi2>, %row: index, %col: index) -> vector<8xi2> {214 %c0 = arith.constant 0 : i2215 %0 = vector.transfer_read %input[%col, %row], %c0 {in_bounds = [false]} : memref<4x8xi2>, vector<8xi2>216 return %0 : vector<8xi2>217}218 219// CHECK-LABEL: func @transfer_read_memref_not_inbound220// CHECK-SAME: (%[[ARG0:.*]]: memref<4x8xi2>, %[[ARG1:.*]]: vector<8xi2>, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index)221// CHECK: vector.transfer_read %[[ARG0]][%[[ARG3]], %[[ARG2]]]222 223// -----224 225func.func @transfer_read_memref_non_id(%input: memref<4x8xi2>, %value: vector<8xi2>, %row: index, %col: index) -> vector<8xi2> {226 %c0 = arith.constant 0 : i2227 %0 = vector.transfer_read %input[%col, %row], %c0 {permutation_map = affine_map<(d0, d1) -> (d0)>, in_bounds = [true]} : memref<4x8xi2>, vector<8xi2>228 return %0 : vector<8xi2>229}230 231// CHECK-LABEL: func @transfer_read_memref_non_id232// CHECK-SAME: (%[[ARG0:.*]]: memref<4x8xi2>, %[[ARG1:.*]]: vector<8xi2>, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index)233// CHECK: vector.transfer_read %[[ARG0]][%[[ARG3]], %[[ARG2]]]234 235// -----236 237func.func @transfer_write_memref(%input: memref<4x8xi2>, %value: vector<8xi2>, %row: index, %col: index) {238 vector.transfer_write %value, %input[%col, %row] {in_bounds = [true]} : vector<8xi2>, memref<4x8xi2>239 return240}241 242// CHECK: #[[MAP:.*]] = affine_map<()[s0, s1] -> (s0 * 8 + s1)>243// CHECK: func @transfer_write_memref244// CHECK-SAME: (%[[ARG0:.*]]: memref<4x8xi2>, %[[ARG1:.*]]: vector<8xi2>, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index)245// CHECK: %[[IDX:.*]] = affine.apply #[[MAP]]()[%[[ARG3]], %[[ARG2]]]246// CHECK: %[[REINT:.*]] = memref.reinterpret_cast %[[ARG0]]247// CHECK: vector.transfer_write %[[ARG1]], %[[REINT]][%[[IDX]]]248 249// -----250 251func.func @alloc() -> memref<4x8xf32> {252 %0 = memref.alloc() : memref<4x8xf32>253 return %0 : memref<4x8xf32>254}255 256// CHECK-LABEL: func @alloc257// CHECK-SAME: () -> memref<4x8xf32>258// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc() : memref<32xf32, strided<[1]>>259// CHECK-NEXT: %[[REINT:.*]] = memref.reinterpret_cast %[[ALLOC]] to offset: [0], sizes: [4, 8], strides: [8, 1] : memref<32xf32, strided<[1]>> to memref<4x8xf32>260 261// -----262 263func.func @alloca() -> memref<4x8xf32> {264 %0 = memref.alloca() : memref<4x8xf32>265 return %0 : memref<4x8xf32>266}267 268// CHECK-LABEL: func.func @alloca() -> memref<4x8xf32>269// CHECK: %[[ALLOC:.*]] = memref.alloca() : memref<32xf32, strided<[1]>>270// CHECK: %[[REINT:.*]] = memref.reinterpret_cast %[[ALLOC]] to offset: [0], sizes: [4, 8], strides: [8, 1] : memref<32xf32, strided<[1]>> to memref<4x8xf32>271 272// -----273 274func.func @chained_alloc_load() -> vector<8xf32> {275 %c3 = arith.constant 3 : index276 %c6 = arith.constant 6 : index277 %0 = memref.alloc() : memref<4x8xf32>278 %value = vector.load %0[%c3, %c6] : memref<4x8xf32>, vector<8xf32>279 return %value : vector<8xf32>280}281 282// CHECK-LABEL: func @chained_alloc_load283// CHECK-SAME: () -> vector<8xf32>284// CHECK-NEXT: %[[C30:.*]] = arith.constant 30 : index285// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc() : memref<32xf32, strided<[1]>>286// CHECK-NEXT: vector.load %[[ALLOC]][%[[C30]]] : memref<32xf32, strided<[1]>>, vector<8xf32>287 288// -----289 290func.func @load_scalar_from_memref_static_dim_col_major(%input: memref<4x8xf32, strided<[1, 4], offset: 100>>, %row: index, %col: index) -> f32 {291 %value = memref.load %input[%col, %row] : memref<4x8xf32, strided<[1, 4], offset: 100>>292 return %value : f32293}294 295// CHECK: #[[MAP:.*]] = affine_map<()[s0, s1] -> (s0 + s1 * 4)>296// CHECK: func @load_scalar_from_memref_static_dim_col_major297// CHECK-SAME: (%[[ARG0:.*]]: memref<4x8xf32, strided<[1, 4], offset: 100>>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)298// CHECK: %[[IDX:.*]] = affine.apply #[[MAP]]()[%[[ARG2]], %[[ARG1]]]299// CHECK: %[[REINT:.*]] = memref.reinterpret_cast %[[ARG0]] to offset: [100], sizes: [32], strides: [1] : memref<4x8xf32, strided<[1, 4], offset: 100>> to memref<32xf32, strided<[1], offset: 100>>300// CHECK: memref.load %[[REINT]][%[[IDX]]] : memref<32xf32, strided<[1], offset: 100>>301