498 lines · plain
1// RUN: mlir-opt %s -split-input-file -affine-data-copy-generate="fast-mem-space=0 skip-non-unit-stride-loops" | FileCheck %s2// Small buffer size to trigger fine copies.3// RUN: mlir-opt %s -split-input-file -affine-data-copy-generate="fast-mem-space=0 fast-mem-capacity=1" | FileCheck --check-prefix=CHECK-SMALL %s4 5// Test affine data copy with a memref filter. We use a test pass that invokes6// affine data copy utility on the input loop nest.7// '-test-affine-data-copy-memref-filter' passes the first memref found in an8// affine.load op in the innermost loop as a filter.9// RUN: mlir-opt %s -split-input-file -test-affine-data-copy='memref-filter' | FileCheck %s --check-prefix=FILTER10// RUN: mlir-opt %s -split-input-file -test-affine-data-copy='for-memref-region' | FileCheck %s --check-prefix=MEMREF_REGION11// RUN: mlir-opt %s -split-input-file -test-affine-data-copy='capacity-kib=32' | FileCheck %s --check-prefix=LIMITED-MEM12 13// -copy-skip-non-stride-loops forces the copies to be placed right inside the14// tile space loops, avoiding the sensitivity of copy placement depth to memory15// footprint -- so that one could write a definite test case and not have to16// update it each time something related to the cost functions change.17 18#id = affine_map<(d0) -> (d0)>19#ub = affine_map<(d0) -> (d0 + 128)>20 21// Map used to index the buffer while computing.22// CHECK-DAG: [[$MAP_IDENTITY:map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0)>23// CHECK-DAG: [[$MAP_PLUS_128:map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 + 128)>24 25// CHECK-LABEL: func @matmul26// FILTER-LABEL: func @matmul27// LIMITED-MEM-LABEL: func @matmul28func.func @matmul(%A: memref<4096x4096xf32>, %B: memref<4096x4096xf32>, %C: memref<4096x4096xf32>) -> memref<4096x4096xf32> {29 affine.for %i = 0 to 4096 step 128 {30 affine.for %j = 0 to 4096 step 128 {31 affine.for %k = 0 to 4096 step 128 {32 affine.for %ii = #id(%i) to #ub(%i) {33 affine.for %jj = #id(%j) to #ub(%j) {34 affine.for %kk = #id(%k) to #ub(%k) {35 %5 = affine.load %A[%ii, %kk] : memref<4096x4096xf32>36 %6 = affine.load %B[%kk, %jj] : memref<4096x4096xf32>37 %7 = affine.load %C[%ii, %jj] : memref<4096x4096xf32>38 %8 = arith.mulf %5, %6 : f3239 %9 = arith.addf %7, %8 : f3240 affine.store %9, %C[%ii, %jj] : memref<4096x4096xf32>41 }42 }43 }44 }45 }46 }47 return %C : memref<4096x4096xf32>48 // LIMITED-MEM: return49}50 51// Buffers of size 128x128 get created here for all three matrices.52 53// CHECK: affine.for %[[I:.*]] = 0 to 4096 step 128 {54// CHECK: affine.for %[[J:.*]] = 0 to 4096 step 128 {55// CHECK: [[BUFC:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<128x128xf32>56// The result matrix's copy gets hoisted out.57// Result matrix copy-in.58// CHECK: affine.for %[[II:.*]] = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {59// CHECK: affine.for %[[JJ:.*]] = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {60// CHECK: affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<4096x4096xf32>61// CHECK: affine.store %{{.*}}, [[BUFC]][%[[II]] - %[[I]], %[[JJ]] - %[[J]]] : memref<128x128xf32>62// CHECK: }63// CHECK: }64 65// LHS matrix copy-in.66// CHECK: affine.for %[[K:.*]] = 0 to 4096 step 128 {67// CHECK: [[BUFA:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<128x128xf32>68// CHECK: affine.for %[[II:.*]] = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {69// CHECK: affine.for %[[KK:.*]] = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {70// CHECK: affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<4096x4096xf32>71// CHECK: affine.store %{{.*}}, [[BUFA]][%[[II]] - %[[I]], %[[KK]] - %[[K]]] : memref<128x128xf32>72// CHECK: }73// CHECK: }74 75// RHS matrix copy-in.76// CHECK: [[BUFB:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<128x128xf32>77// CHECK: affine.for %[[KK:.*]] = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {78// CHECK: affine.for %[[JJ:.*]] = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {79// CHECK: affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<4096x4096xf32>80// CHECK: affine.store %{{.*}}, [[BUFB]][%[[KK]] - %[[K]], %[[JJ]] - %[[J]]] : memref<128x128xf32>81// CHECK: }82// CHECK: }83 84// Computation on the fast buffers.85// CHECK: affine.for %{{.*}} = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {86// CHECK: affine.for %{{.*}} = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {87// CHECK: affine.for %{{.*}} = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {88// CHECK: affine.load [[BUFA]][-%{{.*}} + %{{.*}}, -%{{.*}} + %{{.*}}] : memref<128x128xf32>89// CHECK: affine.load [[BUFB]][-%{{.*}} + %{{.*}}, -%{{.*}} + %{{.*}}] : memref<128x128xf32>90// CHECK: affine.load [[BUFC]][-%{{.*}} + %{{.*}}, -%{{.*}} + %{{.*}}] : memref<128x128xf32>91// CHECK: arith.mulf %{{.*}}, %{{.*}} : f3292// CHECK: arith.addf %{{.*}}, %{{.*}} : f3293// CHECK: affine.store %{{.*}}, [[BUFC]][-%{{.*}} + %{{.*}}, -%{{.*}} + %{{.*}}] : memref<128x128xf32>94// CHECK: }95// CHECK: }96// CHECK: }97// CHECK: memref.dealloc [[BUFB]] : memref<128x128xf32>98// CHECK: memref.dealloc [[BUFA]] : memref<128x128xf32>99// CHECK: }100 101// Result matrix copy out.102// CHECK: affine.for %{{.*}} = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {103// CHECK: affine.for %{{.*}} = #[[$MAP_IDENTITY]](%{{.*}}) to #[[$MAP_PLUS_128]](%{{.*}}) {104// CHECK: affine.load [[BUFC]][%{{.*}} - %{{.*}}, %{{.*}} - %{{.*}}] : memref<128x128xf32>105// CHECK: store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<4096x4096xf32>106// CHECK: }107// CHECK: }108// CHECK: memref.dealloc [[BUFC]] : memref<128x128xf32>109// CHECK: }110// CHECK: }111 112// Check that only one memref is copied when memref filter is used.113 114// FILTER: affine.for %{{.*}} = 0 to 4096 step 128 {115// FILTER: memref.alloc() : memref<128x4096xf32>116// FILTER-NOT: memref.alloc()117// FILTER: affine.for118// FILTER: affine.for %{{.*}} = 0 to 4096 {119// FILTER: affine.for %{{.*}} = 0 to 4096 step 128 {120// FILTER-NEXT: affine.for %{{.*}} = 0 to 4096 step 128 {121// FILTER-NEXT: affine.for %{{.*}} = #map{{.*}}(%{{.*}}) to #map{{.*}}(%{{.*}}) {122// FILTER-NEXT: affine.for %{{.*}} = #map{{.*}}(%{{.*}}) to #map{{.*}}(%{{.*}}) {123// FILTER-NEXT: affine.for %{{.*}} = #map{{.*}}(%{{.*}}) to #map{{.*}}(%{{.*}}) {124// FILTER: memref.dealloc %{{.*}} : memref<128x4096xf32>125// FILTER-NOT: memref.dealloc %{{.*}} : memref<128x4096xf32>126 127// -----128 129//130// This test case will lead to single element buffers. These are eventually131// expected to be turned into registers via alloca and mem2reg.132//133// CHECK-SMALL-LABEL: func @single_elt_buffers134// FILTER-LABEL: func @single_elt_buffers135// MEMREF_REGION-LABEL: func @single_elt_buffers136func.func @single_elt_buffers(%arg0: memref<1024x1024xf32>, %arg1: memref<1024x1024xf32>, %arg2: memref<1024x1024xf32>) -> memref<1024x1024xf32> {137 affine.for %i = 0 to 1024 {138 affine.for %j = 0 to 1024 {139 affine.for %k = 0 to 1024 {140 %6 = affine.load %arg1[%k, %j] : memref<1024x1024xf32>141 %7 = affine.load %arg2[%i, %j] : memref<1024x1024xf32>142 %9 = arith.addf %6, %7 : f32143 affine.store %9, %arg2[%i, %j] : memref<1024x1024xf32>144 }145 }146 }147 return %arg2 : memref<1024x1024xf32>148}149// CHECK-SMALL: affine.for %arg{{.*}} = 0 to 1024 {150// CHECK-SMALL: affine.for %arg{{.*}} = 0 to 1024 {151// CHECK-SMALL: memref.alloc() : memref<1x1xf32>152// CHECK-SMALL: affine.load %arg{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>153// CHECK-SMALL: affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>154// CHECK-SMALL: affine.for %arg{{.*}} = 0 to 1024 {155// CHECK-SMALL: memref.alloc() : memref<1x1xf32>156// CHECK-SMALL: affine.load %arg{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>157// CHECK-SMALL: affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>158// CHECK-SMALL: affine.load %{{.*}}[0, 0] : memref<1x1xf32>159// CHECK-SMALL: affine.load %{{.*}}[0, 0] : memref<1x1xf32>160// CHECK-SMALL: arith.addf %{{.*}}, %{{.*}} : f32161// CHECK-SMALL: affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>162// CHECK-SMALL: memref.dealloc %{{.*}} : memref<1x1xf32>163// CHECK-SMALL: }164// CHECK-SMALL: affine.load %{{.*}}[0, 0] : memref<1x1xf32>165// CHECK-SMALL: affine.store %{{.*}}, %arg{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x1024xf32>166// CHECK-SMALL: memref.dealloc %{{.*}} : memref<1x1xf32>167// CHECK-SMALL: }168// CHECK-SMALL: }169// CHECK-SMALL: return170 171// Check that only one memref is copied when memref filter is used.172 173// FILTER: memref.alloc() : memref<1024x1024xf32>174// FILTER-NOT: memref.alloc()175// FILTER: affine.for %{{.*}} = 0 to 1024 {176// FILTER: affine.for %{{.*}} = 0 to 1024 {177// FILTER: affine.for %{{.*}} = 0 to 1024 {178// FILTER-NEXT: affine.for %{{.*}} = 0 to 1024 {179// FILTER-NEXT: affine.for %{{.*}} = 0 to 1024 {180// FILTER: memref.dealloc %{{.*}} : memref<1024x1024xf32>181// FILTER-NOT: memref.dealloc182// FILTER: return183 184// CHeck that only one memref is copied, because for-memref-region is enabled185// (and the first ever encountered load is analyzed).186// MEMREF_REGION: memref.alloc() : memref<1024x1024xf32>187// MEMREF_REGION-NOT: memref.alloc()188// MEMREF_REGION: affine.for %{{.*}} = 0 to 1024 {189// MEMREF_REGION: affine.for %{{.*}} = 0 to 1024 {190// MEMREF_REGION: }191// MEMREF_REGION: }192// MEMREF_REGION-NEXT: affine.for %{{.*}} = 0 to 1024 {193// MEMREF_REGION-NEXT: affine.for %{{.*}} = 0 to 1024 {194// MEMREF_REGION-NEXT: affine.for %{{.*}} = 0 to 1024 {195// MEMREF_REGION: memref.dealloc %{{.*}} : memref<1024x1024xf32>196// MEMREF_REGION-NOT: memref.dealloc197// MEMREF_REGION-NEXT: return198 199// -----200 201// This pattern typically appears with tiling with tile sizes that don't divide202// the loop trip counts.203 204#map_ub = affine_map<(d0) -> (4096, d0 + 100)>205 206// CHECK-DAG: [[$MAP_IDENTITY:map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0)>207// CHECK-DAG: [[$MAP_MIN_UB1:map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 + 100, 4096)>208// CHECK-DAG: [[$MAP_MIN_UB2:map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (4096, d0 + 100)>209 210// CHECK-LABEL: func @min_upper_bound211func.func @min_upper_bound(%A: memref<4096xf32>) -> memref<4096xf32> {212 affine.for %i = 0 to 4096 step 100 {213 affine.for %ii = affine_map<(d0) -> (d0)>(%i) to min #map_ub(%i) {214 %5 = affine.load %A[%ii] : memref<4096xf32>215 %6 = arith.mulf %5, %5 : f32216 affine.store %6, %A[%ii] : memref<4096xf32>217 }218 }219 return %A : memref<4096xf32>220}221// CHECK: affine.for %[[IV1:.*]] = 0 to 4096 step 100222// CHECK: %[[BUF:.*]] = memref.alloc() : memref<100xf32>223// CHECK-NEXT: affine.for %[[IV2:.*]] = #[[$MAP_IDENTITY]](%[[IV1]]) to min #[[$MAP_MIN_UB1]](%[[IV1]]) {224// CHECK-NEXT: affine.load %{{.*}}[%[[IV2]]] : memref<4096xf32>225// CHECK-NEXT: affine.store %{{.*}}, %[[BUF]][%[[IV2]] - %[[IV1]]] : memref<100xf32>226// CHECK-NEXT: }227// CHECK-NEXT: affine.for %[[IV2:.*]] = #[[$MAP_IDENTITY]](%[[IV1]]) to min #[[$MAP_MIN_UB2]](%[[IV1]]) {228// CHECK-NEXT: affine.load %[[BUF]][-%[[IV1]] + %[[IV2]]] : memref<100xf32>229// CHECK-NEXT: arith.mulf230// CHECK-NEXT: affine.store %{{.*}}, %[[BUF]][-%[[IV1]] + %[[IV2]]] : memref<100xf32>231// CHECK-NEXT: }232// CHECK: affine.for %[[IV2:.*]] = #[[$MAP_IDENTITY]](%[[IV1]]) to min #[[$MAP_MIN_UB1]](%[[IV1]]) {233// CHECK-NEXT: affine.load %[[BUF]][%[[IV2]] - %[[IV1]]] : memref<100xf32>234// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%[[IV2]]] : memref<4096xf32>235// CHECK-NEXT: }236// CHECK-NEXT: memref.dealloc %[[BUF]] : memref<100xf32>237// CHECK-NEXT: }238 239// -----240 241// Lower bound is a max; upper bound is a min. This pattern typically appears242// with multi-level tiling when the tile sizes used don't divide loop trip243// counts.244 245#lb = affine_map<()[s0, s1] -> (s0 * 512, s1 * 6)>246#ub = affine_map<()[s0, s1] -> (s0 * 512 + 512, s1 * 6 + 6)>247 248// CHECK-DAG: #[[$LB:.*]] = affine_map<()[s0, s1] -> (s0 * 512, s1 * 6)>249// CHECK-DAG: #[[$UB:.*]] = affine_map<()[s0, s1] -> (s0 * 512 + 512, s1 * 6 + 6)>250 251// CHECK-LABEL: max_lower_bound(%{{.*}}: memref<2048x516xf64>,252// CHECK-SAME: [[i:arg[0-9a-zA-Z_]+]]253// CHECK-SAME: [[j:arg[0-9a-zA-Z_]+]]254func.func @max_lower_bound(%M: memref<2048x516xf64>, %i : index, %j : index) {255 affine.for %ii = 0 to 2048 {256 affine.for %jj = max #lb()[%i, %j] to min #ub()[%i, %j] {257 affine.load %M[%ii, %jj] : memref<2048x516xf64>258 }259 }260 return261}262 263// CHECK: %[[BUF:.*]] = memref.alloc() : memref<2048x6xf64>264// CHECK-NEXT: affine.for %[[ii:.*]] = 0 to 2048 {265// CHECK-NEXT: affine.for %[[jj:.*]] = max #[[$LB]]()[%[[i]], %[[j]]] to min #[[$UB]]()[%[[i]], %[[j]]] {266// CHECK-NEXT: affine.load %{{.*}}[%[[ii]], %[[jj]]] : memref<2048x516xf64>267// CHECK-NEXT: affine.store %{{.*}}, %[[BUF]][%[[ii]], %[[jj]] - symbol(%[[j]]) * 6] : memref<2048x6xf64>268// CHECK-NEXT: }269// CHECK-NEXT: }270// CHECK-NEXT: affine.for %[[ii_:.*]] = 0 to 2048 {271// CHECK-NEXT: affine.for %[[jj_:.*]] = max #[[$LB]]()[%{{.*}}, %{{.*}}] to min #[[$UB]]()[%{{.*}}, %{{.*}}] {272// CHECK-NEXT: affine.load %[[BUF]][%[[ii_]], %[[jj_]] - symbol(%[[j]]) * 6] : memref<2048x6xf64>273// CHECK-NEXT: }274// CHECK-NEXT: }275// CHECK-NEXT: memref.dealloc %[[BUF]] : memref<2048x6xf64>276 277// -----278 279// CHECK-LABEL: func @empty_loops280func.func @empty_loops(%arg0: memref<1024x1024xf64>) {281 // Empty loops - so no copy generation happens.282 affine.for %i = 0 to 0 {283 affine.load %arg0[0, %i] : memref<1024x1024xf64>284 }285 affine.for %i = 0 to -16 {286 affine.load %arg0[0, %i] : memref<1024x1024xf64>287 }288 return289 // CHECK-NOT: memref.alloc290 // CHECK: return291}292 293#map16 = affine_map<(d0, d1, d2) -> (d0 * 40 + d1 * 8 + d2 * 2)>294#map17 = affine_map<(d0, d1, d2) -> (d0 * 40 + d1 * 8 + d2 * 2 + 2)>295// CHECK-LABEL: func @affine_parallel296func.func @affine_parallel(%85:memref<2x5x4x2xi64>) {297 affine.for %arg0 = 0 to 2 {298 affine.parallel (%arg1) = (0) to (5) {299 affine.parallel (%arg2) = (0) to (4) {300 affine.for %arg3 = #map16(%arg0, %arg1, %arg2) to #map17(%arg0, %arg1, %arg2) {301 %105 = affine.load %85[((%arg3 floordiv 2) floordiv 4) floordiv 5, ((%arg3 floordiv 2) floordiv 4) mod 5, (%arg3 floordiv 2) mod 4, %arg3 mod 2] : memref<2x5x4x2xi64>302 }303 }304 }305 }306 // Lower and upper bounds for the region can't be determined for the outermost307 // dimension. No fast buffer generation.308 // CHECK: affine.for309 // CHECK-NEXT: affine.parallel310 // CHECK-NEXT: affine.parallel311 // CHECK-NEXT: affine.for312 // CHECK-NOT: affine.for313 314 315 return316}317 318// CHECK-LABEL: func @index_elt_type319func.func @index_elt_type(%arg0: memref<1x2x4x8xindex>) {320 affine.for %arg1 = 0 to 1 {321 affine.for %arg2 = 0 to 2 {322 affine.for %arg3 = 0 to 4 {323 affine.for %arg4 = 0 to 8 {324 affine.store %arg4, %arg0[%arg1, %arg2, %arg3, %arg4] : memref<1x2x4x8xindex>325 }326 }327 }328 }329 330 // CHECK: affine.for %{{.*}} = 0 to 1331 // CHECK-NEXT: affine.for %{{.*}} = 0 to 2332 // CHECK-NEXT: affine.for %{{.*}} = 0 to 4333 // CHECK-NEXT: affine.for %{{.*}} = 0 to 8334 335 // CHECK: affine.for %{{.*}} = 0 to 2336 // CHECK-NEXT: affine.for %{{.*}} = 0 to 4337 // CHECK-NEXT: affine.for %{{.*}} = 0 to 8338 return339}340 341#map = affine_map<(d0) -> (d0 + 1)>342 343// CHECK-LABEL: func @arbitrary_memory_space344func.func @arbitrary_memory_space() {345 %alloc = memref.alloc() : memref<256x8xi8, #spirv.storage_class<StorageBuffer>>346 affine.for %arg0 = 0 to 32 step 4 {347 %0 = affine.apply #map(%arg0)348 affine.for %arg1 = 0 to 8 step 2 {349 %1 = affine.apply #map(%arg1)350 affine.for %arg2 = 0 to 8 step 2 {351 // CHECK: memref.alloc() : memref<1x7xi8>352 %2 = affine.apply #map(%arg2)353 %3 = affine.load %alloc[%0, %1] : memref<256x8xi8, #spirv.storage_class<StorageBuffer>>354 affine.store %3, %alloc[%0, %2] : memref<256x8xi8, #spirv.storage_class<StorageBuffer>>355 }356 }357 }358 return359}360 361// CHECK-LABEL: zero_ranked362func.func @zero_ranked(%3:memref<480xi1>) {363 %false = arith.constant false364 %4 = memref.alloc() {alignment = 128 : i64} : memref<i1>365 affine.store %false, %4[] : memref<i1>366 %5 = memref.alloc() {alignment = 128 : i64} : memref<i1>367 memref.copy %4, %5 : memref<i1> to memref<i1>368 affine.for %arg0 = 0 to 480 {369 %11 = affine.load %3[%arg0] : memref<480xi1>370 %12 = affine.load %5[] : memref<i1>371 %13 = arith.cmpi slt, %11, %12 : i1372 %14 = arith.select %13, %11, %12 : i1373 affine.store %14, %5[] : memref<i1>374 }375 return376}377 378// CHECK-LABEL: func @scalar_memref_copy_without_dma379func.func @scalar_memref_copy_without_dma() {380 %false = arith.constant false381 %4 = memref.alloc() {alignment = 128 : i64} : memref<i1>382 affine.store %false, %4[] : memref<i1>383 384 // CHECK: %[[FALSE:.*]] = arith.constant false385 // CHECK: %[[MEMREF:.*]] = memref.alloc() {alignment = 128 : i64} : memref<i1>386 // CHECK: affine.store %[[FALSE]], %[[MEMREF]][] : memref<i1>387 return388}389 390// CHECK-LABEL: func @scalar_memref_copy_in_loop391func.func @scalar_memref_copy_in_loop(%3:memref<480xi1>) {392 %false = arith.constant false393 %4 = memref.alloc() {alignment = 128 : i64} : memref<i1>394 affine.store %false, %4[] : memref<i1>395 %5 = memref.alloc() {alignment = 128 : i64} : memref<i1>396 memref.copy %4, %5 : memref<i1> to memref<i1>397 affine.for %arg0 = 0 to 480 {398 %11 = affine.load %3[%arg0] : memref<480xi1>399 %12 = affine.load %5[] : memref<i1>400 %13 = arith.cmpi slt, %11, %12 : i1401 %14 = arith.select %13, %11, %12 : i1402 affine.store %14, %5[] : memref<i1>403 }404 405 // CHECK: %[[FALSE:.*]] = arith.constant false406 // CHECK: %[[MEMREF:.*]] = memref.alloc() {alignment = 128 : i64} : memref<i1>407 // CHECK: affine.store %[[FALSE]], %[[MEMREF]][] : memref<i1>408 // CHECK: %[[TARGET:.*]] = memref.alloc() {alignment = 128 : i64} : memref<i1>409 // CHECK: memref.copy %alloc, %[[TARGET]] : memref<i1> to memref<i1>410 // CHECK: %[[FAST_MEMREF:.*]] = memref.alloc() : memref<480xi1>411 // CHECK: affine.for %{{.*}} = 0 to 480 {412 // CHECK: %{{.*}} = affine.load %arg0[%{{.*}}] : memref<480xi1>413 // CHECK: affine.store %{{.*}}, %[[FAST_MEMREF]][%{{.*}}] : memref<480xi1>414 // CHECK: }415 // CHECK: affine.for %arg1 = 0 to 480 {416 // CHECK: %[[L0:.*]] = affine.load %[[FAST_MEMREF]][%arg1] : memref<480xi1>417 // CHECK: %[[L1:.*]] = affine.load %[[TARGET]][] : memref<i1>418 // CHECK: %[[CMPI:.*]] = arith.cmpi slt, %[[L0]], %[[L1]] : i1419 // CHECK: %[[SELECT:.*]] = arith.select %[[CMPI]], %[[L0]], %[[L1]] : i1420 // CHECK: affine.store %[[SELECT]], %[[TARGET]][] : memref<i1>421 // CHECK: }422 // CHECK: memref.dealloc %[[FAST_MEMREF]] : memref<480xi1>423 return424}425 426// CHECK-LABEL: func @memref_def_inside427// LIMITED-MEM-LABEL: func @memref_def_inside428func.func @memref_def_inside(%arg0: index) {429 %0 = llvm.mlir.constant(1.000000e+00 : f32) : f32430 // No copy generation can happen at this depth given the definition inside.431 affine.for %arg1 = 0 to 29 {432 %alloc_7 = memref.alloc() : memref<1xf32>433 // CHECK: affine.store {{.*}} : memref<1xf32>434 affine.store %0, %alloc_7[0] : memref<1xf32>435 }436 437 // With the limited capacity specified, buffer generation happens at the438 // innermost depth. Tests that copy-placement is proper and respects the439 // memref definition.440 441 // LIMITED-MEM: affine.for %{{.*}} = 0 to 29442 // LIMITED-MEM-NEXT: memref.alloc() : memref<1xf32>443 // LIMITED-MEM-NEXT: memref.alloc() : memref<1xf32>444 // LIMITED-MEM-NEXT: affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>445 // LIMITED-MEM-NEXT: affine.load %{{.*}}[%c0{{.*}}] : memref<1xf32>446 // LIMITED-MEM-NEXT: affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>447 // LIMITED-MEM-NEXT: memref.dealloc %{{.*}} : memref<1xf32>448 return449}450 451// Test with uses across multiple blocks.452 453memref.global "private" constant @__constant_1x2x1xi32_1 : memref<1x2x1xi32> = dense<0> {alignment = 64 : i64}454 455// CHECK-LABEL: func @multiple_blocks456func.func @multiple_blocks(%arg0: index) -> memref<1x2x1xi32> {457 %c1_i32 = arith.constant 1 : i32458 %c3_i32 = arith.constant 3 : i32459 %0 = memref.get_global @__constant_1x2x1xi32_1 : memref<1x2x1xi32>460 %alloc = memref.alloc() {alignment = 64 : i64} : memref<1x2x1xi32>461 memref.copy %0, %alloc : memref<1x2x1xi32> to memref<1x2x1xi32>462 cf.br ^bb1(%alloc : memref<1x2x1xi32>)463^bb1(%1: memref<1x2x1xi32>): // 2 preds: ^bb0, ^bb2464// CHECK: ^bb1(%[[MEM:.*]]: memref<1x2x1xi32>):465 %alloc_0 = memref.alloc() {alignment = 64 : i64} : memref<1x2x1xi1>466 // CHECK: %[[BUF:.*]] = memref.alloc() : memref<1x2x1xi32>467 affine.for %arg1 = 0 to 1 {468 affine.for %arg2 = 0 to 2 {469 affine.for %arg3 = 0 to 1 {470 // CHECK: affine.load %[[BUF]]471 %3 = affine.load %1[%arg1, %arg2, %arg3] : memref<1x2x1xi32>472 %4 = arith.cmpi slt, %3, %c3_i32 : i32473 affine.store %4, %alloc_0[%arg1, %arg2, %arg3] : memref<1x2x1xi1>474 }475 }476 }477 // CHECK: memref.dealloc %[[BUF]]478 %2 = memref.load %alloc_0[%arg0, %arg0, %arg0] : memref<1x2x1xi1>479 cf.cond_br %2, ^bb2, ^bb3480^bb2: // pred: ^bb1481// CHECK: ^bb2482 %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref<1x2x1xi32>483 affine.for %arg1 = 0 to 1 {484 affine.for %arg2 = 0 to 2 {485 affine.for %arg3 = 0 to 1 {486 // Ensure that this reference isn't replaced.487 %3 = affine.load %1[%arg1, %arg2, %arg3] : memref<1x2x1xi32>488 // CHECK: affine.load %[[MEM]]489 %4 = arith.addi %3, %c1_i32 : i32490 affine.store %4, %alloc_1[%arg1, %arg2, %arg3] : memref<1x2x1xi32>491 }492 }493 }494 cf.br ^bb1(%alloc_1 : memref<1x2x1xi32>)495^bb3: // pred: ^bb1496 return %1 : memref<1x2x1xi32>497}498