| // RUN: mlir-opt -pass-pipeline='builtin.module(func.func(affine-loop-fusion{mode=producer fusion-maximal}))' %s | FileCheck %s |
| |
| // Test fusion of affine nests inside other region-holding ops (scf.for in the |
| // test case below). |
| |
| // CHECK-LABEL: func @fusion_inner_simple |
| func.func @fusion_inner_simple(%A : memref<10xf32>) { |
| %cst = arith.constant 0.0 : f32 |
| |
| affine.for %i = 0 to 100 { |
| %B = memref.alloc() : memref<10xf32> |
| %C = memref.alloc() : memref<10xf32> |
| |
| affine.for %j = 0 to 10 { |
| %v = affine.load %A[%j] : memref<10xf32> |
| affine.store %v, %B[%j] : memref<10xf32> |
| } |
| |
| affine.for %j = 0 to 10 { |
| %v = affine.load %B[%j] : memref<10xf32> |
| affine.store %v, %C[%j] : memref<10xf32> |
| } |
| } |
| |
| // CHECK: affine.for %{{.*}} = 0 to 100 |
| // CHECK-NEXT: memref.alloc |
| // CHECK-NEXT: memref.alloc |
| // CHECK-NEXT: affine.for %{{.*}} = 0 to 10 |
| // CHECK-NOT: affine.for |
| |
| return |
| } |
| |
| // CHECK-LABEL: func @fusion_inner_simple_scf |
| func.func @fusion_inner_simple_scf(%A : memref<10xf32>) { |
| %c0 = arith.constant 0 : index |
| %c1 = arith.constant 1 : index |
| %c100 = arith.constant 100 : index |
| %cst = arith.constant 0.0 : f32 |
| |
| scf.for %i = %c0 to %c100 step %c1 { |
| %B = memref.alloc() : memref<10xf32> |
| %C = memref.alloc() : memref<10xf32> |
| |
| affine.for %j = 0 to 10 { |
| %v = affine.load %A[%j] : memref<10xf32> |
| affine.store %v, %B[%j] : memref<10xf32> |
| } |
| |
| affine.for %j = 0 to 10 { |
| %v = affine.load %B[%j] : memref<10xf32> |
| affine.store %v, %C[%j] : memref<10xf32> |
| } |
| } |
| // CHECK: scf.for |
| // CHECK-NEXT: memref.alloc |
| // CHECK-NEXT: memref.alloc |
| // CHECK-NEXT: affine.for %{{.*}} = 0 to 10 |
| // CHECK-NOT: affine.for |
| return |
| } |
| |
| // CHECK-LABEL: func @fusion_inner_multiple_nests |
| func.func @fusion_inner_multiple_nests() { |
| %alloc_5 = memref.alloc() {alignment = 64 : i64} : memref<4x4xi8> |
| %alloc_10 = memref.alloc() : memref<8x4xi32> |
| affine.for %arg8 = 0 to 4 { |
| %alloc_14 = memref.alloc() : memref<4xi8> |
| %alloc_15 = memref.alloc() : memref<8x4xi8> |
| affine.for %arg9 = 0 to 4 { |
| %0 = affine.load %alloc_5[%arg9, %arg8] : memref<4x4xi8> |
| affine.store %0, %alloc_14[%arg9] : memref<4xi8> |
| } |
| %alloc_16 = memref.alloc() : memref<4xi8> |
| affine.for %arg9 = 0 to 4 { |
| %0 = affine.load %alloc_14[%arg9] : memref<4xi8> |
| affine.store %0, %alloc_16[%arg9] : memref<4xi8> |
| } |
| affine.for %arg9 = 0 to 2 { |
| %0 = affine.load %alloc_15[%arg9 * 4, 0] : memref<8x4xi8> |
| %1 = affine.load %alloc_16[0] : memref<4xi8> |
| %2 = affine.load %alloc_10[%arg9 * 4, %arg8] : memref<8x4xi32> |
| %3 = arith.muli %0, %1 : i8 |
| %4 = arith.extsi %3 : i8 to i32 |
| %5 = arith.addi %4, %2 : i32 |
| affine.store %5, %alloc_10[%arg9 * 4 + 3, %arg8] : memref<8x4xi32> |
| } |
| memref.dealloc %alloc_16 : memref<4xi8> |
| } |
| // CHECK: affine.for %{{.*}} = 0 to 4 { |
| // Everything inside fused into two nests (the second will be DCE'd). |
| // CHECK-NEXT: memref.alloc() : memref<4xi8> |
| // CHECK-NEXT: memref.alloc() : memref<1xi8> |
| // CHECK-NEXT: memref.alloc() : memref<1xi8> |
| // CHECK-NEXT: memref.alloc() : memref<8x4xi8> |
| // CHECK-NEXT: memref.alloc() : memref<4xi8> |
| // CHECK-NEXT: affine.for %{{.*}} = 0 to 2 { |
| // CHECK: } |
| // CHECK: affine.for %{{.*}} = 0 to 4 { |
| // CHECK: } |
| // CHECK-NEXT: memref.dealloc |
| // CHECK-NEXT: } |
| // CHECK-NEXT: return |
| return |
| } |
| |
| // CHECK-LABEL: func @fusion_inside_scf_while |
| func.func @fusion_inside_scf_while(%A : memref<10xf32>) { |
| %c0 = arith.constant 0 : index |
| %c1 = arith.constant 1 : index |
| %c100 = arith.constant 100 : index |
| %cst = arith.constant 0.0 : f32 |
| |
| %0 = scf.while (%arg3 = %cst) : (f32) -> (f32) { |
| %1 = arith.cmpf ult, %arg3, %cst : f32 |
| scf.condition(%1) %arg3 : f32 |
| } do { |
| ^bb0(%arg5: f32): |
| |
| %B = memref.alloc() : memref<10xf32> |
| %C = memref.alloc() : memref<10xf32> |
| |
| affine.for %j = 0 to 10 { |
| %v = affine.load %A[%j] : memref<10xf32> |
| affine.store %v, %B[%j] : memref<10xf32> |
| } |
| |
| affine.for %j = 0 to 10 { |
| %v = affine.load %B[%j] : memref<10xf32> |
| affine.store %v, %C[%j] : memref<10xf32> |
| } |
| %1 = arith.mulf %arg5, %cst : f32 |
| scf.yield %1 : f32 |
| } |
| // CHECK: scf.while |
| // CHECK: affine.for %{{.*}} = 0 to 10 |
| // CHECK-NOT: affine.for |
| // CHECK: scf.yield |
| return |
| } |
| |
| |
| memref.global "private" constant @__constant_10x2xf32 : memref<10x2xf32> = dense<0.000000e+00> |
| |
| // CHECK-LABEL: func @fusion_inner_long |
| func.func @fusion_inner_long(%arg0: memref<10x2xf32>, %arg1: memref<10x10xf32>, %arg2: memref<10x2xf32>, %s: index) { |
| %c0 = arith.constant 0 : index |
| %cst_0 = arith.constant 1.000000e-03 : f32 |
| %c9 = arith.constant 9 : index |
| %c10_i32 = arith.constant 10 : i32 |
| %c1_i32 = arith.constant 1 : i32 |
| %c100_i32 = arith.constant 100 : i32 |
| %c0_i32 = arith.constant 0 : i32 |
| %0 = memref.get_global @__constant_10x2xf32 : memref<10x2xf32> |
| %1 = scf.for %arg3 = %c0_i32 to %c100_i32 step %c1_i32 iter_args(%arg4 = %arg0) -> (memref<10x2xf32>) : i32 { |
| %alloc = memref.alloc() {alignment = 64 : i64} : memref<10xi32> |
| affine.for %arg5 = 0 to 10 { |
| %3 = arith.index_cast %arg5 : index to i32 |
| affine.store %3, %alloc[%arg5] : memref<10xi32> |
| } |
| %2 = scf.for %arg5 = %c0_i32 to %c10_i32 step %c1_i32 iter_args(%arg6 = %0) -> (memref<10x2xf32>) : i32 { |
| %alloc_5 = memref.alloc() : memref<2xf32> |
| affine.for %arg7 = 0 to 2 { |
| %16 = affine.load %arg4[%s, %arg7] : memref<10x2xf32> |
| affine.store %16, %alloc_5[%arg7] : memref<2xf32> |
| } |
| %alloc_6 = memref.alloc() {alignment = 64 : i64} : memref<1x2xf32> |
| affine.for %arg7 = 0 to 2 { |
| %16 = affine.load %alloc_5[%arg7] : memref<2xf32> |
| affine.store %16, %alloc_6[0, %arg7] : memref<1x2xf32> |
| } |
| %alloc_7 = memref.alloc() {alignment = 64 : i64} : memref<10x2xf32> |
| affine.for %arg7 = 0 to 10 { |
| affine.for %arg8 = 0 to 2 { |
| %16 = affine.load %alloc_6[0, %arg8] : memref<1x2xf32> |
| affine.store %16, %alloc_7[%arg7, %arg8] : memref<10x2xf32> |
| } |
| } |
| %alloc_8 = memref.alloc() {alignment = 64 : i64} : memref<10x2xf32> |
| affine.for %arg7 = 0 to 10 { |
| affine.for %arg8 = 0 to 2 { |
| %16 = affine.load %alloc_7[%arg7, %arg8] : memref<10x2xf32> |
| %17 = affine.load %arg4[%arg7, %arg8] : memref<10x2xf32> |
| %18 = arith.subf %16, %17 : f32 |
| affine.store %18, %alloc_8[%arg7, %arg8] : memref<10x2xf32> |
| } |
| } |
| scf.yield %alloc_8 : memref<10x2xf32> |
| // CHECK: scf.for |
| // CHECK: scf.for |
| // CHECK: affine.for %{{.*}} = 0 to 10 |
| // CHECK-NEXT: affine.for %{{.*}} = 0 to 2 |
| // CHECK-NOT: affine.for |
| // CHECK: scf.yield |
| } |
| %alloc_2 = memref.alloc() {alignment = 64 : i64} : memref<10x2xf32> |
| affine.for %arg5 = 0 to 10 { |
| affine.for %arg6 = 0 to 2 { |
| affine.store %cst_0, %alloc_2[%arg5, %arg6] : memref<10x2xf32> |
| } |
| } |
| %alloc_3 = memref.alloc() {alignment = 64 : i64} : memref<10x2xf32> |
| affine.for %arg5 = 0 to 10 { |
| affine.for %arg6 = 0 to 2 { |
| %3 = affine.load %alloc_2[%arg5, %arg6] : memref<10x2xf32> |
| %4 = affine.load %2[%arg5, %arg6] : memref<10x2xf32> |
| %5 = arith.mulf %3, %4 : f32 |
| affine.store %5, %alloc_3[%arg5, %arg6] : memref<10x2xf32> |
| } |
| } |
| scf.yield %alloc_3 : memref<10x2xf32> |
| // The nests above will be fused as well. |
| // CHECK: affine.for %{{.*}} = 0 to 10 |
| // CHECK-NEXT: affine.for %{{.*}} = 0 to 2 |
| // CHECK-NOT: affine.for |
| // CHECK: scf.yield |
| } |
| affine.for %arg3 = 0 to 10 { |
| affine.for %arg4 = 0 to 2 { |
| %2 = affine.load %1[%arg3, %arg4] : memref<10x2xf32> |
| affine.store %2, %arg2[%arg3, %arg4] : memref<10x2xf32> |
| } |
| } |
| return |
| } |
