Skip to content

Commit e245656

Browse files
stevenvarstevenvar
committed
[mlir][Hoisting] Hoisting vector.extract/vector.broadcast pairs
1 parent 2d14ea6 commit e245656

File tree

5 files changed

+275
-0
lines changed

5 files changed

+275
-0
lines changed

‎mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2215,6 +2215,42 @@ def HoistRedundantVectorTransfersOp :
22152215
}];
22162216
}
22172217

2218+
//===----------------------------------------------------------------------===//
2219+
// HoistRedundantVectorBroadcastsOp
2220+
//===----------------------------------------------------------------------===//
2221+
2222+
def HoistRedundantVectorBroadcastsOp :
2223+
Op<Transform_Dialect, "structured.hoist_redundant_vector_broadcasts",
2224+
[FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
2225+
TransformEachOpTrait, TransformOpInterface,
2226+
ReportTrackingListenerFailuresOpTrait]> {
2227+
let description = [{
2228+
Hoist vector.extract / vector.broadcasts pairs out of immediately
2229+
enclosing scf::ForOp iteratively.
2230+
2231+
#### Return modes:
2232+
2233+
The operation always succeeds and returns a handle to the transformed
2234+
function op.
2235+
}];
2236+
2237+
let arguments = (ins TransformHandleTypeInterface:$target);
2238+
let results = (outs TransformHandleTypeInterface:$transformed);
2239+
2240+
let assemblyFormat = "$target attr-dict `:` functional-type(operands, results) ";
2241+
2242+
let builders = [
2243+
OpBuilder<(ins "Value":$target)>,
2244+
];
2245+
let extraClassDeclaration = [{
2246+
::mlir::DiagnosedSilenceableFailure applyToOne(
2247+
::mlir::transform::TransformRewriter &rewriter,
2248+
::mlir::func::FuncOp target,
2249+
::mlir::transform::ApplyToEachResultList &results,
2250+
::mlir::transform::TransformState &state);
2251+
}];
2252+
}
2253+
22182254
//===----------------------------------------------------------------------===//
22192255
// ConvertConv2DToImg2ColOp
22202256
//===----------------------------------------------------------------------===//

‎mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,8 @@ namespace linalg {
4343
/// when used on distributed loops with memref semantics!
4444
void hoistRedundantVectorTransfers(Operation *root);
4545

46+
void hoistRedundantVectorBroadcasts(Operation *root);
47+
4648
} // namespace linalg
4749
} // namespace mlir
4850

‎mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3224,6 +3224,20 @@ transform::HoistRedundantVectorTransfersOp::applyToOne(
32243224
return DiagnosedSilenceableFailure::success();
32253225
}
32263226

3227+
//===----------------------------------------------------------------------===//
3228+
// HoistRedundantVectorBroadcastsOp
3229+
//===----------------------------------------------------------------------===//
3230+
3231+
DiagnosedSilenceableFailure
3232+
transform::HoistRedundantVectorBroadcastsOp::applyToOne(
3233+
transform::TransformRewriter &rewriter, func::FuncOp target,
3234+
transform::ApplyToEachResultList &results,
3235+
transform::TransformState &state) {
3236+
linalg::hoistRedundantVectorBroadcasts(target);
3237+
results.push_back(target);
3238+
return DiagnosedSilenceableFailure::success();
3239+
}
3240+
32273241
//===----------------------------------------------------------------------===//
32283242
// ConvertConv2DToImg2ColOp.
32293243
//===----------------------------------------------------------------------===//

‎mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,120 @@ using llvm::dbgs;
4343
using namespace mlir;
4444
using namespace mlir::linalg;
4545

