Skip to content

Commit bb65dfc

Browse files
HsiangkaiHsiangkai
committed
[mlir][scf] Considering affine.apply when fusing scf::ParallelOp
When checking the load indices of the second loop coincide with the store indices of the first loop, it only considers the index values are the same or not. However, there are some cases the index values come from affine.apply operator. In these cases, it will treat them as different even the affine map is the same and the affine operands are the same. We already check if the iteration space is the same in isFusionLegal(). When checking affine.apply, we only need to consider the operands come from the same induction variables. If so, we know the results of affine.apply are the same.
1 parent 96a3d05 commit bb65dfc

File tree

2 files changed

+120
-2
lines changed

2 files changed

+120
-2
lines changed

‎mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
#include "mlir/IR/Builders.h"
2020
#include "mlir/IR/IRMapping.h"
2121
#include "mlir/IR/OpDefinition.h"
22+
#include "mlir/IR/OperationSupport.h"
2223
#include "mlir/Interfaces/SideEffectInterfaces.h"
2324

2425
namespace mlir {
@@ -102,8 +103,30 @@ static bool haveNoReadsAfterWriteExceptSameIndex(
102103
return WalkResult::interrupt();
103104
for (int i = 0, e = storeIndices.size(); i < e; ++i) {
104105
if (firstToSecondPloopIndices.lookupOrDefault(storeIndices[i]) !=
105-
loadIndices[i])
106-
return WalkResult::interrupt();
106+
loadIndices[i]) {
107+
auto *storeIndexDefOp = storeIndices[i].getDefiningOp();
108+
auto *loadIndexDefOp = loadIndices[i].getDefiningOp();
109+
if (storeIndexDefOp && loadIndexDefOp) {
110+
if (!isMemoryEffectFree(storeIndexDefOp))
111+
return WalkResult::interrupt();
112+
if (!isMemoryEffectFree(loadIndexDefOp))
113+
return WalkResult::interrupt();
114+
if (!OperationEquivalence::isEquivalentTo(
115+
storeIndexDefOp, loadIndexDefOp,
116+
[&](Value storeIndex, Value loadIndex) {
117+
if (firstToSecondPloopIndices.lookupOrDefault(storeIndex) !=
118+
firstToSecondPloopIndices.lookupOrDefault(loadIndex))
119+
return failure();
120+
else
121+
return success();
122+
},
123+
/*markEquivalent=*/nullptr,
124+
OperationEquivalence::Flags::IgnoreLocations)) {
125+
return WalkResult::interrupt();
126+
}
127+
} else
128+
return WalkResult::interrupt();
129+
}
107130
}
108131
return WalkResult::advance();
109132
});

