621 lines · plain
1// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -affine-data-copy-generate="generate-dma=1 fast-mem-space=2 skip-non-unit-stride-loops" -verify-diagnostics | FileCheck %s2// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -affine-data-copy-generate="generate-dma=1 fast-mem-capacity=16 fast-mem-space=2" | FileCheck %s --check-prefix FAST-MEM-16KB3 4// We run most test cases with -copy-skip-non-unit-stride-loops to allow testing5// DMA generation at inner levels easily - since the DMA generation would6// otherwise always generate DMAs at the outermost level (default for fast mem7// capacity is infinite). Using a specific capacity makes it harder to write8// a test case as one would have to calculate total footprints. With9// -copy-skip-non-unit-stride-loops, non-unit strides will always be skipped and10// its inner loops will be traversed till a unit stride loop is found (or the11// innermost block is reached).12 13// -----14 15// CHECK-LABEL: func @loop_nest_1d() {16func.func @loop_nest_1d() {17 %A = memref.alloc() : memref<256 x f32>18 %B = memref.alloc() : memref<512 x f32>19 %F = memref.alloc() : memref<256 x f32, 2>20 // First DMA buffer.21 // CHECK: memref.alloc() : memref<256xf32>22 // CHECK: memref.alloc() : memref<256xf32, 2>23 // Tag for first DMA.24 // CHECK: memref.alloc() : memref<1xi32>25 // First DMA transfer.26 // CHECK: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<256xf32>, memref<256xf32, 2>, memref<1xi32>27 // CHECK: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>28 // Second DMA buffer.29 // CHECK: memref.alloc() : memref<256xf32, 2>30 // Tag for second DMA.31 // CHECK: memref.alloc() : memref<1xi32>32 // Second DMA transfer.33 // CHECK: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<512xf32>, memref<256xf32, 2>, memref<1xi32>34 // CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>35 // CHECK: affine.for %[[IV:.*]] = 0 to 256 {36 // CHECK-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<256xf32, 2>37 // Buffer for '%{{.*}}' in faster memref space is of smaller size: 256xf3238 // Affine map for load on B is composed and becomes identity.39 // CHECK: affine.load %{{.*}}[%[[IV]]] : memref<256xf32, 2>40 // Already in faster memory space.41 // CHECK: affine.load %{{.*}}[%[[IV]]] : memref<256xf32, 2>42 // CHECK-NEXT: }43 // CHECK-NEXT: dealloc %{{.*}} : memref<1xi32>44 // CHECK-NEXT: dealloc %{{.*}} : memref<256xf32, 2>45 // CHECK-NEXT: dealloc %{{.*}} : memref<1xi32>46 // CHECK-NEXT: dealloc %{{.*}} : memref<256xf32, 2>47 // CHECK-NEXT: return48 affine.for %i = 0 to 256 {49 affine.load %A[%i] : memref<256 x f32>50 %idx = affine.apply affine_map<(d0) -> (d0 + 256)>(%i)51 affine.load %B[%idx] : memref<512 x f32>52 affine.load %F[%i] : memref<256 x f32, 2>53 }54 return55}56 57// -----58 59// CHECK-LABEL: func @loop_nest_high_d60// CHECK: %{{.*}} = arith.constant 16384 : index61// CHECK-DAG: [[BUFB:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<512x32xf32, 2>62// CHECK-DAG: [[BUFA:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<512x32xf32, 2>63// CHECK-DAG: [[BUFC:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<512x32xf32, 2>64// CHECK-DAG: [[TAGB:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<1xi32>65// CHECK-DAG: [[TAGA:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<1xi32>66// CHECK-DAG: [[TAGC:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<1xi32>67// CHECK-DAG: [[TAGC_W:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<1xi32>68// INCOMING DMA for B69// CHECK-DAG: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], [[BUFB]][%{{.*}}, %{{.*}}], [[TAGB]][%{{.*}}], %{{.*}} : memref<512x32xf32>, memref<512x32xf32, 2>, memref<1xi32>70// CHECK-DAG: affine.dma_wait [[TAGB]][%{{.*}}], %{{.*}} : memref<1xi32>71// INCOMING DMA for A.72// CHECK-DAG: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], [[BUFA]][%{{.*}}, %{{.*}}], [[TAGA]][%{{.*}}], %{{.*}} : memref<512x32xf32>, memref<512x32xf32, 2>, memref<1xi32>73// CHECK-DAG: affine.dma_wait [[TAGA]][%{{.*}}], %{{.*}} : memref<1xi32>74// INCOMING DMA for C.75// CHECK-DAG: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], [[BUFC]][%{{.*}}, %{{.*}}], [[TAGC]][%{{.*}}], %{{.*}} : memref<512x32xf32>, memref<512x32xf32, 2>, memref<1xi32>76// CHECK-DAG: affine.dma_wait [[TAGC]][%{{.*}}], %{{.*}} : memref<1xi32>77// CHECK-NEXT: affine.for %{{.*}} = 0 to 32 {78// CHECK-NEXT: affine.for %{{.*}} = 0 to 32 {79// CHECK-NEXT: affine.for %{{.*}} = 0 to 32 {80// CHECK-NEXT: affine.for %{{.*}} = 0 to 16 {81// CHECK: affine.load [[BUFB]][%{{.*}} * 16 + %{{.*}}, %{{.*}}] : memref<512x32xf32, 2>82// CHECK-NEXT: "foo"(%{{.*}}) : (f32) -> ()83// CHECK-NEXT: }84// CHECK-NEXT: affine.for %{{.*}} = 0 to 16 {85// CHECK: affine.load [[BUFA]][%{{.*}} * 16 + %{{.*}}, %{{.*}}] : memref<512x32xf32, 2>86// CHECK-NEXT: "bar"(%{{.*}}) : (f32) -> ()87// CHECK-NEXT: }88// CHECK-NEXT: affine.for %{{.*}} = 0 to 16 {89// CHECK-NEXT: "abc_compute"() : () -> f3290// CHECK: affine.load [[BUFC]][%{{.*}} * 16 + %{{.*}}, %{{.*}}] : memref<512x32xf32, 2>91// CHECK-NEXT: "addf32"(%{{.*}}, %{{.*}}) : (f32, f32) -> f3292// CHECK-NEXT: affine.store %{{.*}}, [[BUFC]][%{{.*}} * 16 + %{{.*}}, %{{.*}}] : memref<512x32xf32, 2>93// CHECK-NEXT: }94// CHECK-NEXT: "foobar"() : () -> ()95// CHECK-NEXT: }96// CHECK-NEXT: }97// CHECK-NEXT: }98// OUTGOING DMA for C.99// CHECK-NEXT: affine.dma_start [[BUFC]][%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}, %{{.*}}], [[TAGC_W]][%{{.*}}], %{{.*}} : memref<512x32xf32, 2>, memref<512x32xf32>, memref<1xi32>100// CHECK-NEXT: affine.dma_wait [[TAGC_W]][%{{.*}}], %{{.*}} : memref<1xi32>101// CHECK-NEXT: dealloc [[TAGC_W]] : memref<1xi32>102// CHECK-NEXT: dealloc [[TAGC]] : memref<1xi32>103// CHECK-NEXT: dealloc [[BUFC]] : memref<512x32xf32, 2>104// CHECK-NEXT: dealloc [[TAGA]] : memref<1xi32>105// CHECK-NEXT: dealloc [[BUFA]] : memref<512x32xf32, 2>106// CHECK-NEXT: dealloc [[TAGB]] : memref<1xi32>107// CHECK-NEXT: dealloc [[BUFB]] : memref<512x32xf32, 2>108// CHECK-NEXT: return109// CHECK-NEXT:}110func.func @loop_nest_high_d(%A: memref<512 x 32 x f32>,111 %B: memref<512 x 32 x f32>, %C: memref<512 x 32 x f32>) {112 // DMAs will be performed at this level (jT is the first loop without a stride).113 // A and B are read, while C is both read and written. A total of three new buffers114 // are allocated and existing load's/store's are replaced by accesses to those buffers.115 affine.for %jT = 0 to 32 {116 affine.for %kT = 0 to 32 {117 affine.for %iT = 0 to 32 {118 affine.for %kk = 0 to 16 { // k intratile119 %k = affine.apply affine_map<(d0, d1) -> (16*d0 + d1)> (%kT, %kk)120 %v0 = affine.load %B[%k, %jT] : memref<512 x 32 x f32>121 "foo"(%v0) : (f32) -> ()122 }123 affine.for %ii = 0 to 16 { // i intratile.124 %i = affine.apply affine_map<(d0, d1) -> (16*d0 + d1)>(%iT, %ii)125 %v1 = affine.load %A[%i, %kT] : memref<512 x 32 x f32>126 "bar"(%v1) : (f32) -> ()127 }128 affine.for %ii_ = 0 to 16 { // i intratile.129 %v2 = "abc_compute"() : () -> f32130 %i_ = affine.apply affine_map<(d0, d1) -> (16*d0 + d1)>(%iT, %ii_)131 %v3 = affine.load %C[%i_, %jT] : memref<512 x 32 x f32>132 %v4 = "addf32"(%v2, %v3) : (f32, f32) -> (f32)133 affine.store %v4, %C[%i_, %jT] : memref<512 x 32 x f32>134 }135 "foobar"() : () -> ()136 }137 }138 }139 return140}141 142// -----143 144// A loop nest with a modulo 2 access. A strided DMA is not needed here a 1x2145// region within a 256 x 8 memref.146//147// CHECK-LABEL: func @loop_nest_modulo() {148// CHECK: memref.alloc() : memref<256x8xf32>149// CHECK-NEXT: affine.for %{{.*}} = 0 to 32 step 4 {150// CHECK: memref.alloc() : memref<1x2xf32, 2>151// CHECK-NEXT: memref.alloc() : memref<1xi32>152// Composition of the affine map for '%{{.*}}' causes '%{{.*}}' to be added as a symbol.153// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}, 0], %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<256x8xf32>, memref<1x2xf32, 2>, memref<1xi32>154// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>155// CHECK-NEXT: affine.for %{{.*}} = 0 to 8 {156// ...157// ...158// CHECK: }159// CHECK-NEXT: dealloc %{{.*}} : memref<1xi32>160// CHECK-NEXT: dealloc %{{.*}} : memref<1x2xf32, 2>161// CHECK-NEXT: }162// CHECK-NEXT: return163func.func @loop_nest_modulo() {164 %A = memref.alloc() : memref<256 x 8 x f32>165 affine.for %i = 0 to 32 step 4 {166 // DMAs will be performed at this level (%j is the first unit stride loop)167 affine.for %j = 0 to 8 {168 %idx = affine.apply affine_map<(d0) -> (d0 mod 2)> (%j)169 // A buffer of size 32 x 2 will be allocated (original buffer was 256 x 8).170 %v = affine.load %A[%i, %idx] : memref<256 x 8 x f32>171 }172 }173 return174}175 176// -----177 178// DMA on tiled loop nest. This also tests the case where the bounds are179// dependent on outer loop IVs.180// CHECK-LABEL: func @loop_nest_tiled() -> memref<256x1024xf32> {181func.func @loop_nest_tiled() -> memref<256x1024xf32> {182 %0 = memref.alloc() : memref<256x1024xf32>183 affine.for %i0 = 0 to 256 step 32 {184 affine.for %i1 = 0 to 1024 step 32 {185// CHECK: memref.alloc() : memref<32x32xf32, 2>186// CHECK-NEXT: memref.alloc() : memref<1xi32>187// Strided DMA here: 32 x 32 tile in a 256 x 1024 memref.188// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}, %{{.*}}, %{{.*}} : memref<256x1024xf32>, memref<32x32xf32, 2>, memref<1xi32>189// CHECK-NEXT: affine.dma_wait190// CHECK-NEXT: affine.for %{{.*}} = #map191// CHECK-NEXT: affine.for %{{.*}} = #map192 affine.for %i2 = affine_map<(d0) -> (d0)>(%i0) to affine_map<(d0) -> (d0 + 32)>(%i0) {193 affine.for %i3 = affine_map<(d0) -> (d0)>(%i1) to affine_map<(d0) -> (d0 + 32)>(%i1) {194 // CHECK: affine.load %{{.*}}[-%{{.*}} + %{{.*}}, -%{{.*}} + %{{.*}}] : memref<32x32xf32, 2>195 %1 = affine.load %0[%i2, %i3] : memref<256x1024xf32>196 } // CHECK-NEXT: }197 }198 }199 }200 return %0 : memref<256x1024xf32>201}202 203// -----204 205// CHECK-LABEL: func @dma_constant_dim_access206func.func @dma_constant_dim_access(%A : memref<100x100xf32>) {207 %one = arith.constant 1 : index208 %N = arith.constant 100 : index209 // CHECK: memref.alloc() : memref<1x100xf32, 2>210 // CHECK-NEXT: memref.alloc() : memref<1xi32>211 // No strided DMA needed here.212 // CHECK: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<100x100xf32>, memref<1x100xf32, 2>,213 // CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>214 affine.for %i = 0 to 100 {215 affine.for %j = 0 to affine_map<()[s0] -> (s0)> ()[%N] {216 // CHECK: affine.load %{{.*}}[0, %{{.*}}] : memref<1x100xf32, 2>217 affine.load %A[%one, %j] : memref<100 x 100 x f32>218 }219 }220 return221}222 223// -----224 225// CHECK-LABEL: func @dma_with_symbolic_accesses226func.func @dma_with_symbolic_accesses(%A : memref<100x100xf32>, %M : index) {227 %N = arith.constant 9 : index228 affine.for %i = 0 to 100 {229 affine.for %j = 0 to 100 {230 %idy = affine.apply affine_map<(d0, d1) [s0, s1] -> (d1 + s0 + s1)>(%i, %j)[%M, %N]231 affine.load %A[%i, %idy] : memref<100 x 100 x f32>232 }233 }234 return235// CHECK: memref.alloc() : memref<100x100xf32, 2>236// CHECK-NEXT: memref.alloc() : memref<1xi32>237// CHECK-NEXT: affine.dma_start %{{.*}}[0, symbol(%{{.*}}) + 9], %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}238// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}}239// CHECK-NEXT: affine.for %[[IV0:.*]] = 0 to 100 {240// CHECK-NEXT: affine.for %[[IV1:.*]] = 0 to 100 {241// CHECK: affine.load %{{.*}}[%[[IV0]], %[[IV1]]] : memref<100x100xf32, 2>242// CHECK-NEXT: }243// CHECK-NEXT: }244// CHECK: return245}246 247// -----248 249// CHECK-LABEL: func @dma_with_symbolic_loop_bounds250func.func @dma_with_symbolic_loop_bounds(%A : memref<100x100xf32>, %M : index, %N: index) {251 %K = arith.constant 9 : index252// The buffer size can't be bound by a constant smaller than the original253// memref size; so the DMA buffer is the entire 100x100.254// CHECK: memref.alloc() : memref<100x100xf32, 2>255// CHECK-NEXT: memref.alloc() : memref<1xi32>256// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>257// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>258 affine.for %i = 0 to 100 {259 affine.for %j = %M to %N {260 %idy = affine.apply affine_map<(d1) [s0] -> (d1 + s0)>(%j)[%K]261 affine.load %A[%i, %idy] : memref<100 x 100 x f32>262 }263 }264 return265}266 267// -----268 269// CHECK-LABEL: func @dma_unknown_size270func.func @dma_unknown_size(%arg0: memref<?x?xf32>) {271 %c0 = arith.constant 0 : index272 %M = memref.dim %arg0, %c0 : memref<? x ? x f32>273 %N = memref.dim %arg0, %c0 : memref<? x ? x f32>274 affine.for %i = 0 to %M {275 affine.for %j = 0 to %N {276 // If this loop nest isn't tiled, the access requires a non-constant DMA277 // size -- not yet implemented.278 // CHECK: affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>279 affine.load %arg0[%i, %j] : memref<? x ? x f32>280 }281 }282 return283}284 285// -----286 287// CHECK-LABEL: func @dma_memref_3d288func.func @dma_memref_3d(%arg0: memref<1024x1024x1024xf32>) {289 affine.for %i = 0 to 1024 {290 affine.for %j = 0 to 1024 {291 affine.for %k = 0 to 1024 {292 %idx = affine.apply affine_map<(d0) -> (d0 mod 128)>(%i)293 %idy = affine.apply affine_map<(d0) -> (d0 mod 128)>(%j)294 %idz = affine.apply affine_map<(d0) -> (d0 mod 128)>(%k)295 // DMA with nested striding (or emulating with loop around strided DMA)296 // not yet implemented.297 // CHECK: affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<1024x1024x1024xf32>298 %v = affine.load %arg0[%idx, %idy, %idz] : memref<1024 x 1024 x 1024 x f32>299 }300 }301 }302 return303}304 305// -----306 307// The first load accesses ([2,258), [128,384))308// The second load accesses ([64,320), [2,258))309// The first store writes to ([2,258), [192,448))310// The second store writes to ([128,320), [2,258))311// The union of all these regions is of size 318 x 446 and has its origin at (2,312// 2), i.e., the window ([2,320), [2,448)) in the original space.313 314// CHECK-LABEL: func @multi_load_store_union() {315func.func @multi_load_store_union() {316 %A = memref.alloc() : memref<512 x 512 x f32>317 affine.for %i = 0 to 256 {318 affine.for %j = 0 to 256 {319 %idx = affine.apply affine_map<(d0) -> (d0 + 64)>(%i)320 %idy = affine.apply affine_map<(d0) -> (d0 + 128)>(%j)321 %ishift = affine.apply affine_map<(d0) -> (d0 + 2)>(%i)322 %jshift = affine.apply affine_map<(d0) -> (d0 + 2)>(%j)323 324 %u = affine.load %A[%ishift, %idy] : memref<512 x 512 x f32>325 %v = affine.load %A[%idx, %jshift] : memref<512 x 512 x f32>326 327 %sidx = affine.apply affine_map<(d0) -> (d0 + 128)>(%i)328 %sidy = affine.apply affine_map<(d0) -> (d0 + 192)>(%j)329 330 affine.store %u, %A[%ishift, %sidy] : memref<512 x 512 x f32>331 affine.store %v, %A[%sidx, %jshift] : memref<512 x 512 x f32>332 }333 }334 return335}336// CHECK: memref.alloc() : memref<512x512xf32>337// CHECK-NEXT: memref.alloc() : memref<382x446xf32, 2>338// CHECK-NEXT: memref.alloc() : memref<1xi32>339// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}, %{{.*}}, %{{.*}} : memref<512x512xf32>, memref<382x446xf32, 2>, memref<1xi32>340// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>341// CHECK-NEXT: memref.alloc() : memref<1xi32>342// CHECK-NEXT: affine.for %{{.*}} = 0 to 256 {343// CHECK-NEXT: affine.for %{{.*}} = 0 to 256 {344// CHECK: affine.load %{{.*}}[%{{.*}}, %{{.*}} + 126] : memref<382x446xf32, 2>345// CHECK-NEXT: affine.load %{{.*}}[%{{.*}} + 62, %{{.*}}] : memref<382x446xf32, 2>346// CHECK: affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}} + 190] : memref<382x446xf32, 2>347// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}} + 126, %{{.*}}] : memref<382x446xf32, 2>348// CHECK-NEXT: }349// CHECK-NEXT: }350// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}, %{{.*}}, %{{.*}} : memref<382x446xf32, 2>, memref<512x512xf32>, memref<1xi32>351// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>352// CHECK-NEXT: dealloc %{{.*}} : memref<1xi32>353// CHECK-NEXT: dealloc %{{.*}} : memref<1xi32>354// CHECK-NEXT: dealloc %{{.*}} : memref<382x446xf32, 2>355// CHECK-NEXT: return356// CHECK-NEXT:}357 358// -----359 360// CHECK-LABEL: func @dma_loop_straightline_interspersed() {361func.func @dma_loop_straightline_interspersed() {362 %c0 = arith.constant 0 : index363 %c255 = arith.constant 255 : index364 %A = memref.alloc() : memref<256 x f32>365 %v = affine.load %A[%c0] : memref<256 x f32>366 affine.for %i = 1 to 255 {367 affine.load %A[%i] : memref<256 x f32>368 }369 %l = affine.load %A[%c255] : memref<256 x f32>370 affine.store %l, %A[%c0] : memref<256 x f32>371 return372}373// There are three regions here - the 'load' preceding the loop, the loop374// itself, and the operations appearing after the scf.375// CHECK: memref.alloc() : memref<256xf32>376// CHECK-NEXT: memref.alloc() : memref<1xf32, 2>377// CHECK-NEXT: memref.alloc() : memref<1xi32>378// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<256xf32>, memref<1xf32, 2>, memref<1xi32>379// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>380// CHECK-NEXT: affine.load %{{.*}}[0] : memref<1xf32, 2>381// CHECK-NEXT: dealloc %{{.*}} : memref<1xi32>382// CHECK-NEXT: dealloc %{{.*}} : memref<1xf32, 2>383// CHECK-NEXT: memref.alloc() : memref<254xf32, 2>384// CHECK-NEXT: memref.alloc() : memref<1xi32>385// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<256xf32>, memref<254xf32, 2>, memref<1xi32>386// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>387// CHECK-NEXT: affine.for %{{.*}} = 1 to 255 {388// CHECK-NEXT: affine.load %{{.*}}[%{{.*}} - 1] : memref<254xf32, 2>389// CHECK-NEXT: }390// CHECK-NEXT: dealloc %{{.*}} : memref<1xi32>391// CHECK-NEXT: dealloc %{{.*}} : memref<254xf32, 2>392// CHECK-NEXT: memref.alloc() : memref<256xf32, 2>393// CHECK-NEXT: memref.alloc() : memref<1xi32>394// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<256xf32>, memref<256xf32, 2>, memref<1xi32>395// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>396// CHECK-NEXT: memref.alloc() : memref<1xi32>397// CHECK-NEXT: affine.load %{{.*}}[255] : memref<256xf32, 2>398// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0] : memref<256xf32, 2>399// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<256xf32, 2>, memref<256xf32>, memref<1xi32>400// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>401// CHECK-NEXT: dealloc %{{.*}} : memref<1xi32>402// CHECK-NEXT: dealloc %{{.*}} : memref<1xi32>403// CHECK-NEXT: dealloc %{{.*}} : memref<256xf32, 2>404// CHECK-NEXT: return405 406// -----407 408// CHECK-LABEL: func @dma_mixed_loop_blocks() {409func.func @dma_mixed_loop_blocks() {410 %c0 = arith.constant 0 : index411 %A = memref.alloc() : memref<256 x 256 x vector<8 x f32>>412 affine.for %i = 0 to 256 {413 %v = affine.load %A[%c0, %c0] : memref<256 x 256 x vector<8 x f32>>414 "foo"(%v) : (vector<8 x f32>) -> ()415 affine.for %j = 0 to 256 {416 %w = affine.load %A[%i, %j] : memref<256 x 256 x vector<8 x f32>>417 "bar"(%w) : (vector<8 x f32>) -> ()418 }419 }420 return421}422// CHECK-DAG: [[MEM:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<256x256xvector<8xf32>>423// CHECK-DAG: [[BUF:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<256x256xvector<8xf32>, 2>424// CHECK-DAG: [[TAG:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<1xi32>425// CHECK: affine.dma_start [[MEM]][%{{.*}}, %{{.*}}], [[BUF]][%{{.*}}, %{{.*}}], [[TAG]][%{{.*}}], %{{.*}} : memref<256x256xvector<8xf32>>, memref<256x256xvector<8xf32>, 2>, memref<1xi32>426// CHECK-NEXT: affine.dma_wait [[TAG]][%{{.*}}], %{{.*}} : memref<1xi32>427// CHECK-NEXT: affine.for %{{.*}} = 0 to 256 {428// CHECK: affine.load [[BUF]][0, 0] : memref<256x256xvector<8xf32>, 2>429// CHECK: affine.for %{{.*}} = 0 to 256 {430// CHECK-NEXT: affine.load [[BUF]][%{{.*}}, %{{.*}}] : memref<256x256xvector<8xf32>, 2>431 432// -----433 434// CHECK-LABEL: func @relative_loop_bounds435func.func @relative_loop_bounds(%arg0: memref<1027xf32>) {436 affine.for %i0 = 0 to 1024 {437 affine.for %i2 = affine_map<(d0) -> (d0)>(%i0) to affine_map<(d0) -> (d0 + 4)>(%i0) {438 %0 = arith.constant 0.0 : f32439 affine.store %0, %arg0[%i2] : memref<1027xf32>440 }441 }442 return443}444// CHECK: [[BUF:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<1027xf32, 2>445// CHECK-NEXT: [[MEM:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<1xi32>446// CHECK-NEXT: affine.for %{{.*}} = 0 to 1024 {447// CHECK-NEXT: affine.for %[[I2:.*]] = {{#map[0-9a-zA-Z_]*}}(%{{.*}}) to {{#map[0-9a-zA-Z_]*}}(%{{.*}}) {448// CHECK: affine.store %{{.*}}, [[BUF]][%[[I2]]] : memref<1027xf32, 2>449// CHECK-NEXT: }450// CHECK-NEXT: }451// CHECK-NEXT: affine.dma_start [[BUF]][%{{.*}}], %{{.*}}[%{{.*}}], [[MEM]][%{{.*}}], %{{.*}} : memref<1027xf32, 2>, memref<1027xf32>, memref<1xi32>452// CHECK-NEXT: affine.dma_wait [[MEM]][%{{.*}}], %{{.*}} : memref<1xi32>453 454// -----455 456func.func @test_read_write_region_union() {457 %0 = memref.alloc() : memref<256xf32>458 affine.for %i0 = 0 to 10 {459 // memref dims: [0, 256)460 // read region: [100, 110)461 // write region: [25, 35)462 // union region: [25, 110)463 %a0 = affine.apply affine_map<(d0) -> (d0 + 100)>(%i0)464 %a1 = affine.apply affine_map<(d0) -> (d0 + 25)>(%i0)465 %1 = affine.load %0[%a0] : memref<256xf32>466 affine.store %1, %0[%a1] : memref<256xf32>467 }468 return469}470 471// CHECK: memref.alloc() : memref<256xf32>472// CHECK-NEXT: memref.alloc() : memref<85xf32, 2>473// CHECK-NEXT: memref.alloc() : memref<1xi32>474// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<256xf32>, memref<85xf32, 2>, memref<1xi32>475// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>476// CHECK-NEXT: memref.alloc() : memref<1xi32>477// CHECK-NEXT: affine.for %{{.*}} = 0 to 10 {478// CHECK: affine.load %{{.*}}[%{{.*}} + 75] : memref<85xf32, 2>479// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<85xf32, 2>480// CHECK-NEXT: }481// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<85xf32, 2>, memref<256xf32>, memref<1xi32>482// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>483 484// -----485 486// This should create a buffer of size 2 affine.for %arg2.487 488// CHECK-LABEL: func @test_analysis_util489func.func @test_analysis_util(%arg0: memref<4x4x16x1xf32>, %arg1: memref<144x9xf32>, %arg2: memref<2xf32>) -> (memref<144x9xf32>, memref<2xf32>) {490 %c0 = arith.constant 0 : index491 %0 = memref.alloc() : memref<64x1xf32>492 %1 = memref.alloc() : memref<144x4xf32>493 %2 = arith.constant 0.0 : f32494 affine.for %i8 = 0 to 9 step 3 {495 affine.for %i9 = affine_map<(d0) -> (d0)>(%i8) to affine_map<(d0) -> (d0 + 3)>(%i8) {496 affine.for %i17 = 0 to 64 {497 %25 = affine.load %arg2[%i9 floordiv 8] : memref<2xf32>498 %27 = affine.load %0[%i17, %c0] : memref<64x1xf32>499 affine.store %27, %arg2[%i17] : memref<2xf32>500 }501 }502 }503 return %arg1, %arg2 : memref<144x9xf32>, memref<2xf32>504}505// CHECK: affine.for %{{.*}} = 0 to 9 step 3 {506// CHECK: [[BUF:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<2xf32, 2>507// CHECK: affine.dma_start %{{.*}}[%c0{{.*}}], [[BUF]]508// CHECK: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>509// CHECK: affine.for %{{.*}} =510 511// -----512 513#map3 = affine_map<(d0) -> (d0)>514#map12 = affine_map<(d0) -> (d0 + 3)>515#map14 = affine_map<(d0, d1) -> ((d0 + d1 * 72) floordiv 2304 + ((((d0 + d1 * 72) mod 2304) mod 1152) mod 9) floordiv 3)>516#map15 = affine_map<(d0, d1) -> ((d0 + d1 * 72) mod 2304 - (((d0 + d1 * 72) mod 2304) floordiv 1152) * 1151 - ((((d0 + d1 * 72) mod 2304) mod 1152) floordiv 9) * 9 - (((((d0 + d1 * 72) mod 2304) mod 1152) mod 9) floordiv 3) * 3)>517#map16 = affine_map<(d0, d1) -> (((((d0 + d1 * 72) mod 2304) mod 1152) floordiv 9) floordiv 8)>518// Test for test case in b/128303048 #4.519// CHECK-LABEL: func @test_memref_bounds520func.func @test_memref_bounds(%arg0: memref<4x4x16x1xvector<8x128xf32>>, %arg1: memref<144x9xvector<8x128xf32>>, %arg2: memref<2xvector<8x128xf32>>) -> (memref<144x9xvector<8x128xf32>>, memref<2xvector<8x128xf32>>) {521 %c0 = arith.constant 0 : index522 affine.for %i8 = 0 to 9 step 3 {523 affine.for %i9 = #map3(%i8) to #map12(%i8) {524 affine.for %i10 = 0 to 64 {525 %10 = affine.apply #map14(%i9, %i10)526 %11 = affine.apply #map15(%i9, %i10)527 %12 = affine.apply #map16(%i9, %i10)528 %13 = affine.load %arg0[%10, %11, %12, %c0] : memref<4x4x16x1xvector<8x128xf32>>529 }530 }531 }532 return %arg1, %arg2 : memref<144x9xvector<8x128xf32>>, memref<2xvector<8x128xf32>>533}534 535// CHECK: memref.alloc() : memref<4x4x16x1xvector<8x128xf32>, 2>536// CHECK-NEXT: memref.alloc() : memref<1xi32>537// CHECK-NEXT: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<4x4x16x1xvector<8x128xf32>>, memref<4x4x16x1xvector<8x128xf32>, 2>, memref<1xi32>538// CHECK-NEXT: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>539 540// -----541 542// Since the fast memory size is 4 KB, DMA generation will happen right under543// %i0.544 545// FAST-MEM-16KB-LABEL: func @load_store_same_memref546func.func @load_store_same_memref(%arg0: memref<256x1024xf32>) {547 // FAST-MEM-16KB: affine.for %{{.*}} = 0 to 256 step 4548 affine.for %i0 = 0 to 256 step 4 {549 // FAST-MEM-16KB: [[BUF:%[0-9a-zA-Z_]+]] = memref.alloc() : memref<4x1024xf32, 2>550 // FAST-MEM-16KB: affine.dma_start %{{.*}}551 // FAST-MEM-16KB-NEXT: affine.dma_wait552 // FAST-MEM-16KB: affine.for %{{.*}}553 affine.for %i1 = 0 to 1024 step 4 {554 // FAST-MEM-16KB: affine.for %{{.*}}555 affine.for %i2 = affine_map<(d0) -> (d0)>(%i0) to affine_map<(d0) -> (d0 + 4)>(%i0) {556 // FAST-MEM-16KB: affine.for %{{.*}}557 affine.for %i3 = affine_map<(d0) -> (d0)>(%i1) to affine_map<(d0) -> (d0 + 4)>(%i1) {558 %3 = affine.load %arg0[%i2, %i3] : memref<256x1024xf32>559 %4 = arith.mulf %3, %3 : f32560 affine.store %4, %arg0[%i2, %i3] : memref<256x1024xf32>561 } // FAST-MEM-16KB: }562 } // FAST-MEM-16KB: }563 } // FAST-MEM-16KB: }564 // FAST-MEM-16KB: affine.dma_start [[BUF]]565 // FAST-MEM-16KB-NEXT: affine.dma_wait566 }567 return568}569 570// -----571 572// This a 3-d loop nest tiled by 4 x 4 x 4. Under %i, %j, %k, the size of a573// tile of arg0, arg1, and arg2 accessed is 4 KB (each), i.e., 12 KB in total.574// With fast mem capacity set to 16 KB, the DMAs if placed under %k will fit.575// However, the region of arg2 accessed is invariant w.r.t the %k loop unlike576// %arg0 and %arg1. So, its DMA can be hoisted one level up and placed under577// %j, while the DMAs for arg0 and arg1 appear right under the %k scf.578 579#map0 = affine_map<(d0) -> (d0)>580#map1 = affine_map<(d0) -> (d0 + 4)>581// FAST-MEM-16KB-LABEL: func @simple_matmul582func.func @simple_matmul(%arg0: memref<8x8xvector<64xf32>>, %arg1: memref<8x8xvector<64xf32>>, %arg2: memref<8x8xvector<64xf32>>) -> memref<8x8xvector<64xf32>> {583 affine.for %i = 0 to 8 step 4 {584 affine.for %j = 0 to 8 step 4 {585 affine.for %k = 0 to 8 step 4 {586 affine.for %ii = #map0(%i) to #map1(%i) {587 affine.for %jj = #map0(%j) to #map1(%j) {588 affine.for %kk = #map0(%k) to #map1(%k) {589 %5 = affine.load %arg0[%ii, %kk] : memref<8x8xvector<64xf32>>590 %6 = affine.load %arg1[%kk, %jj] : memref<8x8xvector<64xf32>>591 %7 = affine.load %arg2[%ii, %jj] : memref<8x8xvector<64xf32>>592 %8 = arith.mulf %5, %6 : vector<64xf32>593 %9 = arith.addf %7, %8 : vector<64xf32>594 affine.store %9, %arg2[%ii, %jj] : memref<8x8xvector<64xf32>>595 }596 }597 }598 }599 }600 }601 return %arg2 : memref<8x8xvector<64xf32>>602}603// FAST-MEM-16KB: affine.for %{{.*}} = 0 to 8 step 4 {604// FAST-MEM-16KB: affine.for %{{.*}} = 0 to 8 step 4 {605// FAST-MEM-16KB: affine.dma_start %{{.*}}606// FAST-MEM-16KB: affine.dma_wait607// FAST-MEM-16KB: affine.for %{{.*}} = 0 to 8 step 4 {608// FAST-MEM-16KB: affine.dma_start %{{.*}}609// FAST-MEM-16KB: affine.dma_wait610// FAST-MEM-16KB: affine.dma_start %{{.*}}611// FAST-MEM-16KB: affine.dma_wait612// FAST-MEM-16KB: affine.for %{{.*}} = #map{{[0-9a-zA-Z_]*}}(%{{.*}}) to #map{{[0-9a-zA-Z_]*}}(%{{.*}}) {613// FAST-MEM-16KB-NEXT: affine.for %{{.*}} = #map{{[0-9a-zA-Z_]*}}(%{{.*}}) to #map{{[0-9a-zA-Z_]*}}(%{{.*}}) {614// FAST-MEM-16KB-NEXT: affine.for %{{.*}} = #map{{[0-9a-zA-Z_]*}}(%{{.*}}) to #map{{[0-9a-zA-Z_]*}}(%{{.*}}) {615// FAST-MEM-16KB: }616// FAST-MEM-16KB: }617// FAST-MEM-16KB: }618// FAST-MEM-16KB: }619// FAST-MEM-16KB: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}620// FAST-MEM-16KB: affine.dma_wait621