46+
scf::ForOp replaceWithDifferentYield(RewriterBase &rewriter, scf::ForOp loop,
47+
Value newInitOperand, int index,
48+
Value newYieldValue) {
49+
OpBuilder::InsertionGuard g(rewriter);
50+
rewriter.setInsertionPoint(loop.getOperation());
51+
auto inits = llvm::to_vector(loop.getInits());
52+
53+
// Replace the init value with the new operand
54+
inits[index] = newInitOperand;
55+
56+
scf::ForOp newLoop = rewriter.create<scf::ForOp>(
57+
loop.getLoc(), loop.getLowerBound(), loop.getUpperBound(), loop.getStep(),
58+
inits, [](OpBuilder &, Location, Value, ValueRange) {});
59+
60+
// Generate the new yield with the replaced operand
61+
auto yieldOp = cast<scf::YieldOp>(loop.getBody()->getTerminator());
62+
yieldOp->getOperand(index).replaceAllUsesWith(newYieldValue);
63+
64+
// Move the loop body to the new op.
65+
rewriter.mergeBlocks(loop.getBody(), newLoop.getBody(),
66+
newLoop.getBody()->getArguments().take_front(
67+
loop.getBody()->getNumArguments()));
68+
69+
// Replace the old loop.
70+
rewriter.replaceOp(loop.getOperation(),
71+
newLoop->getResults().take_front(loop.getNumResults()));
72+
return newLoop;
73+
}
74+
75+
// Hoist out a pair of corresponding vector.extract+vector.broadcast
76+
// operations. This function transforms a loop like this:
77+
// %loop = scf.for _ = _ to _ step _ iter_args(%iterarg = %v) -> (t1) {
78+
// %e = vector.extract %iterarg : t1 to t2
79+
// %u = // do something with %e : t2
80+
// %b = vector.broadcast %u : t2 to t1
81+
// scf.yield %b : t1
82+
// }
83+
// into the following:
84+
// %e = vector.extract %v: t1 to t2
85+
// %loop' = scf.for _ = _ to _ step _ iter_args(%iterarg = %e) -> (t2) {
86+
// %u' = // do something with %iterarg : t2
87+
// scf.yield %u' : t2
88+
// }
89+
// %loop = vector.broadcast %loop' : t2 to t1
90+
void mlir::linalg::hoistRedundantVectorBroadcasts(Operation *root) {
91+
bool changed = true;
92+
while (changed) {
93+
changed = false;
94+
// First move loop invariant ops outside of their loop. This needs to be
95+
// done before as we cannot move ops without interrupting the function walk.
96+
root->walk(
97+
[&](LoopLikeOpInterface loopLike) { moveLoopInvariantCode(loopLike); });
98+
99+
root->walk([&](vector::ExtractOp extractOp) {
100+
LLVM_DEBUG(DBGS() << "Candidate for hoisting: "
101+
<< *extractOp.getOperation() << "\n");
102+
103+
auto loop = dyn_cast<scf::ForOp>(extractOp->getParentOp());
104+
if (!loop)
105+
return WalkResult::advance();
106+
107+
// Check that the vector to extract from is an iter_arg
108+
auto blockArg = dyn_cast<BlockArgument>(extractOp.getVector());
109+
if (!blockArg)
110+
return WalkResult::advance();
111+
112+
// If the iter_arg does not have only one use, it won't be possible to
113+
// hoist the extractOp out.
114+
if (!blockArg.hasOneUse())
115+
return WalkResult::advance();
116+
117+
auto initArg = loop.getTiedLoopInit(blockArg)->get();
118+
auto index = blockArg.getArgNumber() - loop.getNumInductionVars();
119+
120+
// Check that the loop yields a broadcast
121+
auto lastOp = loop.getBody()->getTerminator();
122+
auto yieldOp = dyn_cast<scf::YieldOp>(lastOp);
123+
if (!yieldOp)
124+
return WalkResult::advance();
125+
126+
auto broadcast = dyn_cast<vector::BroadcastOp>(
127+
yieldOp->getOperand(index).getDefiningOp());
128+
129+
LLVM_DEBUG(DBGS() << "Candidate broadcast: " << broadcast << "\n");
130+
131+
Type broadcastInputType = broadcast.getSourceType();
132+
if (broadcastInputType != extractOp.getType())
133+
return WalkResult::advance();
134+
135+
// The position of the extract must be defined outside of the loop if
136+
// it is dynamic
137+
for (auto operand : extractOp.getDynamicPosition())
138+
if (!loop.isDefinedOutsideOfLoop(operand))
139+
return WalkResult::advance();
140+
141+
extractOp.getVectorMutable().assign(initArg);
142+
loop.moveOutOfLoop(extractOp);
143+
broadcast->moveAfter(loop);
144+
145+
IRRewriter rewriter(extractOp.getContext());
146+
auto newLoop = replaceWithDifferentYield(
147+
rewriter, loop, extractOp.getResult(), index, broadcast.getSource());
148+
149+
LLVM_DEBUG(DBGS() << "New loop: " << newLoop << "\n");
150+
151+
newLoop.getResult(index).replaceAllUsesWith(broadcast);
152+
broadcast.getSourceMutable().assign(newLoop.getResult(index));
153+
154+
changed = true;
155+
return WalkResult::interrupt();
156+
});
157+
}
158+
}
159+
46160
static bool noAliasingUseInLoop(vector::TransferReadOp transferRead,
47161
LoopLikeOpInterface loop) {
48162
Value source = transferRead.getSource();

‎mlir/test/Dialect/Linalg/hoisting.mlir

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -565,3 +565,112 @@ module attributes {transform.with_named_sequence} {
565565
transform.yield
566566
}
567567
}
568+
569+
// -----
570+
571+
// Test hoisting of vector.extract/vector.broadcast pairs
572+
573+
// CHECK-LABEL: func.func @hoist_vector_broadcasts
574+
// CHECK-SAME: (%{{.+}}: index, %{{.+}}: index, %{{.+}}: index, %[[VEC:.+]]: vector<3x4xf32>) -> vector<3x4xf32> {
575+
// CHECK: %[[EXTRACT:.+]] = vector.extract %[[VEC]][0] : vector<4xf32> from vector<3x4xf32>
576+
// CHECK-NEXT: %[[LOOP:.+]] = scf.for {{.*}} {
577+
// CHECK-NEXT: %[[USE:.+]] = "some_use"({{.*}}) : (vector<4xf32>) -> vector<4xf32>
578+
// CHECK-NEXT: scf.yield %[[USE]] : vector<4xf32>
579+
// CHECK-NEXT: }
580+
// CHECK-NEXT: %[[BCAST:.+]] = vector.broadcast %[[LOOP]] : vector<4xf32> to vector<3x4xf32>
581+
// CHECK-NEXT: return %[[BCAST]] : vector<3x4xf32>
582+
583+
func.func @hoist_vector_broadcasts(%lb : index, %ub : index, %step : index, %vec : vector<3x4xf32>) -> vector<3x4xf32> {
584+
%bcast_vec = scf.for %arg0 = %lb to %ub step %step iter_args(%iarg = %vec) -> vector<3x4xf32> {
585+
%extract = vector.extract %iarg[0] : vector<4xf32> from vector<3x4xf32>
586+
%use = "some_use"(%extract) : (vector<4xf32>) -> vector<4xf32>
587+
%broadcast = vector.broadcast %use : vector<4xf32> to vector<3x4xf32>
588+
scf.yield %broadcast : vector<3x4xf32>
589+
}
590+
return %bcast_vec : vector<3x4xf32>
591+
}
592+
593+
module attributes {transform.with_named_sequence} {
594+
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
595+
%0 = transform.structured.match ops{["func.func"]} in %arg1
596+
: (!transform.any_op) -> !transform.any_op
597+
transform.structured.hoist_redundant_vector_broadcasts %0
598+
: (!transform.any_op) -> !transform.any_op
599+
transform.yield
600+
}
601+
}
602+
603+
// -----
604+
605+
// Test hoisting of vector.extract/vector.broadcast pairs with dynamic position
606+
607+
// CHECK-LABEL: func.func @hoist_vector_broadcasts
608+
// CHECK-SAME: (%{{.+}}: index, %{{.+}}: index, %{{.+}}: index, %[[VEC:.+]]: vector<3x4xf32>, %[[POS:.+]]: index) -> vector<3x4xf32> {
609+
// CHECK: %[[EXTRACT:.+]] = vector.extract %[[VEC]][%[[POS]]] : vector<4xf32> from vector<3x4xf32>
610+
// CHECK-NEXT: %[[LOOP:.+]] = scf.for {{.*}} {
611+
// CHECK-NEXT: %[[USE:.+]] = "some_use"({{.*}}) : (vector<4xf32>) -> vector<4xf32>
612+
// CHECK-NEXT: scf.yield %[[USE]] : vector<4xf32>
613+
// CHECK-NEXT: }
614+
// CHECK-NEXT: %[[BCAST:.+]] = vector.broadcast %[[LOOP]] : vector<4xf32> to vector<3x4xf32>
615+
// CHECK-NEXT: return %[[BCAST]] : vector<3x4xf32>
616+
617+
func.func @hoist_vector_broadcasts_dynamic(%lb : index, %ub : index, %step : index, %vec : vector<3x4xf32>, %pos: index) -> vector<3x4xf32> {
618+
%bcast_vec = scf.for %arg0 = %lb to %ub step %step iter_args(%iarg = %vec) -> vector<3x4xf32> {
619+
%extract = vector.extract %iarg[%pos] : vector<4xf32> from vector<3x4xf32>
620+
%use = "some_use"(%extract) : (vector<4xf32>) -> vector<4xf32>
621+
%broadcast = vector.broadcast %use : vector<4xf32> to vector<3x4xf32>
622+
scf.yield %broadcast : vector<3x4xf32>
623+
}
624+
return %bcast_vec : vector<3x4xf32>
625+
}
626+
627+
module attributes {transform.with_named_sequence} {
628+
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
629+
%0 = transform.structured.match ops{["func.func"]} in %arg1
630+
: (!transform.any_op) -> !transform.any_op
631+
transform.structured.hoist_redundant_vector_broadcasts %0
632+
: (!transform.any_op) -> !transform.any_op
633+
transform.yield
634+
}
635+
}
636+
637+
// -----
638+
639+
// Test hoisting of vector.extract/vector.broadcast pairs with multiple iter_args
640+
641+
// CHECK-LABEL: func.func @hoist_vector_broadcasts_multiple
642+
// CHECK-SAME: (%{{.+}}: index, %{{.+}}: index, %{{.+}}: index, %[[VEC1:.+]]: vector<3x4xf32>,
643+
// CHECK-SAME: %[[VEC2:.+]]: vector<3x5xf32>) -> (vector<3x4xf32>, vector<3x5xf32>) {
644+
// CHECK-DAG: %[[EXTRACT1:.+]] = vector.extract %[[VEC1]][0] : vector<4xf32> from vector<3x4xf32>
645+
// CHECK-DAG: %[[EXTRACT2:.+]] = vector.extract %[[VEC2]][1] : vector<5xf32> from vector<3x5xf32>
646+
// CHECK-NEXT: %[[LOOP:.+]]:2 = scf.for {{.*}} {
647+
// CHECK-DAG: %[[USE1:.+]] = "some_use1"({{.*}}) : (vector<4xf32>) -> vector<4xf32>
648+
// CHECK-DAG: %[[USE2:.+]] = "some_use2"({{.*}}) : (vector<5xf32>) -> vector<5xf32>
649+
// CHECK-NEXT: scf.yield %[[USE1]], %[[USE2]] : vector<4xf32>, vector<5xf32>
650+
// CHECK-NEXT: }
651+
// CHECK-DAG: %[[BCAST1:.+]] = vector.broadcast %[[LOOP]]#0 : vector<4xf32> to vector<3x4xf32>
652+
// CHECK-DAG: %[[BCAST2:.+]] = vector.broadcast %[[LOOP]]#1 : vector<5xf32> to vector<3x5xf32>
653+
// CHECK-NEXT: return %[[BCAST1]], %[[BCAST2]] : vector<3x4xf32>, vector<3x5xf32>
654+
655+
func.func @hoist_vector_broadcasts_multiple(%lb : index, %ub : index, %step : index, %vec1 : vector<3x4xf32>, %vec2 : vector<3x5xf32>) -> (vector<3x4xf32>, vector<3x5xf32>) {
656+
%bcast_vec:2 = scf.for %arg0 = %lb to %ub step %step iter_args(%iarg = %vec1, %iarg2 = %vec2) -> (vector<3x4xf32>, vector<3x5xf32>) {
657+
%extract1 = vector.extract %iarg[0] : vector<4xf32> from vector<3x4xf32>
658+
%extract2 = vector.extract %iarg2[1] : vector<5xf32> from vector<3x5xf32>
659+
%use1 = "some_use1"(%extract1) : (vector<4xf32>) -> vector<4xf32>
660+
%use2 = "some_use2"(%extract2) : (vector<5xf32>) -> vector<5xf32>
661+
%broadcast1 = vector.broadcast %use1 : vector<4xf32> to vector<3x4xf32>
662+
%broadcast2 = vector.broadcast %use2 : vector<5xf32> to vector<3x5xf32>
663+
scf.yield %broadcast1, %broadcast2 : vector<3x4xf32>,vector<3x5xf32>
664+
}
665+
return %bcast_vec#0, %bcast_vec#1 : vector<3x4xf32>, vector<3x5xf32>
666+
}
667+
668+
module attributes {transform.with_named_sequence} {
669+
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
670+
%0 = transform.structured.match ops{["func.func"]} in %arg1
671+
: (!transform.any_op) -> !transform.any_op
672+
transform.structured.hoist_redundant_vector_broadcasts %0
673+
: (!transform.any_op) -> !transform.any_op
674+
transform.yield
675+
}
676+
}

0 commit comments

Comments
 (0)