‎mlir/test/Dialect/SCF/parallel-loop-fusion.mlir

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -480,3 +480,98 @@ func.func @do_not_fuse_multiple_stores_on_diff_indices(
480480
// CHECK: scf.reduce
481481
// CHECK: }
482482
// CHECK: memref.dealloc [[SUM]]
483+
484+
// -----
485+
486+
func.func @fuse_same_indices_by_affine_apply(
487+
%A: memref<2x2xf32>, %B: memref<2x2xf32>) {
488+
%c0 = arith.constant 0 : index
489+
%c1 = arith.constant 1 : index
490+
%c2 = arith.constant 2 : index
491+
%sum = memref.alloc() : memref<2x3xf32>
492+
scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
493+
%B_elem = memref.load %B[%i, %j] : memref<2x2xf32>
494+
%1 = affine.apply affine_map<(d0, d1) -> (d0 + d1)>(%i, %j)
495+
memref.store %B_elem, %sum[%i, %1] : memref<2x3xf32>
496+
scf.reduce
497+
}
498+
scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
499+
%1 = affine.apply affine_map<(d0, d1) -> (d0 + d1)>(%i, %j)
500+
%sum_elem = memref.load %sum[%i, %1] : memref<2x3xf32>
501+
%A_elem = memref.load %A[%i, %j] : memref<2x2xf32>
502+
%product = arith.mulf %sum_elem, %A_elem : f32
503+
memref.store %product, %B[%i, %j] : memref<2x2xf32>
504+
scf.reduce
505+
}
506+
memref.dealloc %sum : memref<2x3xf32>
507+
return
508+
}
509+
// CHECK: #[[$MAP:.*]] = affine_map<(d0, d1) -> (d0 + d1)>
510+
// CHECK-LABEL: fuse_same_indices_by_affine_apply
511+
// CHECK-SAME: (%[[ARG0:.*]]: memref<2x2xf32>, %[[ARG1:.*]]: memref<2x2xf32>) {
512+
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
513+
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
514+
// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
515+
// CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<2x3xf32>
516+
// CHECK-NEXT: scf.parallel (%[[ARG2:.*]], %[[ARG3:.*]]) = (%[[C0]], %[[C0]]) to (%[[C2]], %[[C2]]) step (%[[C1]], %[[C1]]) {
517+
// CHECK-NEXT: %[[S0:.*]] = memref.load %[[ARG1]][%[[ARG2]], %[[ARG3]]] : memref<2x2xf32>
518+
// CHECK-NEXT: %[[S1:.*]] = affine.apply #[[$MAP]](%[[ARG2]], %[[ARG3]])
519+
// CHECK-NEXT: memref.store %[[S0]], %[[ALLOC]][%[[ARG2]], %[[S1]]] : memref<2x3xf32>
520+
// CHECK-NEXT: %[[S2:.*]] = affine.apply #[[$MAP]](%[[ARG2]], %[[ARG3]])
521+
// CHECK-NEXT: %[[S3:.*]] = memref.load %[[ALLOC]][%[[ARG2]], %[[S2]]] : memref<2x3xf32>
522+
// CHECK-NEXT: %[[S4:.*]] = memref.load %[[ARG0]][%[[ARG2]], %[[ARG3]]] : memref<2x2xf32>
523+
// CHECK-NEXT: %[[S5:.*]] = arith.mulf %[[S3]], %[[S4]] : f32
524+
// CHECK-NEXT: memref.store %[[S5]], %[[ARG1]][%[[ARG2]], %[[ARG3]]] : memref<2x2xf32>
525+
// CHECK-NEXT: scf.reduce
526+
// CHECK-NEXT: }
527+
// CHECK-NEXT: memref.dealloc %[[ALLOC]] : memref<2x3xf32>
528+
// CHECK-NEXT: return
529+
530+
// -----
531+
532+
func.func @do_not_fuse_affine_apply_to_non_ind_var(
533+
%A: memref<2x2xf32>, %B: memref<2x2xf32>, %OffsetA: index, %OffsetB: index) {
534+
%c0 = arith.constant 0 : index
535+
%c1 = arith.constant 1 : index
536+
%c2 = arith.constant 2 : index
537+
%sum = memref.alloc() : memref<2x3xf32>
538+
scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
539+
%B_elem = memref.load %B[%i, %j] : memref<2x2xf32>
540+
%1 = affine.apply affine_map<(d0, d1) -> (d0 + d1)>(%i, %OffsetA)
541+
memref.store %B_elem, %sum[%i, %1] : memref<2x3xf32>
542+
scf.reduce
543+
}
544+
scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
545+
%1 = affine.apply affine_map<(d0, d1) -> (d0 + d1)>(%i, %OffsetB)
546+
%sum_elem = memref.load %sum[%i, %1] : memref<2x3xf32>
547+
%A_elem = memref.load %A[%i, %j] : memref<2x2xf32>
548+
%product = arith.mulf %sum_elem, %A_elem : f32
549+
memref.store %product, %B[%i, %j] : memref<2x2xf32>
550+
scf.reduce
551+
}
552+
memref.dealloc %sum : memref<2x3xf32>
553+
return
554+
}
555+
// CHECK: #[[$MAP:.*]] = affine_map<(d0, d1) -> (d0 + d1)>
556+
// CHECK-LABEL: do_not_fuse_affine_apply_to_non_ind_var
557+
// CHECK-SAME: (%[[ARG0:.*]]: memref<2x2xf32>, %[[ARG1:.*]]: memref<2x2xf32>, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index) {
558+
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
559+
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
560+
// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
561+
// CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<2x3xf32>
562+
// CHECK-NEXT: scf.parallel (%[[ARG4:.*]], %[[ARG5:.*]]) = (%[[C0]], %[[C0]]) to (%[[C2]], %[[C2]]) step (%[[C1]], %[[C1]]) {
563+
// CHECK-NEXT: %[[S0:.*]] = memref.load %[[ARG1]][%[[ARG4]], %[[ARG5]]] : memref<2x2xf32>
564+
// CHECK-NEXT: %[[S1:.*]] = affine.apply #[[$MAP]](%[[ARG4]], %[[ARG2]])
565+
// CHECK-NEXT: memref.store %[[S0]], %[[ALLOC]][%[[ARG4]], %[[S1]]] : memref<2x3xf32>
566+
// CHECK-NEXT: scf.reduce
567+
// CHECK-NEXT: }
568+
// CHECK-NEXT: scf.parallel (%[[ARG4:.*]], %[[ARG5:.*]]) = (%[[C0]], %[[C0]]) to (%[[C2]], %[[C2]]) step (%[[C1]], %[[C1]]) {
569+
// CHECK-NEXT: %[[S0:.*]] = affine.apply #[[$MAP]](%[[ARG4]], %[[ARG3]])
570+
// CHECK-NEXT: %[[S1:.*]] = memref.load %[[ALLOC]][%[[ARG4]], %[[S0]]] : memref<2x3xf32>
571+
// CHECK-NEXT: %[[S2:.*]] = memref.load %[[ARG0]][%[[ARG4]], %[[ARG5]]] : memref<2x2xf32>
572+
// CHECK-NEXT: %[[S3:.*]] = arith.mulf %[[S1]], %[[S2]] : f32
573+
// CHECK-NEXT: memref.store %[[S3]], %[[ARG1]][%[[ARG4]], %[[ARG5]]] : memref<2x2xf32>
574+
// CHECK-NEXT: scf.reduce
575+
// CHECK-NEXT: }
576+
// CHECK-NEXT: memref.dealloc %[[ALLOC]] : memref<2x3xf32>
577+
// CHECK-NEXT: return

0 commit comments

Comments
 (0)