1299 lines · plain
1// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion))' -split-input-file | FileCheck %s2// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{maximal}))' -split-input-file | FileCheck %s --check-prefix=MAXIMAL3 4// Part I of fusion tests in mlir/test/Transforms/loop-fusion.mlir.5// Part II of fusion tests in mlir/test/Transforms/loop-fusion-2.mlir6// Part IV of fusion tests in mlir/test/Transforms/loop-fusion-4.mlir7 8// -----9 10// Test case from github bug 777.11// CHECK-LABEL: func @mul_add_012func.func @mul_add_0(%arg0: memref<3x4xf32>, %arg1: memref<4x3xf32>, %arg2: memref<3x3xf32>, %arg3: memref<3x3xf32>) {13 %cst = arith.constant 0.000000e+00 : f3214 %0 = memref.alloc() : memref<3x3xf32>15 affine.for %arg4 = 0 to 3 {16 affine.for %arg5 = 0 to 3 {17 affine.store %cst, %0[%arg4, %arg5] : memref<3x3xf32>18 }19 }20 affine.for %arg4 = 0 to 3 {21 affine.for %arg5 = 0 to 3 {22 affine.for %arg6 = 0 to 4 {23 %1 = affine.load %arg1[%arg6, %arg5] : memref<4x3xf32>24 %2 = affine.load %arg0[%arg4, %arg6] : memref<3x4xf32>25 %3 = arith.mulf %2, %1 : f3226 %4 = affine.load %0[%arg4, %arg5] : memref<3x3xf32>27 %5 = arith.addf %4, %3 : f3228 affine.store %5, %0[%arg4, %arg5] : memref<3x3xf32>29 }30 }31 }32 affine.for %arg4 = 0 to 3 {33 affine.for %arg5 = 0 to 3 {34 %6 = affine.load %arg2[%arg4, %arg5] : memref<3x3xf32>35 %7 = affine.load %0[%arg4, %arg5] : memref<3x3xf32>36 %8 = arith.addf %7, %6 : f3237 affine.store %8, %arg3[%arg4, %arg5] : memref<3x3xf32>38 }39 }40 // CHECK: affine.for %[[i0:.*]] = 0 to 3 {41 // CHECK-NEXT: affine.for %[[i1:.*]] = 0 to 3 {42 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>43 // CHECK-NEXT: affine.for %[[i2:.*]] = 0 to 4 {44 // CHECK-NEXT: affine.load %{{.*}}[%[[i2]], %[[i1]]] : memref<4x3xf32>45 // CHECK-NEXT: affine.load %{{.*}}[%[[i0]], %[[i2]]] : memref<3x4xf32>46 // CHECK-NEXT: arith.mulf %{{.*}}, %{{.*}} : f3247 // CHECK-NEXT: affine.load %{{.*}}[0, 0] : memref<1x1xf32>48 // CHECK-NEXT: arith.addf %{{.*}}, %{{.*}} : f3249 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>50 // CHECK-NEXT: }51 // CHECK-NEXT: affine.load %{{.*}}[%[[i0]], %[[i1]]] : memref<3x3xf32>52 // CHECK-NEXT: affine.load %{{.*}}[0, 0] : memref<1x1xf32>53 // CHECK-NEXT: arith.addf %{{.*}}, %{{.*}} : f3254 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%[[i0]], %[[i1]]] : memref<3x3xf32>55 // CHECK-NEXT: }56 // CHECK-NEXT: }57 // CHECK-NEXT: return58 return59}60 61// -----62 63// Verify that 'fuseProducerConsumerNodes' fuse a producer loop with a store64// that has multiple outgoing edges.65 66// CHECK-LABEL: func @should_fuse_multi_outgoing_edge_store_producer67func.func @should_fuse_multi_outgoing_edge_store_producer(%a : memref<1xf32>) {68 %cst = arith.constant 0.000000e+00 : f3269 affine.for %arg0 = 0 to 1 {70 affine.store %cst, %a[%arg0] : memref<1xf32>71 }72 73 affine.for %arg0 = 0 to 1 {74 %0 = affine.load %a[%arg0] : memref<1xf32>75 }76 77 affine.for %arg0 = 0 to 1 {78 %0 = affine.load %a[%arg0] : memref<1xf32>79 }80 // CHECK: affine.for %{{.*}} = 0 to 1 {81 // CHECK-NEXT: affine.store82 // CHECK-NEXT: affine.load83 // CHECK-NEXT: affine.load84 // CHECK-NEXT: }85 86 return87}88 89// -----90 91// Verify that 'fuseProducerConsumerNodes' fuses a producer loop that: 1) has92// multiple outgoing edges, 2) producer store has a single outgoing edge.93// Sibling loop fusion should not fuse any of these loops due to94// dependencies on external memrefs '%a' and '%b'.95 96// CHECK-LABEL: func @should_fuse_producer_with_multi_outgoing_edges97func.func @should_fuse_producer_with_multi_outgoing_edges(%a : memref<1xf32>, %b : memref<1xf32>) {98 %cst = arith.constant 0.000000e+00 : f3299 affine.for %arg0 = 0 to 1 {100 %0 = affine.load %a[%arg0] : memref<1xf32>101 affine.store %cst, %b[%arg0] : memref<1xf32>102 }103 104 affine.for %arg0 = 0 to 1 {105 affine.store %cst, %a[%arg0] : memref<1xf32>106 %1 = affine.load %b[%arg0] : memref<1xf32>107 }108 // CHECK: affine.for %{{.*}} = 0 to 1109 // CHECK-NEXT: affine.load %[[A:.*]][{{.*}}]110 // CHECK-NEXT: affine.store %{{.*}}, %[[B:.*]][{{.*}}]111 // CHECK-NEXT: affine.store %{{.*}}, %[[A]]112 // CHECK-NEXT: affine.load %[[B]]113 // CHECK-NOT: affine.for %{{.*}}114 // CHECK: return115 return116}117 118// MAXIMAL-LABEL: func @reshape_into_matmul119func.func @reshape_into_matmul(%lhs : memref<1024x1024xf32>,120 %R: memref<16x64x1024xf32>, %out: memref<1024x1024xf32>) {121 %rhs = memref.alloc() : memref<1024x1024xf32>122 123 // Reshape from 3-d to 2-d.124 affine.for %i0 = 0 to 16 {125 affine.for %i1 = 0 to 64 {126 affine.for %k = 0 to 1024 {127 %v = affine.load %R[%i0, %i1, %k] : memref<16x64x1024xf32>128 affine.store %v, %rhs[64*%i0 + %i1, %k] : memref<1024x1024xf32>129 }130 }131 }132 133 // Matmul.134 affine.for %i = 0 to 1024 {135 affine.for %j = 0 to 1024 {136 affine.for %k = 0 to 1024 {137 %0 = affine.load %rhs[%k, %j] : memref<1024x1024xf32>138 %1 = affine.load %lhs[%i, %k] : memref<1024x1024xf32>139 %2 = arith.mulf %1, %0 : f32140 %3 = affine.load %out[%i, %j] : memref<1024x1024xf32>141 %4 = arith.addf %3, %2 : f32142 affine.store %4, %out[%i, %j] : memref<1024x1024xf32>143 }144 }145 }146 return147}148// MAXIMAL-NEXT: memref.alloc149// MAXIMAL-NEXT: affine.for150// MAXIMAL-NEXT: affine.for151// MAXIMAL-NEXT: affine.for152// MAXIMAL-NOT: affine.for153// MAXIMAL: return154 155// -----156 157// CHECK-LABEL: func @vector_loop158func.func @vector_loop(%a : memref<10x20xf32>, %b : memref<10x20xf32>,159 %c : memref<10x20xf32>) {160 affine.for %j = 0 to 10 {161 affine.for %i = 0 to 5 {162 %ld0 = affine.vector_load %a[%j, %i*4] : memref<10x20xf32>, vector<4xf32>163 affine.vector_store %ld0, %b[%j, %i*4] : memref<10x20xf32>, vector<4xf32>164 }165 }166 167 affine.for %j = 0 to 10 {168 affine.for %i = 0 to 5 {169 %ld0 = affine.vector_load %b[%j, %i*4] : memref<10x20xf32>, vector<4xf32>170 affine.vector_store %ld0, %c[%j, %i*4] : memref<10x20xf32>, vector<4xf32>171 }172 }173 174 return175}176// CHECK: affine.for177// CHECK-NEXT: affine.for178// CHECK-NEXT: affine.vector_load179// CHECK-NEXT: affine.vector_store180// CHECK-NEXT: affine.vector_load181// CHECK-NEXT: affine.vector_store182// CHECK-NOT: affine.for183 184// -----185 186// CHECK-LABEL: func @multi_outgoing_edges187func.func @multi_outgoing_edges(%in0 : memref<32xf32>,188 %in1 : memref<32xf32>) {189 affine.for %d = 0 to 32 {190 %lhs = affine.load %in0[%d] : memref<32xf32>191 %rhs = affine.load %in1[%d] : memref<32xf32>192 %add = arith.addf %lhs, %rhs : f32193 affine.store %add, %in0[%d] : memref<32xf32>194 }195 affine.for %d = 0 to 32 {196 %lhs = affine.load %in0[%d] : memref<32xf32>197 %rhs = affine.load %in1[%d] : memref<32xf32>198 %add = arith.subf %lhs, %rhs : f32199 affine.store %add, %in0[%d] : memref<32xf32>200 }201 affine.for %d = 0 to 32 {202 %lhs = affine.load %in0[%d] : memref<32xf32>203 %rhs = affine.load %in1[%d] : memref<32xf32>204 %add = arith.mulf %lhs, %rhs : f32205 affine.store %add, %in0[%d] : memref<32xf32>206 }207 affine.for %d = 0 to 32 {208 %lhs = affine.load %in0[%d] : memref<32xf32>209 %rhs = affine.load %in1[%d] : memref<32xf32>210 %add = arith.divf %lhs, %rhs : f32211 affine.store %add, %in0[%d] : memref<32xf32>212 }213 return214}215 216// CHECK: affine.for217// CHECK-NOT: affine.for218// CHECK: arith.addf219// CHECK-NOT: affine.for220// CHECK: arith.subf221// CHECK-NOT: affine.for222// CHECK: arith.mulf223// CHECK-NOT: affine.for224// CHECK: arith.divf225 226// -----227 228// Test fusion when dynamically shaped memrefs are used with constant trip count loops.229 230// CHECK-LABEL: func @calc231func.func @calc(%arg0: memref<?xf32>, %arg1: memref<?xf32>, %arg2: memref<?xf32>, %len: index) {232 %c1 = arith.constant 1 : index233 %1 = memref.alloc(%len) : memref<?xf32>234 affine.for %arg4 = 1 to 10 {235 %7 = affine.load %arg0[%arg4] : memref<?xf32>236 %8 = affine.load %arg1[%arg4] : memref<?xf32>237 %9 = arith.addf %7, %8 : f32238 affine.store %9, %1[%arg4] : memref<?xf32>239 }240 affine.for %arg4 = 1 to 10 {241 %7 = affine.load %1[%arg4] : memref<?xf32>242 %8 = affine.load %arg1[%arg4] : memref<?xf32>243 %9 = arith.mulf %7, %8 : f32244 affine.store %9, %arg2[%arg4] : memref<?xf32>245 }246 return247}248// CHECK: memref.alloc() : memref<1xf32>249// CHECK: affine.for %arg{{.*}} = 1 to 10 {250// CHECK-NEXT: affine.load %arg{{.*}}251// CHECK-NEXT: affine.load %arg{{.*}}252// CHECK-NEXT: arith.addf253// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>254// CHECK-NEXT: affine.load %{{.*}}[0] : memref<1xf32>255// CHECK-NEXT: affine.load %arg{{.*}}[%arg{{.*}}] : memref<?xf32>256// CHECK-NEXT: arith.mulf257// CHECK-NEXT: affine.store %{{.*}}, %arg{{.*}}[%arg{{.*}}] : memref<?xf32>258// CHECK-NEXT: }259// CHECK-NEXT: return260 261// -----262 263// CHECK-LABEL: func @should_not_fuse_since_non_affine_users264func.func @should_not_fuse_since_non_affine_users(%in0 : memref<32xf32>,265 %in1 : memref<32xf32>) {266 affine.for %d = 0 to 32 {267 %lhs = affine.load %in0[%d] : memref<32xf32>268 %rhs = affine.load %in1[%d] : memref<32xf32>269 %add = arith.addf %lhs, %rhs : f32270 affine.store %add, %in0[%d] : memref<32xf32>271 }272 affine.for %d = 0 to 32 {273 %lhs = memref.load %in0[%d] : memref<32xf32>274 %rhs = memref.load %in1[%d] : memref<32xf32>275 %add = arith.subf %lhs, %rhs : f32276 memref.store %add, %in0[%d] : memref<32xf32>277 }278 affine.for %d = 0 to 32 {279 %lhs = affine.load %in0[%d] : memref<32xf32>280 %rhs = affine.load %in1[%d] : memref<32xf32>281 %add = arith.mulf %lhs, %rhs : f32282 affine.store %add, %in0[%d] : memref<32xf32>283 }284 return285}286 287// CHECK: affine.for288// CHECK: arith.addf289// CHECK: affine.for290// CHECK: arith.subf291// CHECK: affine.for292// CHECK: arith.mulf293 294// -----295 296// CHECK-LABEL: func @should_not_fuse_since_top_level_non_affine_users297func.func @should_not_fuse_since_top_level_non_affine_users(%in0 : memref<32xf32>,298 %in1 : memref<32xf32>) {299 %sum = memref.alloc() : memref<f32>300 affine.for %d = 0 to 32 {301 %lhs = affine.load %in0[%d] : memref<32xf32>302 %rhs = affine.load %in1[%d] : memref<32xf32>303 %add = arith.addf %lhs, %rhs : f32304 memref.store %add, %sum[] : memref<f32>305 affine.store %add, %in0[%d] : memref<32xf32>306 }307 %load_sum = memref.load %sum[] : memref<f32>308 affine.for %d = 0 to 32 {309 %lhs = affine.load %in0[%d] : memref<32xf32>310 %rhs = affine.load %in1[%d] : memref<32xf32>311 %add = arith.mulf %lhs, %rhs : f32312 %sub = arith.subf %add, %load_sum: f32313 affine.store %sub, %in0[%d] : memref<32xf32>314 }315 memref.dealloc %sum : memref<f32>316 return317}318 319// CHECK: affine.for320// CHECK: arith.addf321// CHECK: affine.for322// CHECK: arith.mulf323// CHECK: arith.subf324 325// -----326 327// CHECK-LABEL: func @should_not_fuse_since_top_level_non_affine_mem_write_users328func.func @should_not_fuse_since_top_level_non_affine_mem_write_users(329 %in0 : memref<32xf32>, %in1 : memref<32xf32>) {330 %c0 = arith.constant 0 : index331 %cst_0 = arith.constant 0.000000e+00 : f32332 333 affine.for %d = 0 to 32 {334 %lhs = affine.load %in0[%d] : memref<32xf32>335 %rhs = affine.load %in1[%d] : memref<32xf32>336 %add = arith.addf %lhs, %rhs : f32337 affine.store %add, %in0[%d] : memref<32xf32>338 }339 memref.store %cst_0, %in0[%c0] : memref<32xf32>340 affine.for %d = 0 to 32 {341 %lhs = affine.load %in0[%d] : memref<32xf32>342 %rhs = affine.load %in1[%d] : memref<32xf32>343 %add = arith.addf %lhs, %rhs: f32344 affine.store %add, %in0[%d] : memref<32xf32>345 }346 return347}348 349// CHECK: affine.for350// CHECK: arith.addf351// CHECK: affine.for352// CHECK: arith.addf353 354// Tests that fusion isn't prevented by the presence of a dealloc op in355// between since we can move the fused nest.356 357// CHECK-LABEL: func @fuse_non_affine_intervening_op358func.func @fuse_non_affine_intervening_op() {359 %cst = arith.constant 0.0 : f32360 361 %a = memref.alloc() : memref<100xf32>362 %b = memref.alloc() : memref<100xf32>363 %c = memref.alloc() : memref<100xf32>364 365 affine.for %i = 0 to 100 {366 affine.store %cst, %a[%i] : memref<100xf32>367 affine.store %cst, %c[%i] : memref<100xf32>368 }369 370 // The source is fused into the destination while being moved here.371 // CHECK: affine.for %{{.*}} = 0 to 100372 // CHECK-NEXT: affine.store %cst373 // CHECK-NEXT: affine.store %cst{{.*}}374 // CHECK-NEXT: affine.load375 // CHECK-NEXT: affine.store376 // CHECK-NEXT: }377 // CHECK-NEXT: memref.dealloc378 379 memref.dealloc %c : memref<100xf32>380 381 affine.for %i = 0 to 100 {382 %v = affine.load %a[%i] : memref<100xf32>383 affine.store %v, %b[%i] : memref<100xf32>384 }385 386 return387}388 389// Tests that fusion happens in the presence of intervening non-affine reads.390 391// CHECK-LABEL: func @fuse_non_affine_intervening_read392func.func @fuse_non_affine_intervening_read() {393 %cst = arith.constant 0.0 : f32394 395 %a = memref.alloc() : memref<100xf32>396 %b = memref.alloc() : memref<100xf32>397 %c = memref.alloc() : memref<100xf32>398 399 affine.for %i = 0 to 100 {400 affine.store %cst, %a[%i] : memref<100xf32>401 }402 403 // The source is fused into the destination while being moved here.404 // CHECK: affine.for %{{.*}} = 0 to 100405 // CHECK-NEXT: affine.store %cst406 // CHECK-NEXT: affine.load407 // CHECK-NEXT: affine.store408 // CHECK-NEXT: }409 410 // CHECK: affine.for %{{.*}} = 0 to 100411 // CHECK-NEXT: memref.load412 affine.for %i = 0 to 100 {413 memref.load %a[%i] : memref<100xf32>414 }415 416 affine.for %i = 0 to 100 {417 %v = affine.load %a[%i] : memref<100xf32>418 affine.store %v, %b[%i] : memref<100xf32>419 }420 421 return422}423 424// Tests that fusion happens in the presence of intervening non-affine region425// ops.426 427// CHECK-LABEL: func @fuse_non_affine_intervening_read_nest428func.func @fuse_non_affine_intervening_read_nest() {429 %cst = arith.constant 0.0 : f32430 %c0 = arith.constant 0 : index431 %c1 = arith.constant 1 : index432 %c100 = arith.constant 100 : index433 434 %a = memref.alloc() : memref<100xf32>435 %b = memref.alloc() : memref<100xf32>436 %c = memref.alloc() : memref<100xf32>437 438 affine.for %i = 0 to 100 {439 affine.store %cst, %a[%i] : memref<100xf32>440 }441 442 // The source is fused into the destination while being moved here.443 // CHECK: affine.for %{{.*}} = 0 to 100444 // CHECK-NEXT: affine.store %cst445 // CHECK-NEXT: affine.load446 // CHECK-NEXT: affine.store447 // CHECK-NEXT: }448 449 // CHECK: scf.for450 // CHECK-NEXT: memref.load451 scf.for %i = %c0 to %c100 step %c1 {452 memref.load %a[%i] : memref<100xf32>453 }454 455 affine.for %i = 0 to 100 {456 %v = affine.load %a[%i] : memref<100xf32>457 affine.store %v, %b[%i] : memref<100xf32>458 }459 460 return461}462 463// Tests that fusion does not happen when there are non-affine sources464// intervening.465 466// CHECK-LABEL: func @no_fusion_scf_for_store467func.func @no_fusion_scf_for_store() {468 %cst = arith.constant 0.0 : f32469 %cst1 = arith.constant 1.0 : f32470 %c0 = arith.constant 0 : index471 %c1 = arith.constant 1 : index472 %c100 = arith.constant 100 : index473 474 %a = memref.alloc() : memref<100xf32>475 %b = memref.alloc() : memref<100xf32>476 %c = memref.alloc() : memref<100xf32>477 478 // CHECK: affine.for %{{.*}} = 0 to 100479 // CHECK-NEXT: affine.store480 // CHECK-NEXT: }481 affine.for %i = 0 to 100 {482 affine.store %cst, %a[%i] : memref<100xf32>483 }484 485 // CHECK: scf.for486 scf.for %i = %c0 to %c100 step %c1 {487 memref.store %cst1, %a[%i] : memref<100xf32>488 }489 490 // CHECK: affine.for %{{.*}} = 0 to 100491 // CHECK-NEXT: affine.load492 // CHECK-NEXT: affine.store493 affine.for %i = 0 to 100 {494 // Non-affine source for this load.495 %v = affine.load %a[%i] : memref<100xf32>496 affine.store %v, %b[%i] : memref<100xf32>497 }498 499 return500}501 502// -----503 504// CHECK-LABEL: func @fuse_minor_affine_map505// MAXIMAL-LABEL: func @fuse_minor_affine_map506func.func @fuse_minor_affine_map(%in: memref<128xf32>, %out: memref<20x512xf32>) {507 %tmp = memref.alloc() : memref<128xf32>508 509 affine.for %arg4 = 0 to 128 {510 %ld = affine.load %in[%arg4] : memref<128xf32>511 affine.store %ld, %tmp[%arg4] : memref<128xf32>512 }513 514 affine.for %arg3 = 0 to 20 {515 affine.for %arg4 = 0 to 512 {516 %ld = affine.load %tmp[%arg4 mod 128] : memref<128xf32>517 affine.store %ld, %out[%arg3, %arg4] : memref<20x512xf32>518 }519 }520 521 return522}523 524// MAXIMAL: memref.alloc() : memref<1xf32>525// MAXIMAL: affine.for526// MAXIMAL-NEXT: affine.for527// MAXIMAL-NOT: affine.for528// MAXIMAL: return529 530// -----531 532// CHECK-LABEL: func @should_fuse_multi_store_producer_and_privatize_memfefs533func.func @should_fuse_multi_store_producer_and_privatize_memfefs() {534 %a = memref.alloc() : memref<10xf32>535 %b = memref.alloc() : memref<10xf32>536 %c = memref.alloc() : memref<10xf32>537 %cst = arith.constant 0.000000e+00 : f32538 affine.for %arg0 = 0 to 10 {539 affine.store %cst, %a[%arg0] : memref<10xf32>540 affine.store %cst, %b[%arg0] : memref<10xf32>541 affine.store %cst, %c[%arg0] : memref<10xf32>542 %0 = affine.load %c[%arg0] : memref<10xf32>543 }544 545 affine.for %arg0 = 0 to 10 {546 %0 = affine.load %a[%arg0] : memref<10xf32>547 }548 549 affine.for %arg0 = 0 to 10 {550 %0 = affine.load %b[%arg0] : memref<10xf32>551 }552 553 // All the memrefs should be privatized except '%c', which is not involved in554 // the producer-consumer fusion.555 // CHECK: affine.for %{{.*}} = 0 to 10 {556 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>557 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>558 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>559 // CHECK-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<10xf32>560 // CHECK-NEXT: affine.load %{{.*}}[0] : memref<1xf32>561 // CHECK-NEXT: affine.load %{{.*}}[0] : memref<1xf32>562 // CHECK-NEXT: }563 564 return565}566 567// CHECK-LABEL: func @should_fuse_multi_store_producer_with_escaping_memrefs_and_remove_src568func.func @should_fuse_multi_store_producer_with_escaping_memrefs_and_remove_src(569 %a : memref<10xf32>, %b : memref<10xf32>) {570 %cst = arith.constant 0.000000e+00 : f32571 affine.for %i0 = 0 to 10 {572 affine.store %cst, %a[%i0] : memref<10xf32>573 affine.store %cst, %b[%i0] : memref<10xf32>574 }575 576 affine.for %i1 = 0 to 10 {577 %0 = affine.load %a[%i1] : memref<10xf32>578 }579 580 affine.for %i2 = 0 to 10 {581 %0 = affine.load %b[%i2] : memref<10xf32>582 }583 584 // Producer loop '%i0' should be removed after fusion since fusion is maximal.585 // No memref should be privatized since they escape the function, and the586 // producer is removed after fusion.587 // CHECK: affine.for %{{.*}} = 0 to 10 {588 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>589 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>590 // CHECK-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<10xf32>591 // CHECK-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<10xf32>592 // CHECK-NEXT: }593 // CHECK-NOT: affine.for594 595 return596}597 598// -----599 600func.func @should_fuse_multi_store_producer_with_escaping_memrefs_and_preserve_src(601 %a : memref<10xf32>, %b : memref<10xf32>) {602 %cst = arith.constant 0.000000e+00 : f32603 affine.for %i0 = 0 to 10 {604 affine.store %cst, %a[%i0] : memref<10xf32>605 affine.store %cst, %b[%i0] : memref<10xf32>606 }607 608 affine.for %i1 = 0 to 5 {609 %0 = affine.load %a[%i1] : memref<10xf32>610 }611 612 affine.for %i2 = 0 to 10 {613 %0 = affine.load %b[%i2] : memref<10xf32>614 }615 616 // Loops '%i0' and '%i2' should be fused first and '%i0' should be removed617 // since fusion is maximal. Then the fused loop and '%i1' should be fused618 // and the fused loop shouldn't be removed since fusion is not maximal.619 // CHECK: affine.for %{{.*}} = 0 to 10 {620 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>621 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>622 // CHECK-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<10xf32>623 // CHECK-NEXT: }624 // CHECK: affine.for %{{.*}} = 0 to 5 {625 // CHECK-NEXT: affine.store %{{.*}} : memref<1xf32>626 // CHECK-NEXT: affine.store %{{.*}} : memref<10xf32>627 // CHECK-NEXT: affine.load %{{.*}} : memref<10xf32>628 // CHECK-NEXT: affine.load %{{.*}} : memref<1xf32>629 // CHECK-NEXT: }630 // CHECK-NOT: affine.for631 632 return633}634 635 636func.func @should_not_fuse_due_to_dealloc(%arg0: memref<16xf32>){637 %A = memref.alloc() : memref<16xf32>638 %C = memref.alloc() : memref<16xf32>639 %cst_1 = arith.constant 1.000000e+00 : f32640 affine.for %arg1 = 0 to 16 {641 %a = affine.load %arg0[%arg1] : memref<16xf32>642 affine.store %a, %A[%arg1] : memref<16xf32>643 affine.store %a, %C[%arg1] : memref<16xf32>644 }645 memref.dealloc %C : memref<16xf32>646 %B = memref.alloc() : memref<16xf32>647 affine.for %arg1 = 0 to 16 {648 %a = affine.load %A[%arg1] : memref<16xf32>649 %b = arith.addf %cst_1, %a : f32650 affine.store %b, %B[%arg1] : memref<16xf32>651 }652 memref.dealloc %A : memref<16xf32>653 return654}655// CHECK-LABEL: func @should_not_fuse_due_to_dealloc656// CHECK: affine.for657// CHECK-NEXT: affine.load658// CHECK-NEXT: affine.store659// CHECK-NEXT: affine.store660// CHECK: memref.dealloc661// CHECK: affine.for662// CHECK-NEXT: affine.load663// CHECK-NEXT: arith.addf664// CHECK-NEXT: affine.store665 666// CHECK-LABEL: func @cannot_fuse_intervening_deallocs667func.func @cannot_fuse_intervening_deallocs(%arg0: memref<16xf32>){668 %A = memref.alloc() : memref<16xf32>669 %C = memref.alloc() : memref<16xf32>670 %cst_1 = arith.constant 1.000000e+00 : f32671 // CHECK: affine.for %{{.*}} = 0 to 16672 affine.for %arg1 = 0 to 16 {673 %a = affine.load %arg0[%arg1] : memref<16xf32>674 affine.store %a, %A[%arg1] : memref<16xf32>675 affine.store %a, %C[%arg1] : memref<16xf32>676 }677 // The presence of B's alloc prevents placement of the fused nest above C's678 // dealloc. No fusion here.679 memref.dealloc %C : memref<16xf32>680 %B = memref.alloc() : memref<16xf32>681 // CHECK: affine.for %{{.*}} = 0 to 16682 affine.for %arg1 = 0 to 16 {683 %a = affine.load %A[%arg1] : memref<16xf32>684 %b = arith.addf %cst_1, %a : f32685 affine.store %b, %B[%arg1] : memref<16xf32>686 }687 memref.dealloc %A : memref<16xf32>688 return689}690 691// -----692 693// CHECK-LABEL: func @should_fuse_defining_node_has_no_dependence_from_source_node694func.func @should_fuse_defining_node_has_no_dependence_from_source_node(695 %a : memref<10xf32>, %b : memref<f32>) -> () {696 affine.for %i0 = 0 to 10 {697 %0 = affine.load %b[] : memref<f32>698 affine.store %0, %a[%i0] : memref<10xf32>699 }700 %0 = affine.load %b[] : memref<f32>701 affine.for %i1 = 0 to 10 {702 %1 = affine.load %a[%i1] : memref<10xf32>703 %2 = arith.divf %0, %1 : f32704 }705 706 // Loops '%i0' and '%i1' should be fused even though there is a defining node707 // between the loops. It is because the node has no dependence from '%i0'.708 // CHECK: affine.load %{{.*}}[] : memref<f32>709 // CHECK-NEXT: affine.for %{{.*}} = 0 to 10 {710 // CHECK-NEXT: affine.load %{{.*}}[] : memref<f32>711 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>712 // CHECK-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<10xf32>713 // CHECK-NEXT: arith.divf714 // CHECK-NEXT: }715 // CHECK-NOT: affine.for716 return717}718 719// -----720 721// CHECK-LABEL: func @should_not_fuse_defining_node_has_dependence_from_source_loop722func.func @should_not_fuse_defining_node_has_dependence_from_source_loop(723 %a : memref<10xf32>, %b : memref<f32>) -> () {724 %cst = arith.constant 0.000000e+00 : f32725 affine.for %i0 = 0 to 10 {726 affine.store %cst, %b[] : memref<f32>727 affine.store %cst, %a[%i0] : memref<10xf32>728 }729 %0 = affine.load %b[] : memref<f32>730 affine.for %i1 = 0 to 10 {731 %1 = affine.load %a[%i1] : memref<10xf32>732 %2 = arith.divf %0, %1 : f32733 }734 735 // Loops '%i0' and '%i1' should not be fused because the defining node of '%0'736 // used in '%i1' has dependence from loop '%i0'.737 // CHECK: affine.for %{{.*}} = 0 to 10 {738 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[] : memref<f32>739 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>740 // CHECK-NEXT: }741 // CHECK-NEXT: affine.load %{{.*}}[] : memref<f32>742 // CHECK: affine.for %{{.*}} = 0 to 10 {743 // CHECK-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<10xf32>744 // CHECK-NEXT: arith.divf745 // CHECK-NEXT: }746 return747}748 749// -----750 751// CHECK-LABEL: func @should_not_fuse_defining_node_has_transitive_dependence_from_source_loop752func.func @should_not_fuse_defining_node_has_transitive_dependence_from_source_loop(753 %a : memref<10xf32>, %b : memref<10xf32>, %c : memref<f32>) -> () {754 %cst = arith.constant 0.000000e+00 : f32755 affine.for %i0 = 0 to 10 {756 affine.store %cst, %a[%i0] : memref<10xf32>757 affine.store %cst, %b[%i0] : memref<10xf32>758 }759 affine.for %i1 = 0 to 10 {760 %1 = affine.load %b[%i1] : memref<10xf32>761 affine.store %1, %c[] : memref<f32>762 }763 %0 = affine.load %c[] : memref<f32>764 affine.for %i2 = 0 to 10 {765 %1 = affine.load %a[%i2] : memref<10xf32>766 %2 = arith.divf %0, %1 : f32767 }768 769 // When loops '%i0' and '%i2' are evaluated first, they should not be770 // fused. The defining node of '%0' in loop '%i2' has transitive dependence771 // from loop '%i0'. After that, loops '%i0' and '%i1' are evaluated, and they772 // will be fused as usual.773 // CHECK: affine.for %{{.*}} = 0 to 10 {774 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>775 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>776 // CHECK-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<10xf32>777 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[] : memref<f32>778 // CHECK-NEXT: }779 // CHECK-NEXT: affine.load %{{.*}}[] : memref<f32>780 // CHECK: affine.for %{{.*}} = 0 to 10 {781 // CHECK-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<10xf32>782 // CHECK-NEXT: arith.divf783 // CHECK-NEXT: }784 // CHECK-NOT: affine.for785 return786}787 788// -----789 790// CHECK-LABEL: func @should_not_fuse_dest_loop_nest_return_value791func.func @should_not_fuse_dest_loop_nest_return_value(792 %a : memref<10xf32>) -> () {793 %cst = arith.constant 0.000000e+00 : f32794 affine.for %i0 = 0 to 10 {795 affine.store %cst, %a[%i0] : memref<10xf32>796 }797 %b = affine.for %i1 = 0 to 10 step 2 iter_args(%b_iter = %cst) -> f32 {798 %load_a = affine.load %a[%i1] : memref<10xf32>799 affine.yield %load_a: f32800 }801 802 // CHECK: affine.for %{{.*}} = 0 to 10 {803 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>804 // CHECK-NEXT: }805 // CHECK: affine.for %{{.*}} = 0 to 10 step 2 iter_args(%{{.*}} = %{{.*}}) -> (f32) {806 // CHECK-NEXT: affine.load807 // CHECK-NEXT: affine.yield808 // CHECK-NEXT: }809 810 return811}812 813// -----814 815// CHECK-LABEL: func @should_not_fuse_src_loop_nest_return_value816func.func @should_not_fuse_src_loop_nest_return_value(817 %a : memref<10xf32>) -> () {818 %cst = arith.constant 1.000000e+00 : f32819 %b = affine.for %i = 0 to 10 step 2 iter_args(%b_iter = %cst) -> f32 {820 %c = arith.addf %b_iter, %b_iter : f32821 affine.store %c, %a[%i] : memref<10xf32>822 affine.yield %c: f32823 }824 affine.for %i1 = 0 to 10 {825 %1 = affine.load %a[%i1] : memref<10xf32>826 }827 828 // CHECK: %{{.*}} = affine.for %{{.*}} = 0 to 10 step 2 iter_args(%{{.*}} = %{{.*}}) -> (f32) {829 // CHECK-NEXT: %{{.*}} = arith.addf %{{.*}}, %{{.*}} : f32830 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>831 // CHECK-NEXT: affine.yield %{{.*}} : f32832 // CHECK-NEXT: }833 // CHECK: affine.for %{{.*}} = 0 to 10 {834 // CHECK-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<10xf32>835 // CHECK-NEXT: }836 837 return838}839 840// -----841 842func.func private @some_function(memref<16xf32>)843func.func @call_op_prevents_fusion(%arg0: memref<16xf32>){844 %A = memref.alloc() : memref<16xf32>845 %cst_1 = arith.constant 1.000000e+00 : f32846 affine.for %arg1 = 0 to 16 {847 %a = affine.load %arg0[%arg1] : memref<16xf32>848 affine.store %a, %A[%arg1] : memref<16xf32>849 }850 call @some_function(%A) : (memref<16xf32>) -> ()851 %B = memref.alloc() : memref<16xf32>852 affine.for %arg1 = 0 to 16 {853 %a = affine.load %A[%arg1] : memref<16xf32>854 %b = arith.addf %cst_1, %a : f32855 affine.store %b, %B[%arg1] : memref<16xf32>856 }857 return858}859// CHECK-LABEL: func @call_op_prevents_fusion860// CHECK: affine.for861// CHECK-NEXT: affine.load862// CHECK-NEXT: affine.store863// CHECK: call864// CHECK: affine.for865// CHECK-NEXT: affine.load866// CHECK-NEXT: arith.addf867// CHECK-NEXT: affine.store868 869// -----870 871func.func private @some_function()872func.func @call_op_does_not_prevent_fusion(%arg0: memref<16xf32>){873 %A = memref.alloc() : memref<16xf32>874 %cst_1 = arith.constant 1.000000e+00 : f32875 affine.for %arg1 = 0 to 16 {876 %a = affine.load %arg0[%arg1] : memref<16xf32>877 affine.store %a, %A[%arg1] : memref<16xf32>878 }879 call @some_function() : () -> ()880 %B = memref.alloc() : memref<16xf32>881 affine.for %arg1 = 0 to 16 {882 %a = affine.load %A[%arg1] : memref<16xf32>883 %b = arith.addf %cst_1, %a : f32884 affine.store %b, %B[%arg1] : memref<16xf32>885 }886 return887}888// CHECK-LABEL: func @call_op_does_not_prevent_fusion889// CHECK: affine.for890// CHECK-NOT: affine.for891 892// -----893 894// Test for source that writes to an escaping memref and has two consumers.895// Fusion should create private memrefs in place of `%arg0` since the source is896// not to be removed after fusion and the destinations do not write to `%arg0`.897// This should enable both the consumers to benefit from fusion, which would not898// be possible if private memrefs were not created.899func.func @should_fuse_with_both_consumers_separately(%arg0: memref<10xf32>) {900 %cf7 = arith.constant 7.0 : f32901 affine.for %i0 = 0 to 10 {902 affine.store %cf7, %arg0[%i0] : memref<10xf32>903 }904 affine.for %i1 = 0 to 7 {905 %v0 = affine.load %arg0[%i1] : memref<10xf32>906 }907 affine.for %i1 = 5 to 9 {908 %v0 = affine.load %arg0[%i1] : memref<10xf32>909 }910 return911}912 913// CHECK-LABEL: func @should_fuse_with_both_consumers_separately914// CHECK: affine.for915// CHECK-NEXT: affine.store916// CHECK: affine.for917// CHECK-NEXT: affine.store918// CHECK-NEXT: affine.load919// CHECK: affine.for920// CHECK-NEXT: affine.store921// CHECK-NEXT: affine.load922 923// -----924 925// Fusion is avoided when the slice computed is invalid. Comments below describe926// incorrect backward slice computation. Similar logic applies for forward slice927// as well.928func.func @no_fusion_cannot_compute_valid_slice() {929 %A = memref.alloc() : memref<5xf32>930 %B = memref.alloc() : memref<6xf32>931 %C = memref.alloc() : memref<5xf32>932 %cst = arith.constant 0. : f32933 934 affine.for %arg0 = 0 to 5 {935 %a = affine.load %A[%arg0] : memref<5xf32>936 affine.store %a, %B[%arg0 + 1] : memref<6xf32>937 }938 939 affine.for %arg0 = 0 to 5 {940 // Backward slice computed will be:941 // slice ( src loop: 0, dst loop: 1, depth: 1 : insert point: (1, 0)942 // loop bounds: [(d0) -> (d0 - 1), (d0) -> (d0)] )943 944 // Resulting fusion would be as below. It is easy to note the out-of-bounds945 // access by 'affine.load'.946 947 // #map0 = affine_map<(d0) -> (d0 - 1)>948 // #map1 = affine_map<(d0) -> (d0)>949 // affine.for %arg1 = #map0(%arg0) to #map1(%arg0) {950 // %5 = affine.load %1[%arg1] : memref<5xf32>951 // ...952 // ...953 // }954 955 %a = affine.load %B[%arg0] : memref<6xf32>956 %b = arith.mulf %a, %cst : f32957 affine.store %b, %C[%arg0] : memref<5xf32>958 }959 return960}961// CHECK-LABEL: func @no_fusion_cannot_compute_valid_slice962// CHECK: affine.for963// CHECK-NEXT: affine.load964// CHECK-NEXT: affine.store965// CHECK: affine.for966// CHECK-NEXT: affine.load967// CHECK-NEXT: arith.mulf968// CHECK-NEXT: affine.store969 970// MAXIMAL-LABEL: func @reduce_add_f32_f32(971func.func @reduce_add_f32_f32(%arg0: memref<64x64xf32, 1>, %arg1: memref<1x64xf32, 1>, %arg2: memref<1x64xf32, 1>) {972 %cst_0 = arith.constant 0.000000e+00 : f32973 %cst_1 = arith.constant 1.000000e+00 : f32974 %0 = memref.alloca() : memref<f32, 1>975 %1 = memref.alloca() : memref<f32, 1>976 affine.for %arg3 = 0 to 1 {977 affine.for %arg4 = 0 to 64 {978 %accum = affine.for %arg5 = 0 to 64 iter_args (%prevAccum = %cst_0) -> f32 {979 %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>980 %5 = arith.addf %prevAccum, %4 : f32981 affine.yield %5 : f32982 }983 %accum_dbl = arith.addf %accum, %accum : f32984 affine.store %accum_dbl, %arg1[%arg3, %arg4] : memref<1x64xf32, 1>985 }986 }987 affine.for %arg3 = 0 to 1 {988 affine.for %arg4 = 0 to 64 {989 %accum = affine.for %arg5 = 0 to 64 iter_args (%prevAccum = %cst_1) -> f32 {990 %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>991 %5 = arith.mulf %prevAccum, %4 : f32992 affine.yield %5 : f32993 }994 %accum_sqr = arith.mulf %accum, %accum : f32995 affine.store %accum_sqr, %arg2[%arg3, %arg4] : memref<1x64xf32, 1>996 }997 }998 return999}1000// The two loops here get maximally sibling-fused at the innermost1001// insertion point. Test checks if the innermost reduction loop of the fused loop1002// gets promoted into its outerloop.1003// MAXIMAL-SAME: %[[arg_0:.*]]: memref<64x64xf32, 1>,1004// MAXIMAL-SAME: %[[arg_1:.*]]: memref<1x64xf32, 1>,1005// MAXIMAL-SAME: %[[arg_2:.*]]: memref<1x64xf32, 1>) {1006// MAXIMAL: %[[cst:.*]] = arith.constant 0 : index1007// MAXIMAL-NEXT: %[[cst_0:.*]] = arith.constant 0.000000e+00 : f321008// MAXIMAL-NEXT: %[[cst_1:.*]] = arith.constant 1.000000e+00 : f321009// MAXIMAL: affine.for %[[idx_0:.*]] = 0 to 1 {1010// MAXIMAL-NEXT: affine.for %[[idx_1:.*]] = 0 to 64 {1011// MAXIMAL-NEXT: %[[results:.*]]:2 = affine.for %[[idx_2:.*]] = 0 to 64 iter_args(%[[iter_0:.*]] = %[[cst_1]], %[[iter_1:.*]] = %[[cst_0]]) -> (f32, f32) {1012// MAXIMAL-NEXT: %[[val_0:.*]] = affine.load %[[arg_0]][%[[idx_2]], %[[idx_1]]] : memref<64x64xf32, 1>1013// MAXIMAL-NEXT: %[[reduc_0:.*]] = arith.addf %[[iter_1]], %[[val_0]] : f321014// MAXIMAL-NEXT: %[[val_1:.*]] = affine.load %[[arg_0]][%[[idx_2]], %[[idx_1]]] : memref<64x64xf32, 1>1015// MAXIMAL-NEXT: %[[reduc_1:.*]] = arith.mulf %[[iter_0]], %[[val_1]] : f321016// MAXIMAL-NEXT: affine.yield %[[reduc_1]], %[[reduc_0]] : f32, f321017// MAXIMAL-NEXT: }1018// MAXIMAL-NEXT: %[[reduc_0_dbl:.*]] = arith.addf %[[results:.*]]#1, %[[results]]#1 : f321019// MAXIMAL-NEXT: affine.store %[[reduc_0_dbl]], %[[arg_1]][%[[cst]], %[[idx_1]]] : memref<1x64xf32, 1>1020// MAXIMAL-NEXT: %[[reduc_1_sqr:.*]] = arith.mulf %[[results]]#0, %[[results]]#0 : f321021// MAXIMAL-NEXT: affine.store %[[reduc_1_sqr]], %[[arg_2]][%[[idx_0]], %[[idx_1]]] : memref<1x64xf32, 1>1022// MAXIMAL-NEXT: }1023// MAXIMAL-NEXT: }1024// MAXIMAL-NEXT: return1025// MAXIMAL-NEXT: }1026 1027// -----1028 1029// CHECK-LABEL: func @reduce_add_non_innermost1030func.func @reduce_add_non_innermost(%arg0: memref<64x64xf32, 1>, %arg1: memref<1x64xf32, 1>, %arg2: memref<1x64xf32, 1>) {1031 %cst = arith.constant 0.000000e+00 : f321032 %cst_0 = arith.constant 1.000000e+00 : f321033 %0 = memref.alloca() : memref<f32, 1>1034 %1 = memref.alloca() : memref<f32, 1>1035 affine.for %arg3 = 0 to 1 {1036 affine.for %arg4 = 0 to 64 {1037 %accum = affine.for %arg5 = 0 to 64 iter_args (%prevAccum = %cst) -> f32 {1038 %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>1039 %5 = arith.addf %prevAccum, %4 : f321040 affine.yield %5 : f321041 }1042 %accum_dbl = arith.addf %accum, %accum : f321043 affine.store %accum_dbl, %arg1[%arg3, %arg4] : memref<1x64xf32, 1>1044 }1045 }1046 affine.for %arg3 = 0 to 1 {1047 affine.for %arg4 = 0 to 64 {1048 %accum = affine.for %arg5 = 0 to 64 iter_args (%prevAccum = %cst_0) -> f32 {1049 %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>1050 %5 = arith.mulf %prevAccum, %4 : f321051 affine.yield %5 : f321052 }1053 %accum_sqr = arith.mulf %accum, %accum : f321054 affine.store %accum_sqr, %arg2[%arg3, %arg4] : memref<1x64xf32, 1>1055 }1056 }1057 return1058}1059// Test checks the loop structure is preserved after sibling fusion.1060// CHECK: affine.for1061// CHECK-NEXT: affine.for1062// CHECK-NEXT: affine.for1063// CHECK: affine.for1064 1065 1066 1067// -----1068 1069// CHECK-LABEL: func @fuse_large_number_of_loops1070func.func @fuse_large_number_of_loops(%arg0: memref<20x10xf32, 1>, %arg1: memref<20x10xf32, 1>, %arg2: memref<20x10xf32, 1>, %arg3: memref<20x10xf32, 1>, %arg4: memref<20x10xf32, 1>, %arg5: memref<f32, 1>, %arg6: memref<f32, 1>, %arg7: memref<f32, 1>, %arg8: memref<f32, 1>, %arg9: memref<20x10xf32, 1>, %arg10: memref<20x10xf32, 1>, %arg11: memref<20x10xf32, 1>, %arg12: memref<20x10xf32, 1>) {1071 %cst = arith.constant 1.000000e+00 : f321072 %0 = memref.alloc() : memref<f32, 1>1073 affine.store %cst, %0[] : memref<f32, 1>1074 %1 = memref.alloc() : memref<20x10xf32, 1>1075 affine.for %arg13 = 0 to 20 {1076 affine.for %arg14 = 0 to 10 {1077 %21 = affine.load %arg6[] : memref<f32, 1>1078 affine.store %21, %1[%arg13, %arg14] : memref<20x10xf32, 1>1079 }1080 }1081 %2 = memref.alloc() : memref<20x10xf32, 1>1082 affine.for %arg13 = 0 to 20 {1083 affine.for %arg14 = 0 to 10 {1084 %21 = affine.load %1[%arg13, %arg14] : memref<20x10xf32, 1>1085 %22 = affine.load %arg3[%arg13, %arg14] : memref<20x10xf32, 1>1086 %23 = arith.mulf %22, %21 : f321087 affine.store %23, %2[%arg13, %arg14] : memref<20x10xf32, 1>1088 }1089 }1090 %3 = memref.alloc() : memref<f32, 1>1091 %4 = affine.load %arg6[] : memref<f32, 1>1092 %5 = affine.load %0[] : memref<f32, 1>1093 %6 = arith.subf %5, %4 : f321094 affine.store %6, %3[] : memref<f32, 1>1095 %7 = memref.alloc() : memref<20x10xf32, 1>1096 affine.for %arg13 = 0 to 20 {1097 affine.for %arg14 = 0 to 10 {1098 %21 = affine.load %3[] : memref<f32, 1>1099 affine.store %21, %7[%arg13, %arg14] : memref<20x10xf32, 1>1100 }1101 }1102 %8 = memref.alloc() : memref<20x10xf32, 1>1103 affine.for %arg13 = 0 to 20 {1104 affine.for %arg14 = 0 to 10 {1105 %21 = affine.load %arg1[%arg13, %arg14] : memref<20x10xf32, 1>1106 %22 = affine.load %7[%arg13, %arg14] : memref<20x10xf32, 1>1107 %23 = arith.mulf %22, %21 : f321108 affine.store %23, %8[%arg13, %arg14] : memref<20x10xf32, 1>1109 }1110 }1111 %9 = memref.alloc() : memref<20x10xf32, 1>1112 affine.for %arg13 = 0 to 20 {1113 affine.for %arg14 = 0 to 10 {1114 %21 = affine.load %arg1[%arg13, %arg14] : memref<20x10xf32, 1>1115 %22 = affine.load %8[%arg13, %arg14] : memref<20x10xf32, 1>1116 %23 = arith.mulf %22, %21 : f321117 affine.store %23, %9[%arg13, %arg14] : memref<20x10xf32, 1>1118 }1119 }1120 affine.for %arg13 = 0 to 20 {1121 affine.for %arg14 = 0 to 10 {1122 %21 = affine.load %9[%arg13, %arg14] : memref<20x10xf32, 1>1123 %22 = affine.load %2[%arg13, %arg14] : memref<20x10xf32, 1>1124 %23 = arith.addf %22, %21 : f321125 affine.store %23, %arg11[%arg13, %arg14] : memref<20x10xf32, 1>1126 }1127 }1128 %10 = memref.alloc() : memref<20x10xf32, 1>1129 affine.for %arg13 = 0 to 20 {1130 affine.for %arg14 = 0 to 10 {1131 %21 = affine.load %1[%arg13, %arg14] : memref<20x10xf32, 1>1132 %22 = affine.load %arg2[%arg13, %arg14] : memref<20x10xf32, 1>1133 %23 = arith.mulf %22, %21 : f321134 affine.store %23, %10[%arg13, %arg14] : memref<20x10xf32, 1>1135 }1136 }1137 affine.for %arg13 = 0 to 20 {1138 affine.for %arg14 = 0 to 10 {1139 %21 = affine.load %8[%arg13, %arg14] : memref<20x10xf32, 1>1140 %22 = affine.load %10[%arg13, %arg14] : memref<20x10xf32, 1>1141 %23 = arith.addf %22, %21 : f321142 affine.store %23, %arg10[%arg13, %arg14] : memref<20x10xf32, 1>1143 }1144 }1145 %11 = memref.alloc() : memref<20x10xf32, 1>1146 affine.for %arg13 = 0 to 20 {1147 affine.for %arg14 = 0 to 10 {1148 %21 = affine.load %arg10[%arg13, %arg14] : memref<20x10xf32, 1>1149 %22 = affine.load %arg10[%arg13, %arg14] : memref<20x10xf32, 1>1150 %23 = arith.mulf %22, %21 : f321151 affine.store %23, %11[%arg13, %arg14] : memref<20x10xf32, 1>1152 }1153 }1154 %12 = memref.alloc() : memref<20x10xf32, 1>1155 affine.for %arg13 = 0 to 20 {1156 affine.for %arg14 = 0 to 10 {1157 %21 = affine.load %11[%arg13, %arg14] : memref<20x10xf32, 1>1158 %22 = affine.load %arg11[%arg13, %arg14] : memref<20x10xf32, 1>1159 %23 = arith.subf %22, %21 : f321160 affine.store %23, %12[%arg13, %arg14] : memref<20x10xf32, 1>1161 }1162 }1163 %13 = memref.alloc() : memref<20x10xf32, 1>1164 affine.for %arg13 = 0 to 20 {1165 affine.for %arg14 = 0 to 10 {1166 %21 = affine.load %arg7[] : memref<f32, 1>1167 affine.store %21, %13[%arg13, %arg14] : memref<20x10xf32, 1>1168 }1169 }1170 %14 = memref.alloc() : memref<20x10xf32, 1>1171 affine.for %arg13 = 0 to 20 {1172 affine.for %arg14 = 0 to 10 {1173 %21 = affine.load %arg4[%arg13, %arg14] : memref<20x10xf32, 1>1174 %22 = affine.load %13[%arg13, %arg14] : memref<20x10xf32, 1>1175 %23 = arith.mulf %22, %21 : f321176 affine.store %23, %14[%arg13, %arg14] : memref<20x10xf32, 1>1177 }1178 }1179 %15 = memref.alloc() : memref<20x10xf32, 1>1180 affine.for %arg13 = 0 to 20 {1181 affine.for %arg14 = 0 to 10 {1182 %21 = affine.load %arg8[] : memref<f32, 1>1183 affine.store %21, %15[%arg13, %arg14] : memref<20x10xf32, 1>1184 }1185 }1186 %16 = memref.alloc() : memref<20x10xf32, 1>1187 affine.for %arg13 = 0 to 20 {1188 affine.for %arg14 = 0 to 10 {1189 %21 = affine.load %15[%arg13, %arg14] : memref<20x10xf32, 1>1190 %22 = affine.load %12[%arg13, %arg14] : memref<20x10xf32, 1>1191 %23 = arith.addf %22, %21 : f321192 affine.store %23, %16[%arg13, %arg14] : memref<20x10xf32, 1>1193 }1194 }1195 %17 = memref.alloc() : memref<20x10xf32, 1>1196 affine.for %arg13 = 0 to 20 {1197 affine.for %arg14 = 0 to 10 {1198 %21 = affine.load %16[%arg13, %arg14] : memref<20x10xf32, 1>1199 %22 = math.sqrt %21 : f321200 affine.store %22, %17[%arg13, %arg14] : memref<20x10xf32, 1>1201 }1202 }1203 %18 = memref.alloc() : memref<20x10xf32, 1>1204 affine.for %arg13 = 0 to 20 {1205 affine.for %arg14 = 0 to 10 {1206 %21 = affine.load %arg5[] : memref<f32, 1>1207 affine.store %21, %18[%arg13, %arg14] : memref<20x10xf32, 1>1208 }1209 }1210 %19 = memref.alloc() : memref<20x10xf32, 1>1211 affine.for %arg13 = 0 to 20 {1212 affine.for %arg14 = 0 to 10 {1213 %21 = affine.load %arg1[%arg13, %arg14] : memref<20x10xf32, 1>1214 %22 = affine.load %18[%arg13, %arg14] : memref<20x10xf32, 1>1215 %23 = arith.mulf %22, %21 : f321216 affine.store %23, %19[%arg13, %arg14] : memref<20x10xf32, 1>1217 }1218 }1219 %20 = memref.alloc() : memref<20x10xf32, 1>1220 affine.for %arg13 = 0 to 20 {1221 affine.for %arg14 = 0 to 10 {1222 %21 = affine.load %17[%arg13, %arg14] : memref<20x10xf32, 1>1223 %22 = affine.load %19[%arg13, %arg14] : memref<20x10xf32, 1>1224 %23 = arith.divf %22, %21 : f321225 affine.store %23, %20[%arg13, %arg14] : memref<20x10xf32, 1>1226 }1227 }1228 affine.for %arg13 = 0 to 20 {1229 affine.for %arg14 = 0 to 10 {1230 %21 = affine.load %20[%arg13, %arg14] : memref<20x10xf32, 1>1231 %22 = affine.load %14[%arg13, %arg14] : memref<20x10xf32, 1>1232 %23 = arith.addf %22, %21 : f321233 affine.store %23, %arg12[%arg13, %arg14] : memref<20x10xf32, 1>1234 }1235 }1236 affine.for %arg13 = 0 to 20 {1237 affine.for %arg14 = 0 to 10 {1238 %21 = affine.load %arg12[%arg13, %arg14] : memref<20x10xf32, 1>1239 %22 = affine.load %arg0[%arg13, %arg14] : memref<20x10xf32, 1>1240 %23 = arith.subf %22, %21 : f321241 affine.store %23, %arg9[%arg13, %arg14] : memref<20x10xf32, 1>1242 }1243 }1244 return1245}1246// CHECK: affine.for1247// CHECK: affine.for1248// CHECK-NOT: affine.for1249 1250// CHECK-LABEL: func @alias_escaping_memref1251func.func @alias_escaping_memref(%a : memref<2x5xf32>) {1252 %cst = arith.constant 0.000000e+00 : f321253 %alias = memref.reinterpret_cast %a to offset: [0], sizes: [10], strides: [1] : memref<2x5xf32> to memref<10xf32>1254 affine.for %i0 = 0 to 10 {1255 affine.store %cst, %alias[%i0] : memref<10xf32>1256 }1257 1258 affine.for %i1 = 0 to 10 {1259 %0 = affine.load %alias[%i1] : memref<10xf32>1260 }1261 // Fusion happens, but memref isn't privatized since %alias is an alias of a1262 // function argument.1263 // CHECK: memref.reinterpret_cast1264 // CHECK-NEXT: affine.for1265 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>1266 // CHECK-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<10xf32>1267 // CHECK-NEXT: }1268 // CHECK-NOT: affine.for1269 1270 return1271}1272 1273// CHECK-LABEL: func @unknown_memref_def_op1274func.func @unknown_memref_def_op() {1275 %cst = arith.constant 0.000000e+00 : f321276 %may_alias = call @bar() : () -> memref<10xf32>1277 affine.for %i0 = 0 to 10 {1278 affine.store %cst, %may_alias[%i0] : memref<10xf32>1279 }1280 1281 affine.for %i1 = 0 to 10 {1282 %0 = affine.load %may_alias[%i1] : memref<10xf32>1283 }1284 // Fusion happens, but memref isn't privatized since %may_alias's origin is1285 // unknown.1286 // CHECK: call1287 // CHECK-NEXT: affine.for1288 // CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>1289 // CHECK-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<10xf32>1290 // CHECK-NEXT: }1291 // CHECK-NOT: affine.for1292 1293 return1294}1295func.func private @bar() -> memref<10xf32>1296 1297 1298// Add further tests in mlir/test/Transforms/loop-fusion-4.mlir1299