774 lines · plain
1// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{mode=producer}))' -split-input-file | FileCheck %s --check-prefix=PRODUCER-CONSUMER2// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{compute-tolerance=0.0}))' -split-input-file | FileCheck %s --check-prefix=ZERO-TOLERANCE3// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{mode=producer maximal}))' -split-input-file | FileCheck %s --check-prefix=PRODUCER-CONSUMER-MAXIMAL4// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion{maximal mode=sibling}))' -split-input-file | FileCheck %s --check-prefix=SIBLING-MAXIMAL5// All fusion: producer-consumer and sibling.6// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(affine-loop-fusion))' -split-input-file | FileCheck %s --check-prefix=ALL7// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(spirv.func(affine-loop-fusion{mode=producer}))' -split-input-file | FileCheck %s --check-prefix=SPIRV8 9// Part I of fusion tests in mlir/test/Transforms/loop-fusion.mlir.10// Part II of fusion tests in mlir/test/Transforms/loop-fusion-2.mlir11// Part III of fusion tests in mlir/test/Transforms/loop-fusion-3.mlir12 13// Expects fusion of producer into consumer at depth 4 and subsequent removal of14// source loop.15// PRODUCER-CONSUMER-LABEL: func @unflatten4d16func.func @unflatten4d(%arg1: memref<7x8x9x10xf32>) {17 %m = memref.alloc() : memref<5040xf32>18 %cf7 = arith.constant 7.0 : f3219 20 affine.for %i0 = 0 to 7 {21 affine.for %i1 = 0 to 8 {22 affine.for %i2 = 0 to 9 {23 affine.for %i3 = 0 to 10 {24 affine.store %cf7, %m[720 * %i0 + 90 * %i1 + 10 * %i2 + %i3] : memref<5040xf32>25 }26 }27 }28 }29 affine.for %i0 = 0 to 7 {30 affine.for %i1 = 0 to 8 {31 affine.for %i2 = 0 to 9 {32 affine.for %i3 = 0 to 10 {33 %v0 = affine.load %m[720 * %i0 + 90 * %i1 + 10 * %i2 + %i3] : memref<5040xf32>34 affine.store %v0, %arg1[%i0, %i1, %i2, %i3] : memref<7x8x9x10xf32>35 }36 }37 }38 }39 return40}41 42// PRODUCER-CONSUMER: affine.for43// PRODUCER-CONSUMER-NEXT: affine.for44// PRODUCER-CONSUMER-NEXT: affine.for45// PRODUCER-CONSUMER-NEXT: affine.for46// PRODUCER-CONSUMER-NOT: affine.for47// PRODUCER-CONSUMER: return48 49// -----50 51// Expects fusion of producer into consumer at depth 2 and subsequent removal of52// source loop.53// PRODUCER-CONSUMER-LABEL: func @unflatten2d_with_transpose54func.func @unflatten2d_with_transpose(%arg1: memref<8x7xf32>) {55 %m = memref.alloc() : memref<56xf32>56 %cf7 = arith.constant 7.0 : f3257 58 affine.for %i0 = 0 to 7 {59 affine.for %i1 = 0 to 8 {60 affine.store %cf7, %m[8 * %i0 + %i1] : memref<56xf32>61 }62 }63 affine.for %i0 = 0 to 8 {64 affine.for %i1 = 0 to 7 {65 %v0 = affine.load %m[%i0 + 8 * %i1] : memref<56xf32>66 affine.store %v0, %arg1[%i0, %i1] : memref<8x7xf32>67 }68 }69 return70}71 72// PRODUCER-CONSUMER: affine.for73// PRODUCER-CONSUMER-NEXT: affine.for74// PRODUCER-CONSUMER-NOT: affine.for75// PRODUCER-CONSUMER: return76 77// -----78 79// Expects fusion of producer into consumer at depth 1 and source loop to not80// be removed due to difference in loop steps.81// PRODUCER-CONSUMER-LABEL: func @check_src_dst_step82func.func @check_src_dst_step(%m : memref<100xf32>,83 %src: memref<100xf32>,84 %out: memref<100xf32>) {85 affine.for %i0 = 0 to 100 {86 %r1 = affine.load %src[%i0]: memref<100xf32>87 affine.store %r1, %m[%i0] : memref<100xf32>88 }89 affine.for %i2 = 0 to 100 step 2 {90 %r2 = affine.load %m[%i2] : memref<100xf32>91 affine.store %r2, %out[%i2] : memref<100xf32>92 }93 return94}95 96// Check if the fusion did take place as well as that the source loop was97// not removed. To check if fusion took place, the read instruction from the98// original source loop is checked to be in the fused loop.99//100// PRODUCER-CONSUMER: affine.for %[[idx_0:.*]] = 0 to 100 {101// PRODUCER-CONSUMER-NEXT: %[[result_0:.*]] = affine.load %[[arr1:.*]][%[[idx_0]]] : memref<100xf32>102// PRODUCER-CONSUMER-NEXT: affine.store %[[result_0]], %{{.*}}[%[[idx_0]]] : memref<100xf32>103// PRODUCER-CONSUMER-NEXT: }104// PRODUCER-CONSUMER: affine.for %[[idx_1:.*]] = 0 to 100 step 2 {105// PRODUCER-CONSUMER: affine.load %[[arr1]][%[[idx_1]]] : memref<100xf32>106// PRODUCER-CONSUMER: }107// PRODUCER-CONSUMER: return108 109// -----110 111// SIBLING-MAXIMAL-LABEL: func @reduce_add_non_maximal_f32_f32(112func.func @reduce_add_non_maximal_f32_f32(%arg0: memref<64x64xf32, 1>, %arg1 : memref<1x64xf32, 1>, %arg2 : memref<1x64xf32, 1>) {113 %cst_0 = arith.constant 0.000000e+00 : f32114 %cst_1 = arith.constant 1.000000e+00 : f32115 // This nest writes to %arg1 but can be eliminated post sibling fusion.116 affine.for %arg3 = 0 to 1 {117 affine.for %arg4 = 0 to 64 {118 %accum = affine.for %arg5 = 0 to 64 iter_args (%prevAccum = %cst_0) -> f32 {119 %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>120 %5 = arith.addf %prevAccum, %4 : f32121 affine.yield %5 : f32122 }123 %accum_dbl = arith.addf %accum, %accum : f32124 affine.store %accum_dbl, %arg1[%arg3, %arg4] : memref<1x64xf32, 1>125 }126 }127 affine.for %arg3 = 0 to 1 {128 affine.for %arg4 = 0 to 64 {129 // Following loop trip count does not match the corresponding source trip count.130 %accum = affine.for %arg5 = 0 to 32 iter_args (%prevAccum = %cst_1) -> f32 {131 %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>132 %5 = arith.mulf %prevAccum, %4 : f32133 affine.yield %5 : f32134 }135 %accum_sqr = arith.mulf %accum, %accum : f32136 affine.store %accum_sqr, %arg2[%arg3, %arg4] : memref<1x64xf32, 1>137 }138 }139 return140}141// Test checks the loop structure is preserved after sibling fusion142// since the destination loop and source loop trip counts do not143// match.144// SIBLING-MAXIMAL: %[[cst_0:.*]] = arith.constant 0.000000e+00 : f32145// SIBLING-MAXIMAL-NEXT: %[[cst_1:.*]] = arith.constant 1.000000e+00 : f32146// SIBLING-MAXIMAL-NEXT: affine.for %{{.*}} = 0 to 1 {147// SIBLING-MAXIMAL-NEXT: affine.for %{{.*}} = 0 to 64 {148// SIBLING-MAXIMAL-NEXT: affine.for %{{.*}} = 0 to 32 iter_args(%{{.*}} = %[[cst_1]]) -> (f32) {149// SIBLING-MAXIMAL-NEXT: affine.for %{{.*}} = 0 to 64 iter_args(%{{.*}} = %[[cst_0]]) -> (f32) {150 151// -----152 153// SIBLING-MAXIMAL-LABEL: func @sibling_load_only154func.func @sibling_load_only(%arg0: memref<10xf32>) {155 affine.for %arg1 = 0 to 10 {156 %0 = affine.load %arg0[%arg1] : memref<10xf32>157 }158 affine.for %arg1 = 0 to 10 {159 %0 = affine.load %arg0[%arg1] : memref<10xf32>160 }161 // SIBLING-MAXIMAL-NEXT: affine.for162 // SIBLING-MAXIMAL-NEXT: affine.load163 // SIBLING-MAXIMAL-NEXT: affine.load164 return165}166 167// -----168 169// PRODUCER-CONSUMER-LABEL: func @fusion_for_multiple_blocks() {170func.func @fusion_for_multiple_blocks() {171^bb0:172 %m = memref.alloc() : memref<10xf32>173 %cf7 = arith.constant 7.0 : f32174 175 affine.for %i0 = 0 to 10 {176 affine.store %cf7, %m[%i0] : memref<10xf32>177 }178 affine.for %i1 = 0 to 10 {179 %v0 = affine.load %m[%i1] : memref<10xf32>180 }181 // PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 10 {182 // PRODUCER-CONSUMER-NEXT: affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>183 // PRODUCER-CONSUMER-NEXT: affine.load %{{.*}}[0] : memref<1xf32>184 // PRODUCER-CONSUMER-NEXT: }185 cf.br ^bb1186^bb1:187 affine.for %i0 = 0 to 10 {188 affine.store %cf7, %m[%i0] : memref<10xf32>189 }190 affine.for %i1 = 0 to 10 {191 %v0 = affine.load %m[%i1] : memref<10xf32>192 }193 // PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 10 {194 // PRODUCER-CONSUMER-NEXT: affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>195 // PRODUCER-CONSUMER-NEXT: affine.load %{{.*}}[0] : memref<1xf32>196 // PRODUCER-CONSUMER-NEXT: }197 return198}199 200// -----201 202// PRODUCER-CONSUMER-LABEL: @fuse_higher_dim_nest_into_lower_dim_nest203func.func @fuse_higher_dim_nest_into_lower_dim_nest() {204 %A = memref.alloc() : memref<8x12x128x64xf32>205 %B = memref.alloc() : memref<8x128x12x64xf32>206 affine.for %arg205 = 0 to 8 {207 affine.for %arg206 = 0 to 128 {208 affine.for %arg207 = 0 to 12 {209 affine.for %arg208 = 0 to 64 {210 %a = affine.load %A[%arg205, %arg207, %arg206, %arg208] : memref<8x12x128x64xf32>211 affine.store %a, %B[%arg205, %arg206, %arg207, %arg208] : memref<8x128x12x64xf32>212 }213 }214 }215 }216 %C = memref.alloc() : memref<8x128x768xf16>217 affine.for %arg205 = 0 to 8 {218 affine.for %arg206 = 0 to 128 {219 affine.for %arg207 = 0 to 768 {220 %b = affine.load %B[%arg205, %arg206, %arg207 floordiv 64, %arg207 mod 64] : memref<8x128x12x64xf32>221 %c = arith.truncf %b : f32 to f16222 affine.store %c, %C[%arg205, %arg206, %arg207] : memref<8x128x768xf16>223 }224 }225 }226 227 // Check that fusion happens into the innermost loop of the consumer.228 // PRODUCER-CONSUMER: affine.for229 // PRODUCER-CONSUMER-NEXT: affine.for %{{.*}} = 0 to 128230 // PRODUCER-CONSUMER-NEXT: affine.for %{{.*}} = 0 to 768231 // PRODUCER-CONSUMER-NOT: affine.for232 // PRODUCER-CONSUMER: return233 return234}235 236// -----237 238// Basic test to ensure fusion works inside other func ops like spirv.func.239 240#map = affine_map<(d0, d1) -> (d0 + d1)>241module {242 // SPIRV-LABEL: func @test_avgpool2d_pad_right243 spirv.func @test_avgpool2d_pad_right(%arg0: !spirv.array<8192 x f32>) -> !spirv.array<8192 x f32> "None" {244 %cst_f32 = spirv.Constant 0.000000e+00 : f32245 %0 = builtin.unrealized_conversion_cast %arg0 : !spirv.array<8192 x f32> to tensor<1x32x32x8xf32>246 %padded = tensor.pad %0 low[0, 4, 4, 0] high[0, 4, 8193, 0] {247 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):248 tensor.yield %cst_f32 : f32249 } : tensor<1x32x32x8xf32> to tensor<1x40x8229x8xf32>250 %1 = bufferization.to_buffer %padded : tensor<1x40x8229x8xf32> to memref<1x40x8229x8xf32>251 %alloc_0 = memref.alloc() {alignment = 64 : i64} : memref<1x32x32x8xf32>252 affine.for %arg1 = 0 to 1 {253 affine.for %arg2 = 0 to 32 {254 affine.for %arg3 = 0 to 32 {255 affine.for %arg4 = 0 to 8 {256 affine.for %arg5 = 0 to 1 {257 affine.for %arg6 = 0 to 1 {258 %4 = affine.apply #map(%arg2, %arg5)259 %5 = affine.apply #map(%arg3, %arg6)260 %6 = affine.load %1[%arg1, %4, %5, %arg4] : memref<1x40x8229x8xf32>261 %7 = affine.load %alloc_0[%arg1, %arg2, %arg3, %arg4] : memref<1x32x32x8xf32>262 %8 = arith.addf %7, %6 : f32263 affine.store %8, %alloc_0[%arg1, %arg2, %arg3, %arg4] : memref<1x32x32x8xf32>264 }265 }266 }267 }268 }269 }270 %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref<1x32x32x8xf32>271 affine.for %arg1 = 0 to 1 {272 affine.for %arg2 = 0 to 32 {273 affine.for %arg3 = 0 to 32 {274 affine.for %arg4 = 0 to 8 {275 %4 = affine.load %alloc_0[%arg1, %arg2, %arg3, %arg4] : memref<1x32x32x8xf32>276 }277 }278 }279 }280 // Test fusion.281 // SPIRV: affine.for %{{.*}} = 0 to 1 {282 // SPIRV-NEXT: affine.for %{{.*}} = 0 to 32 {283 // SPIRV-NEXT: affine.for %{{.*}} = 0 to 32 {284 // SPIRV-NEXT: affine.for %{{.*}} = 0 to 8 {285 // SPIRV-NOT: affine.for %{{.*}}286 287 // SPIRV: ReturnValue288 %2 = bufferization.to_tensor %alloc_1 : memref<1x32x32x8xf32> to tensor<1x32x32x8xf32>289 %3 = builtin.unrealized_conversion_cast %2 : tensor<1x32x32x8xf32> to !spirv.array<8192 x f32>290 spirv.ReturnValue %3 : !spirv.array<8192 x f32>291 }292}293 294// -----295 296// PRODUCER-CONSUMER-LABEL: func @same_memref_load_store297func.func @same_memref_load_store(%producer : memref<32xf32>, %consumer: memref<16xf32>){298 %cst = arith.constant 2.000000e+00 : f32299 // Source isn't removed.300 // PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 32301 affine.for %arg3 = 0 to 32 {302 %0 = affine.load %producer[%arg3] : memref<32xf32>303 %2 = arith.mulf %0, %cst : f32304 affine.store %2, %producer[%arg3] : memref<32xf32>305 }306 affine.for %arg3 = 0 to 16 {307 %0 = affine.load %producer[%arg3] : memref<32xf32>308 %2 = arith.addf %0, %cst : f32309 affine.store %2, %consumer[%arg3] : memref<16xf32>310 }311 // Fused nest.312 // PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 16313 // PRODUCER-CONSUMER-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<32xf32>314 // PRODUCER-CONSUMER-NEXT: arith.mulf315 // PRODUCER-CONSUMER-NEXT: affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>316 // PRODUCER-CONSUMER-NEXT: affine.load %{{.*}}[0] : memref<1xf32>317 // PRODUCER-CONSUMER-NEXT: arith.addf318 // PRODUCER-CONSUMER-NEXT: affine.store319 // PRODUCER-CONSUMER-NEXT: }320 return321}322 323// -----324 325// PRODUCER-CONSUMER-LABEL: func @same_memref_load_multiple_stores326// ALL-LABEL: func @same_memref_load_multiple_stores327func.func @same_memref_load_multiple_stores(%producer : memref<32xf32>, %producer_2 : memref<32xf32>, %consumer: memref<16xf32>){328 %cst = arith.constant 2.000000e+00 : f32329 // Ensure that source isn't removed during both producer-consumer fusion and330 // sibling fusion.331 // PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 32332 // ALL: affine.for %{{.*}} = 0 to 32333 affine.for %arg3 = 0 to 32 {334 %0 = affine.load %producer[%arg3] : memref<32xf32>335 %2 = arith.mulf %0, %cst : f32336 affine.store %2, %producer[%arg3] : memref<32xf32>337 affine.store %2, %producer_2[%arg3] : memref<32xf32>338 }339 affine.for %arg3 = 0 to 16 {340 %0 = affine.load %producer[%arg3] : memref<32xf32>341 %1 = affine.load %producer_2[%arg3] : memref<32xf32>342 %2 = arith.addf %0, %1 : f32343 affine.store %2, %consumer[%arg3] : memref<16xf32>344 }345 // Fused nest.346 // PRODUCER-CONSUMER: affine.for %{{.*}} = 0 to 16347 // PRODUCER-CONSUMER-NEXT: affine.load %{{.*}}[%{{.*}}] : memref<32xf32>348 // PRODUCER-CONSUMER-NEXT: arith.mulf349 // PRODUCER-CONSUMER-NEXT: affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>350 // PRODUCER-CONSUMER-NEXT: affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>351 // PRODUCER-CONSUMER-NEXT: affine.load %{{.*}}[0] : memref<1xf32>352 // PRODUCER-CONSUMER-NEXT: affine.load %{{.*}}[0] : memref<1xf32>353 // PRODUCER-CONSUMER-NEXT: arith.addf354 // PRODUCER-CONSUMER-NEXT: affine.store355 // PRODUCER-CONSUMER-NEXT: }356 // ALL: affine.for %{{.*}} = 0 to 16357 // ALL: mulf358 // ALL: addf359 return360}361 362// -----363 364#map = affine_map<()[s0] -> (s0 + 5)>365#map1 = affine_map<()[s0] -> (s0 + 17)>366 367// Test with non-int/float memref types.368 369// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @memref_index_type370func.func @memref_index_type() {371 %0 = llvm.mlir.constant(2 : index) : i64372 %2 = llvm.mlir.constant(0 : index) : i64373 %3 = builtin.unrealized_conversion_cast %2 : i64 to index374 %alloc = memref.alloc() {alignment = 64 : i64} : memref<8x18xf32>375 %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref<3xf32>376 %alloc_2 = memref.alloc() {alignment = 64 : i64} : memref<3xindex>377 affine.for %arg3 = 0 to 3 {378 %4 = affine.load %alloc_2[%arg3] : memref<3xindex>379 %5 = builtin.unrealized_conversion_cast %4 : index to i64380 %6 = llvm.sub %0, %5 : i64381 %7 = builtin.unrealized_conversion_cast %6 : i64 to index382 affine.store %7, %alloc_2[%arg3] : memref<3xindex>383 }384 affine.for %arg3 = 0 to 3 {385 %4 = affine.load %alloc_2[%arg3] : memref<3xindex>386 %5 = affine.apply #map()[%4]387 %6 = affine.apply #map1()[%3]388 %7 = memref.load %alloc[%5, %6] : memref<8x18xf32>389 affine.store %7, %alloc_1[%arg3] : memref<3xf32>390 }391 // Expect fusion.392 // PRODUCER-CONSUMER-MAXIMAL: affine.for393 // PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for394 // PRODUCER-CONSUMER-MAXIMAL: return395 return396}397 398// -----399 400#map = affine_map<(d0) -> (d0)>401#map1 =affine_map<(d0) -> (d0 + 1)>402 403// Test non-integer memory spaces.404 405// PRODUCER-CONSUMER-LABEL: func @non_int_memory_space406func.func @non_int_memory_space() {407 %alloc = memref.alloc() : memref<256x8xf32, #spirv.storage_class<StorageBuffer>>408 affine.for %arg0 = 0 to 64 {409 affine.for %arg1 = 0 to 8 {410 %0 = affine.apply #map(%arg1)411 %1 = affine.load %alloc[%arg0, %0] : memref<256x8xf32, #spirv.storage_class<StorageBuffer>>412 affine.store %1, %alloc[%arg0, %arg1] : memref<256x8xf32, #spirv.storage_class<StorageBuffer>>413 }414 }415 affine.for %arg0 = 16 to 32 {416 affine.for %arg1 = 0 to 8 {417 %0 = affine.apply #map(%arg1)418 %1 = affine.load %alloc[%arg0, %0] : memref<256x8xf32, #spirv.storage_class<StorageBuffer>>419 affine.store %1, %alloc[%arg0, %arg1] : memref<256x8xf32, #spirv.storage_class<StorageBuffer>>420 }421 }422 // Fused nest.423 // PRODUCER-CONSUMER-NEXT: memref.alloc()424 // PRODUCER-CONSUMER-NEXT: memref.alloc()425 // PRODUCER-CONSUMER-NEXT: affine.for %{{.*}} = 16 to 32426 // PRODUCER-CONSUMER-NEXT: affine.for %{{.*}} = 0 to 8427 return428}429 430// -----431 432#map = affine_map<(d0) -> (d0)>433#map1 = affine_map<(d0) -> (d0 + 1)>434 435// Exercises fix for crash reported at https://github.com/llvm/llvm-project/issues/119525436 437// No fusion of producer into consumer happens here as the slice is determined438// to be invalid. This is a limitation and it is possible to compute a slice439// (reduction along %arg4) and fuse.440 441// PRODUCER-CONSUMER-LABEL: func @slice_compute_check442func.func @slice_compute_check(%arg0: memref<1x8x26xi32, strided<[?, ?, ?], offset: ?>>, %arg1: memref<1x8x26xi32, strided<[?, ?, ?], offset: ?>>, %arg2: memref<1x8x26xi32, strided<[?, ?, ?], offset: ?>>) {443 %alloc_14 = memref.alloc() : memref<1x8x26xi32>444 %alloc_15 = memref.alloc() : memref<1x26xi32>445 affine.for %arg3 = 0 to 1 {446 affine.for %arg4 = 0 to 8 {447 affine.for %arg5 = 0 to 26 {448 affine.for %arg6 = #map(%arg3) to #map1(%arg3) {449 affine.for %arg7 = #map(%arg4) to #map1(%arg4) {450 affine.for %arg8 = #map(%arg5) to #map1(%arg5) {451 %61 = affine.load %alloc_14[%arg6, %arg7, %arg8] : memref<1x8x26xi32>452 %62 = affine.load %alloc_15[%arg6, %arg8] : memref<1x26xi32>453 %63 = llvm.intr.smin(%61, %62) : (i32, i32) -> i32454 affine.store %63, %alloc_15[%arg6, %arg8] : memref<1x26xi32>455 }456 }457 }458 }459 }460 }461 affine.for %arg3 = 0 to 26 {462 %61 = affine.load %alloc_15[0, %arg3] : memref<1x26xi32>463 }464 memref.dealloc %alloc_15 : memref<1x26xi32>465 memref.dealloc %alloc_14 : memref<1x8x26xi32>466 return467}468 469// -----470 471// Exercises fix for crash reported at https://github.com/llvm/llvm-project/issues/108374472 473// No fusion of producer into consumer happens here. The slice will not be474// valid as the producer doesn't supply to all of the consumer.475 476#map = affine_map<(d0) -> (d0)>477#map1 = affine_map<(d0) -> (d0 + 1)>478// PRODUCER-CONSUMER-LABEL: func @test_add_slice_bounds479func.func @test_add_slice_bounds() {480 %alloc = memref.alloc() : memref<10xf32>481 %cst = arith.constant 0.619152 : f32482 affine.for %arg0 = 0 to 10 {483 affine.for %arg1 = #map(%arg0) to #map1(%arg0) {484 affine.store %cst, %alloc[%arg1] : memref<10xf32>485 }486 }487 affine.for %arg0 = 0 to 3 {488 affine.for %arg1 = 0 to 10 {489 affine.for %arg2 = #map(%arg0) to #map1(%arg0) {490 affine.for %arg3 = #map(%arg1) to #map1(%arg1) {491 %0 = affine.apply #map1(%arg3)492 %1 = affine.load %alloc[%0] : memref<10xf32>493 }494 }495 }496 }497 return498}499 500// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @producer_reduction_no_fusion501func.func @producer_reduction_no_fusion(%input : memref<10xf32>, %output : memref<10xf32>, %reduc : memref<1xf32>) {502 %zero = arith.constant 0. : f32503 %one = arith.constant 1. : f32504 // This producer can't be fused into inside %i without a violation of505 // semantics.506 // PRODUCER-CONSUMER-MAXIMAL: affine.for %{{.*}} = 0 to 10507 affine.for %i = 0 to 10 {508 %0 = affine.load %input[%i] : memref<10xf32>509 %1 = affine.load %reduc[0] : memref<1xf32>510 %2 = arith.addf %0, %1 : f32511 affine.store %2, %reduc[0] : memref<1xf32>512 }513 // PRODUCER-CONSUMER-MAXIMAL: affine.for %{{.*}} = 0 to 10514 affine.for %i = 0 to 10 {515 %0 = affine.load %reduc[0] : memref<1xf32>516 %2 = arith.addf %0, %one : f32517 affine.store %2, %output[%i] : memref<10xf32>518 }519 return520}521 522// SIBLING-MAXIMAL-LABEL: func @sibling_reduction523func.func @sibling_reduction(%input : memref<10xf32>, %output : memref<10xf32>, %reduc : memref<10xf32>) {524 %zero = arith.constant 0. : f32525 %one = arith.constant 1. : f32526 affine.for %i = 0 to 10 {527 %0 = affine.load %input[%i] : memref<10xf32>528 %2 = arith.addf %0, %one : f32529 affine.store %2, %output[%i] : memref<10xf32>530 }531 // Ensure that the fusion happens at the right depth.532 affine.for %i = 0 to 10 {533 %0 = affine.load %input[%i] : memref<10xf32>534 %1 = affine.load %reduc[0] : memref<10xf32>535 %2 = arith.addf %0, %1 : f32536 affine.store %2, %reduc[0] : memref<10xf32>537 }538 // SIBLING-MAXIMAL: affine.for %{{.*}} = 0 to 10539 // SIBLING-MAXIMAL-NEXT: affine.load540 // SIBLING-MAXIMAL-NEXT: addf541 // SIBLING-MAXIMAL-NEXT: affine.store542 // SIBLING-MAXIMAL-NEXT: affine.load543 // SIBLING-MAXIMAL-NEXT: affine.load544 // SIBLING-MAXIMAL-NEXT: addf545 // SIBLING-MAXIMAL-NEXT: affine.store546 return547}548 549// -----550 551// From https://github.com/llvm/llvm-project/issues/54541552 553#map = affine_map<(d0) -> (d0 mod 65536)>554// ZERO-TOLERANCE-LABEL: func @zero_tolerance555func.func @zero_tolerance(%arg0: memref<65536xcomplex<f64>>, %arg1: memref<30x131072xi64>,556%3 : memref<30xi64>,557%4 : memref<30xi64>,558%5 : memref<30xi64>,559%6 : memref<30xi64>560) {561 %c65536 = arith.constant 65536 : index562 %cst = arith.constant 0.000000e+00 : f64563 %cst_0 = arith.constant 0x4320000000380004 : f64564 %cst_1 = arith.constant 5.000000e-01 : f64565 %0 = memref.alloc() {alignment = 128 : i64} : memref<30x131072xi64>566 %1 = memref.alloc() {alignment = 128 : i64} : memref<131072xi1>567 %2 = memref.alloc() {alignment = 128 : i64} : memref<131072xi128>568 // This nest nest shouldn't be fused in when a zero tolerance is specified.569 // ZERO-TOLERANCE: affine.for %{{.*}} = 0 to 131072570 affine.for %arg2 = 0 to 131072 {571 %7 = affine.apply #map(%arg2)572 %8 = affine.load %arg0[%7] : memref<65536xcomplex<f64>>573 %9 = arith.cmpi ult, %arg2, %c65536 : index574 %10 = complex.im %8 : complex<f64>575 %11 = complex.re %8 : complex<f64>576 %12 = arith.select %9, %11, %10 : f64577 %13 = arith.cmpf olt, %12, %cst : f64578 %14 = arith.negf %12 : f64579 %15 = arith.select %13, %14, %12 : f64580 %16 = arith.mulf %15, %cst_0 : f64581 %17 = arith.addf %16, %cst_1 : f64582 %18 = arith.fptosi %17 : f64 to i128583 affine.store %18, %2[%arg2] : memref<131072xi128>584 affine.store %13, %1[%arg2] : memref<131072xi1>585 }586 // The next two nests are fused.587 // ZERO-TOLERANCE: affine.for %{{.*}} = 0 to 30588 // ZERO-TOLERANCE-NEXT: affine.for %{{.*}} = 0 to 131072589 // ZERO-TOLERANCE: func.call @__external_reduce_barrett590 // ZERO-TOLERANCE: affine.store591 // ZERO-TOLERANCE: affine.load592 // ZERO-TOLERANCE-NEXT: affine.store593 affine.for %arg2 = 0 to 30 {594 affine.for %arg3 = 0 to 131072 {595 %7 = affine.load %6[%arg2] : memref<30xi64>596 %8 = affine.load %3[%arg2] : memref<30xi64>597 %9 = affine.load %5[%arg2] : memref<30xi64>598 %10 = affine.load %4[%arg2] : memref<30xi64>599 %11 = affine.load %2[%arg3] : memref<131072xi128>600 %12 = affine.load %1[%arg3] : memref<131072xi1>601 %13 = func.call @__external_reduce_barrett(%7, %8, %9, %10, %11) {outputModFac = 1 : i64} : (i64, i64, i64, i64, i128) -> i64602 %14 = arith.subi %7, %13 : i64603 %15 = arith.select %12, %14, %13 : i64604 affine.store %15, %0[%arg2, %arg3] : memref<30x131072xi64>605 }606 }607 func.call @__external_levelwise_forward_ntt(%0) : (memref<30x131072xi64>) -> ()608 affine.for %arg2 = 0 to 30 {609 affine.for %arg3 = 0 to 131072 {610 %7 = affine.load %0[%arg2, %arg3] : memref<30x131072xi64>611 affine.store %7, %arg1[%arg2, %arg3] : memref<30x131072xi64>612 }613 }614 // Under maximal fusion, just one nest.615 // PRODUCER-CONSUMER-MAXIMAL: affine.for %{{.*}} = 0 to 30616 // PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for %{{.*}} = 0 to 131072617 // PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for %{{.*}}618 memref.dealloc %2 : memref<131072xi128>619 memref.dealloc %1 : memref<131072xi1>620 memref.dealloc %0 : memref<30x131072xi64>621 return622}623func.func private @__external_levelwise_forward_ntt(memref<30x131072xi64>)624func.func private @__external_reduce_barrett(i64, i64, i64, i64, i128) -> i64625 626// An unrolled loop nest. Fusion here should correctly fuse while preserving627// dependences between store-load pairs of the same memref. A private memref628// of size 1x1x1 can't be created.629 630// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @unrolled631func.func @unrolled(%arg0: memref<2x4xf32>, %arg1: memref<1x2x4xf32>) {632 %alloc = memref.alloc() : memref<1x2x4xf32>633 affine.for %i = 0 to 1 {634 %0 = affine.load %arg0[0, 0] : memref<2x4xf32>635 %1 = affine.load %arg0[0, 1] : memref<2x4xf32>636 %2 = affine.load %arg0[0, 2] : memref<2x4xf32>637 %3 = affine.load %arg0[0, 3] : memref<2x4xf32>638 %4 = affine.load %arg0[1, 0] : memref<2x4xf32>639 %5 = affine.load %arg0[1, 1] : memref<2x4xf32>640 %6 = affine.load %arg0[1, 2] : memref<2x4xf32>641 %7 = affine.load %arg0[1, 3] : memref<2x4xf32>642 643 affine.store %0, %alloc[0, 0, 0] : memref<1x2x4xf32>644 affine.store %1, %alloc[0, 0, 1] : memref<1x2x4xf32>645 affine.store %2, %alloc[0, 0, 2] : memref<1x2x4xf32>646 affine.store %3, %alloc[0, 0, 3] : memref<1x2x4xf32>647 affine.store %4, %alloc[0, 1, 0] : memref<1x2x4xf32>648 affine.store %5, %alloc[0, 1, 1] : memref<1x2x4xf32>649 affine.store %6, %alloc[0, 1, 2] : memref<1x2x4xf32>650 affine.store %7, %alloc[0, 1, 3] : memref<1x2x4xf32>651 }652 653 affine.for %i = 0 to 2 {654 affine.for %j = 0 to 4 {655 %8 = affine.load %alloc[0, %i, %j] : memref<1x2x4xf32>656 %9 = arith.negf %8 : f32657 affine.store %9, %arg1[0, %i, %j] : memref<1x2x4xf32>658 }659 }660 // PRODUCER-CONSUMER-MAXIMAL: affine.for %{{.*}} = 0 to 2 {661 // PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.for %{{.*}} = 0 to 4 {662 // PRODUCER-CONSUMER-MAXIMAL-NEXT: affine.load %{{.*}}[0, 0]663 // PRODUCER-CONSUMER-MAXIMAL: affine.load %{{.*}}[1, 3]664 // PRODUCER-CONSUMER-MAXIMAL: affine.store %{{.*}}[0, 0, 0]665 // PRODUCER-CONSUMER-MAXIMAL: affine.store %{{.*}}[0, 1, 3]666 // PRODUCER-CONSUMER-MAXIMAL: affine.load %{{.*}}[0, %{{.*}}, %{{.*}}]667 return668}669 670// -----671 672// Exercises fix for crash reported at https://github.com/llvm/llvm-project/issues/139231673 674#map = affine_map<(d0, d1) -> (d0 + d1)>675#map1 = affine_map<(d0, d1) -> (d0 * 2 + d1 * 2)>676module {677 func.func @zero_candidates() {678 %cst = arith.constant 2.221140e+03 : f32679 %cst_0 = arith.constant 2.606200e+03 : f32680 %cst_1 = arith.constant 3.224000e+03 : f32681 %cst_2 = arith.constant 0.000000e+00 : f32682 %alloc = memref.alloc() {alignment = 64 : i64} : memref<3x7x5x6xf32>683 affine.for %arg0 = 0 to 3 {684 affine.for %arg1 = 0 to 7 {685 affine.for %arg2 = 0 to 5 {686 affine.for %arg3 = 0 to 6 {687 affine.store %cst_1, %alloc[%arg0, %arg1, %arg2, %arg3] : memref<3x7x5x6xf32>688 }689 }690 }691 }692 %alloc_3 = memref.alloc() {alignment = 64 : i64} : memref<3x10x7x6xf32>693 %subview = memref.subview %alloc_3[0, 2, 1, 0] [3, 7, 5, 6] [1, 1, 1, 1] : memref<3x10x7x6xf32> to memref<3x7x5x6xf32, strided<[420, 42, 6, 1], offset: 90>>694 memref.copy %alloc, %subview : memref<3x7x5x6xf32> to memref<3x7x5x6xf32, strided<[420, 42, 6, 1], offset: 90>>695 %alloc_4 = memref.alloc() {alignment = 64 : i64} : memref<3x10x3x6x1xf32>696 affine.for %arg0 = 0 to 3 {697 affine.for %arg1 = 0 to 10 {698 affine.for %arg2 = 0 to 3 {699 affine.for %arg3 = 0 to 6 {700 affine.for %arg4 = 0 to 1 {701 affine.store %cst_2, %alloc_4[%arg0, %arg1, %arg2, %arg3, %arg4] : memref<3x10x3x6x1xf32>702 }703 }704 }705 }706 }707 affine.for %arg0 = 0 to 3 {708 affine.for %arg1 = 0 to 10 {709 affine.for %arg2 = 0 to 3 {710 affine.for %arg3 = 0 to 6 {711 affine.for %arg4 = 0 to 1 {712 affine.for %arg5 = 0 to 1 {713 affine.for %arg6 = 0 to 2 {714 %0 = affine.apply #map(%arg1, %arg5)715 %1 = affine.apply #map1(%arg2, %arg6)716 %2 = affine.load %alloc_3[%arg0, %0, %1, %arg3] : memref<3x10x7x6xf32>717 %3 = affine.load %alloc_4[%arg0, %arg1, %arg2, %arg3, %arg4] : memref<3x10x3x6x1xf32>718 %4 = arith.mulf %2, %cst_0 : f32719 %5 = arith.addf %3, %4 : f32720 affine.store %5, %alloc_4[%arg0, %arg1, %arg2, %arg3, %arg4] : memref<3x10x3x6x1xf32>721 }722 }723 }724 }725 }726 }727 }728 %alloc_5 = memref.alloc() {alignment = 64 : i64} : memref<3x10x3x6xf32>729 %expand_shape = memref.expand_shape %alloc_5 [[0], [1], [2], [3, 4]] output_shape [3, 10, 3, 6, 1] : memref<3x10x3x6xf32> into memref<3x10x3x6x1xf32>730 affine.for %arg0 = 0 to 3 {731 affine.for %arg1 = 0 to 10 {732 affine.for %arg2 = 0 to 3 {733 affine.for %arg3 = 0 to 6 {734 affine.for %arg4 = 0 to 1 {735 %0 = affine.load %alloc_4[%arg0, %arg1, %arg2, %arg3, %arg4] : memref<3x10x3x6x1xf32>736 %1 = arith.addf %0, %cst : f32737 affine.store %1, %expand_shape[%arg0, %arg1, %arg2, %arg3, %arg4] : memref<3x10x3x6x1xf32>738 }739 }740 }741 }742 }743 return744 }745}746 747// SIBLING-MAXIMAL-LABEL: memref_cast_reused748func.func @memref_cast_reused(%arg: memref<*xf32>) {749 %alloc = memref.cast %arg : memref<*xf32> to memref<10xf32>750 %alloc_0 = memref.alloc() : memref<10xf32>751 %alloc_1 = memref.alloc() : memref<10xf32>752 %cst = arith.constant 0.000000e+00 : f32753 %cst_2 = arith.constant 1.000000e+00 : f32754 affine.for %arg0 = 0 to 10 {755 %0 = affine.load %alloc[%arg0] : memref<10xf32>756 %1 = arith.addf %0, %cst_2 : f32757 affine.store %1, %alloc_0[%arg0] : memref<10xf32>758 }759 affine.for %arg0 = 0 to 10 {760 %0 = affine.load %alloc[%arg0] : memref<10xf32>761 %1 = affine.load %alloc_1[0] : memref<10xf32>762 %2 = arith.addf %0, %1 : f32763 affine.store %2, %alloc_1[0] : memref<10xf32>764 }765 // SIBLING-MAXIMAL: affine.for %{{.*}} = 0 to 10766 // SIBLING-MAXIMAL: addf767 // SIBLING-MAXIMAL-NEXT: affine.store768 // SIBLING-MAXIMAL-NEXT: affine.load769 // SIBLING-MAXIMAL-NEXT: affine.load770 // SIBLING-MAXIMAL-NEXT: addf771 // SIBLING-MAXIMAL-NEXT: affine.store772 return773}774