xref: /llvm-project/mlir/test/Examples/mlir-opt/loop_fusion_default.mlir (revision 7f1968625a607fb49a2b9a67f3c8fb2892cf4839)
1// RUN: mlir-opt --pass-pipeline="builtin.module(affine-loop-fusion)" %s | FileCheck %s
2
3
4// CHECK-LABEL: @producer_consumer_fusion
5// CHECK-COUNT-1: affine.for
6module {
7  func.func @producer_consumer_fusion(%arg0: memref<10xf32>, %arg1: memref<10xf32>) {
8    %0 = memref.alloc() : memref<10xf32>
9    %1 = memref.alloc() : memref<10xf32>
10    %cst = arith.constant 0.000000e+00 : f32
11    affine.for %arg2 = 0 to 10 {
12      affine.store %cst, %0[%arg2] : memref<10xf32>
13      affine.store %cst, %1[%arg2] : memref<10xf32>
14    }
15    affine.for %arg2 = 0 to 10 {
16      %2 = affine.load %0[%arg2] : memref<10xf32>
17      %3 = arith.addf %2, %2 : f32
18      affine.store %3, %arg0[%arg2] : memref<10xf32>
19    }
20    affine.for %arg2 = 0 to 10 {
21      %2 = affine.load %1[%arg2] : memref<10xf32>
22      %3 = arith.mulf %2, %2 : f32
23      affine.store %3, %arg1[%arg2] : memref<10xf32>
24    }
25    return
26  }
27}
28