已合并
[VFFusion] Plan-time move and conflict update for all fused members #1838
[VFFusion] Plan-time move and conflict update for all fused members #1838
已合并
hujiajun创建于 26 天前
5 个文件变更+166-86
Mbishengir/include/bishengir/Dialect/HFusion/Transforms/AutoVectorize/FusedNode.h+2-2
@@ -107,8 +107,8 @@ private:
107 bool canFitStack(Operation *candidate) const;107 bool canFitStack(Operation *candidate) const;
108 bool hasRankReducing() const;108 bool hasRankReducing() const;
109 109 
110 void consolidateToTail(Block *block) const;110 void consolidateToTail() const;
111 void updateConflicts(Block *block) const;111 void updateConflicts(Operation *newOp) const;
112};112};
113 113 
114} // namespace hfusion114} // namespace hfusion
Mbishengir/include/bishengir/Dialect/HFusion/Transforms/AutoVectorize/PlanContext.h+17-1
@@ -25,6 +25,9 @@ namespace hfusion {
25// Every fusable op(including LinalgOp and interleave/deinterleave) corresponds25// Every fusable op(including LinalgOp and interleave/deinterleave) corresponds
26// to a FusableOpInfo.26// to a FusableOpInfo.
27struct FusableOpInfo {27struct FusableOpInfo {
28 template <typename PivotTy>
29 using ConflictPivotMap = DenseMap<Operation *, DenseSet<PivotTy>>;
30 
28 std::string label;31 std::string label;
29 int64_t numLoops = 0;32 int64_t numLoops = 0;
30 unsigned numReductionLoops = 0;33 unsigned numReductionLoops = 0;
@@ -32,7 +35,13 @@ struct FusableOpInfo {
32 SmallVector<int64_t> tileSize;35 SmallVector<int64_t> tileSize;
33 SmallVector<int64_t> tileInterchange;36 SmallVector<int64_t> tileInterchange;
34 unsigned maxElemBitWidth = 1;37 unsigned maxElemBitWidth = 1;
35 DenseSet<Operation *> conflictList;38 // TODO: Conflict storage (copy/sync/op/group) currently lives in PlanContext
39 // for convenience. Move to a dedicated ConflictTracker class to separate plan
40 // state from conflict tracking, making incremental update semantics clearer.
41 DenseSet<Operation *> copyConflicts;
42 DenseSet<Operation *> syncConflicts;
43 ConflictPivotMap<Operation *> opConflicts;
44 ConflictPivotMap<const FusedNode *> groupConflicts;
36 std::shared_ptr<FusedNode> fusedNode = nullptr;45 std::shared_ptr<FusedNode> fusedNode = nullptr;
37};46};
38 47 
@@ -109,9 +118,16 @@ public:
109 // Compute tile sizes for all fused ops after fusion planning.118 // Compute tile sizes for all fused ops after fusion planning.
110 void computeTileSize();119 void computeTileSize();
111 120 
121 void addSyncConflict(Operation *a, Operation *b);
122 void addCopyConflict(Operation *a, Operation *b);
123 void addOpConflict(Operation *a, Operation *b, Operation *pivot);
124 void addGroupConflict(Operation *a, Operation *b, const FusedNode *pivot);
125 
112 bool hasConflict(Operation *a, Operation *b) const;126 bool hasConflict(Operation *a, Operation *b) const;
113 bool hasConflict(const FusableOpInfo &a, Operation *b) const;127 bool hasConflict(const FusableOpInfo &a, Operation *b) const;
114 128 
129 void dissolvePivot(Operation *newOp, const FusedNode *node, Block *block);
130 
115private:131private:
116 // Register a fusable op with its label; computes numLoops/shape/bitWidth.132 // Register a fusable op with its label; computes numLoops/shape/bitWidth.
117 void registerAndAnalyzeOp(Operation *op, const std::string &label);133 void registerAndAnalyzeOp(Operation *op, const std::string &label);
Mbishengir/lib/Dialect/HFusion/Transforms/AutoVectorize/FusedNode.cpp+52-53
@@ -15,6 +15,7 @@
15#include "bishengir/Dialect/Utils/Util.h"15#include "bishengir/Dialect/Utils/Util.h"
16#include "mlir/Dialect/Linalg/IR/Linalg.h"16#include "mlir/Dialect/Linalg/IR/Linalg.h"
17#include "mlir/IR/Dominance.h"17#include "mlir/IR/Dominance.h"
18#include "mlir/IR/Verifier.h"
18#include "llvm/Support/Debug.h"19#include "llvm/Support/Debug.h"
19 20 
20using namespace mlir;21using namespace mlir;
@@ -29,13 +30,15 @@ namespace hfusion {
29void FusedNode::addProducer(Operation *op) {30void FusedNode::addProducer(Operation *op) {
30 producers.insert(op);31 producers.insert(op);
31 ctx.getInfo(op).fusedNode = shared_from_this();32 ctx.getInfo(op).fusedNode = shared_from_this();
33 consolidateToTail();
34 updateConflicts(op);
32}35}
33 36 
34void FusedNode::addLeaf(Operation *op) {37void FusedNode::addLeaf(Operation *op) {
35 leaf.insert(op);38 leaf.insert(op);
36 ctx.getInfo(op).fusedNode = shared_from_this();39 ctx.getInfo(op).fusedNode = shared_from_this();
37 consolidateToTail(op->getBlock());40 consolidateToTail();
38 updateConflicts(op->getBlock());41 updateConflicts(op);
39}42}
40 43 
41bool FusedNode::hasCommonAxis(Operation *op) const {44bool FusedNode::hasCommonAxis(Operation *op) const {
@@ -94,72 +97,68 @@ bool FusedNode::canAccept(Operation *candidate, AcceptContext actx) const {
94// Helpers for consolidateToTail97// Helpers for consolidateToTail
95//===----------------------------------------------------------------------===//98//===----------------------------------------------------------------------===//
96 99 
97static void100static void climbedForwardSlice(Operation *op,
98visitUsersOfLeafNodeRecursively(Operation *op, Operation *lastLeafNode,101 SetVector<Operation *> &forwardSlice,
99 Block *block,102 const DominanceInfo &domInfo,
100 DenseSet<Operation *> &usersToBeMovedSet) {103 Operation *fusedLoop) {
101 DominanceInfo domInfo;104 while (op->getParentOp() != fusedLoop->getParentOp())
102 for (Operation *user : op->getUsers()) {105 op = op->getParentOp();
103 if (!user || user->hasTrait<OpTrait::IsTerminator>() ||106 if (forwardSlice.count(op) || domInfo.dominates(fusedLoop, op))
104 domInfo.properlyDominates(lastLeafNode, user)) {107 return;
105 continue;108 for (Operation *userOp : op->getUsers())
106 }109 climbedForwardSlice(userOp, forwardSlice, domInfo, fusedLoop);
107 while (!isOpInBlock(user, block)) {110 forwardSlice.insert(op);
108 user = user->getParentOp();
109 }
110 if (usersToBeMovedSet.contains(user)) {
111 continue;
112 }
113 usersToBeMovedSet.insert(user);
114 visitUsersOfLeafNodeRecursively(user, lastLeafNode, block,
115 usersToBeMovedSet);
116 }
117}111}
118 112 
119void FusedNode::consolidateToTail(Block *block) const {113// NOTE: Called after each addLeaf/addProducer, so DominanceInfo is recomputed
120 SmallVector<Operation *> leafNodeGroup{leafOps()};114// per call. Fused-node sizes are bounded by maxFusedOps (small), so the
121 if (leafNodeGroup.size() == 1)115// overhead is negligible in practice.
116void FusedNode::consolidateToTail() const {
117 SmallVector<Operation *> opGroup{ops().begin(), ops().end()};
118 if (opGroup.size() == 1)
122 return;119 return;
123 120 
124 llvm::sort(leafNodeGroup,121 // Sort in reverse block order so that opGroup.front() is the last op in the
125 [](Operation *a, Operation *b) { return a->isBeforeInBlock(b); });122 // block and iteration proceeds from back to front, gathering each earlier op
126 Operation *lastLeafNode = leafNodeGroup.back();123 // and its forward-slice users toward the tail of the block.
127 DenseSet<Operation *> usersToBeMovedSet;124 llvm::sort(opGroup,
128 for (Operation *leafNode : llvm::drop_end(leafNodeGroup)) {125 [](Operation *a, Operation *b) { return b->isBeforeInBlock(a); });
129 visitUsersOfLeafNodeRecursively(leafNode, lastLeafNode, block,126 Operation *firstOp = opGroup.front();
130 usersToBeMovedSet);127 Operation *lastOp = opGroup.front();
131 }
132 SmallVector<Operation *> usersToBeMoved(usersToBeMovedSet.begin(),
133 usersToBeMovedSet.end());
134 llvm::sort(usersToBeMoved,
135 [](Operation *a, Operation *b) { return a->isBeforeInBlock(b); });
136 128 
137 Operation *prevMovedLeafNode = lastLeafNode;129 DominanceInfo domInfo(opGroup.back());
138 for (Operation *leafNodeToBeMoved : llvm::drop_end(leafNodeGroup)) {130 for (auto *curr : llvm::drop_begin(opGroup)) {
139 leafNodeToBeMoved->moveBefore(prevMovedLeafNode);131 SetVector<Operation *> forwardSlice;
140 prevMovedLeafNode = leafNodeToBeMoved;132 climbedForwardSlice(curr, forwardSlice, domInfo, firstOp);
141 }133 
142 Operation *prevMovedUser = lastLeafNode;134 for (auto *userOp : drop_end(forwardSlice))
143 for (Operation *userToBeMoved : usersToBeMoved) {135 userOp->moveAfter(lastOp);
144 userToBeMoved->moveAfter(prevMovedUser);136 curr->moveBefore(firstOp);
145 prevMovedUser = userToBeMoved;137 firstOp = curr;
146 }138 }
139 if (failed(verify(opGroup.back()->getParentOp())))
140 llvm::report_fatal_error("consolidateToTail produced invalid IR");
147}141}
148 142 
149void FusedNode::updateConflicts(Block *block) const {143void FusedNode::updateConflicts(Operation *newOp) const {
144 Block *block = newOp->getBlock();
145 ctx.dissolvePivot(newOp, this, block);
146 
150 DenseSet<Operation *> upstreamOps;147 DenseSet<Operation *> upstreamOps;
151 DenseSet<Operation *> visitedUpstreamOps;148 DenseSet<Operation *> visitedUpstreamOps;
152 DenseSet<Operation *> downstreamOps;149 DenseSet<Operation *> downstreamOps;
153 DenseSet<Operation *> visitedDownstreamOps;150 DenseSet<Operation *> visitedDownstreamOps;
154 for (Operation *leafNode : leafOps()) {151 for (Operation *op : ops()) {
155 findUpstreamFusableOpOf(leafNode, block, upstreamOps, visitedUpstreamOps);152 findUpstreamFusableOpOf(op, block, upstreamOps, visitedUpstreamOps);
156 findDownstreamFusableOpOf(leafNode, block, downstreamOps,153 findDownstreamFusableOpOf(op, block, downstreamOps, visitedDownstreamOps);
157 visitedDownstreamOps);154 }
155 for (Operation *op : ops()) {
156 upstreamOps.erase(op);
157 downstreamOps.erase(op);
158 }158 }
159 for (auto upstreamOp : upstreamOps) {159 for (auto upstreamOp : upstreamOps) {
160 for (auto downstreamOp : downstreamOps) {160 for (auto downstreamOp : downstreamOps) {
161 ctx.getInfo(upstreamOp).conflictList.insert(downstreamOp);161 ctx.addGroupConflict(upstreamOp, downstreamOp, this);
162 ctx.getInfo(downstreamOp).conflictList.insert(upstreamOp);
163 }162 }
164 }163 }
165}164}
Mbishengir/lib/Dialect/HFusion/Transforms/AutoVectorize/PlanContext.cpp+95-25
@@ -118,11 +118,21 @@ void PlanContext::initFusableOpInfoFrom(func::FuncOp func) {
118 computeConflictLists(func);118 computeConflictLists(func);
119#ifndef NDEBUG119#ifndef NDEBUG
120 LLVM_DEBUG(llvm::dbgs() << "========Dumping conflict lists begin========\n");120 LLVM_DEBUG(llvm::dbgs() << "========Dumping conflict lists begin========\n");
121 for (auto info : opInfos()) {121 for (auto &[op, info] : opInfos()) {
122 LLVM_DEBUG(llvm::dbgs() << "========Dumping op========\n");122 LLVM_DEBUG(llvm::dbgs() << "========Dumping op========\n");
123 LLVM_DEBUG(llvm::dbgs() << *info.first << "\n");123 LLVM_DEBUG(llvm::dbgs() << *op << "\n");
124 LLVM_DEBUG(llvm::dbgs() << "========Dumping conflict list========\n");124 LLVM_DEBUG(llvm::dbgs() << "========Dumping conflict list========\n");
125 for (auto op : info.second.conflictList)125 LLVM_DEBUG(llvm::dbgs() << "++++++++sync conflicts++++++++\n");
126 for (auto *op : info.syncConflicts)
127 LLVM_DEBUG(llvm::dbgs() << *op << "\n");
128 LLVM_DEBUG(llvm::dbgs() << "++++++++copy conflicts++++++++\n");
129 for (auto *op : info.copyConflicts)
130 LLVM_DEBUG(llvm::dbgs() << *op << "\n");
131 LLVM_DEBUG(llvm::dbgs() << "++++++++op conflicts++++++++\n");
132 for (auto &[op, _] : info.opConflicts)
133 LLVM_DEBUG(llvm::dbgs() << *op << "\n");
134 LLVM_DEBUG(llvm::dbgs() << "++++++++group conflicts++++++++\n");
135 for (auto &[op, _] : info.groupConflicts)
126 LLVM_DEBUG(llvm::dbgs() << *op << "\n");136 LLVM_DEBUG(llvm::dbgs() << *op << "\n");
127 }137 }
128 LLVM_DEBUG(llvm::dbgs() << "\n");138 LLVM_DEBUG(llvm::dbgs() << "\n");
@@ -148,6 +158,22 @@ void PlanContext::computeTileSize() {
148 158 
149namespace {159namespace {
150 160 
161template <typename PivotTy>
162static void
163dissolvePivotImpl(FusableOpInfo::ConflictPivotMap<PivotTy> &conflictA,
164 Operation *a,
165 FusableOpInfo::ConflictPivotMap<PivotTy> &conflictB,
166 Operation *b, PivotTy pivot) {
167 if (!conflictA.contains(b) || !conflictB.contains(a) ||
168 conflictA[b].erase(pivot) != conflictB[a].erase(pivot) ||
169 conflictA[b].empty() != conflictB[a].empty())
170 llvm::report_fatal_error("inconsistent conflict state");
171 if (!conflictA[b].empty())
172 return;
173 conflictA.erase(b);
174 conflictB.erase(a);
175}
176 
151void findPreviousAndFollowingFusableOpOf(Operation *barrierOp, Block *block,177void findPreviousAndFollowingFusableOpOf(Operation *barrierOp, Block *block,
152 DenseSet<Operation *> &previousOps,178 DenseSet<Operation *> &previousOps,
153 DenseSet<Operation *> &followingOps) {179 DenseSet<Operation *> &followingOps) {
@@ -187,8 +213,7 @@ void computeConflictListsForCopyOpOperand(PlanContext &ctx,
187 for (auto previousOp : previousOps) {213 for (auto previousOp : previousOps) {
188 if (hasMemRefInOperands(previousOp, memRef)) {214 if (hasMemRefInOperands(previousOp, memRef)) {
189 for (auto followingOp : followingOps) {215 for (auto followingOp : followingOps) {
190 ctx.getInfo(previousOp).conflictList.insert(followingOp);216 ctx.addCopyConflict(previousOp, followingOp);
191 ctx.getInfo(followingOp).conflictList.insert(previousOp);
192 }217 }
193 }218 }
194 }219 }
@@ -196,8 +221,7 @@ void computeConflictListsForCopyOpOperand(PlanContext &ctx,
196 for (auto followingOp : followingOps) {221 for (auto followingOp : followingOps) {
197 if (hasMemRefInOperands(followingOp, memRef)) {222 if (hasMemRefInOperands(followingOp, memRef)) {
198 for (auto previousOp : previousOps) {223 for (auto previousOp : previousOps) {
199 ctx.getInfo(previousOp).conflictList.insert(followingOp);224 ctx.addCopyConflict(previousOp, followingOp);
200 ctx.getInfo(followingOp).conflictList.insert(previousOp);
201 }225 }
202 }226 }
203 }227 }
@@ -205,26 +229,55 @@ void computeConflictListsForCopyOpOperand(PlanContext &ctx,
205 229 
206} // namespace230} // namespace
207 231 
232void PlanContext::dissolvePivot(Operation *newOp, const FusedNode *node,
233 Block *block) {
234 // Step 1 — discharge op-level def-use conflicts induced by `newOp`.
235 // When `newOp` was standalone it sat between upstreamOps and their
236 // downstream, creating pivot entries in opConflicts. Now that `newOp`
237 // is fused into `node`, erase it from every (ancestor, descendant)
238 // pair so that pairs that lose their last pivot also lose the conflict.
239 DenseSet<Operation *> upstreamOps;
240 DenseSet<Operation *> visitedUpstreamOps;
241 findUpstreamFusableOpOf(newOp, block, upstreamOps, visitedUpstreamOps);
242 for (Operation *upstreamOp : upstreamOps) {
243 auto &conflicts = getInfo(upstreamOp).opConflicts;
244 for (auto [otherOp, pivotSet] : llvm::make_early_inc_range(conflicts)) {
245 auto &otherConflicts = getInfo(otherOp).opConflicts;
246 dissolvePivotImpl(conflicts, upstreamOp, otherConflicts, otherOp, newOp);
247 }
248 }
249 
250 // Step 2 — discharge group-level conflicts where `node` blocks
251 // `newOp` from sharing a loop with other ops. Since `newOp` just
252 // joined `node`, the group is no longer a barrier between them.
253 auto &conflicts = getInfo(newOp).groupConflicts;
254 for (auto [otherOp, pivotSet] : llvm::make_early_inc_range(conflicts)) {
255 auto &otherConflicts = getInfo(otherOp).groupConflicts;
256 dissolvePivotImpl(conflicts, newOp, otherConflicts, otherOp, node);
257 }
258}
259 
208void PlanContext::computeConflictLists(func::FuncOp func) {260void PlanContext::computeConflictLists(func::FuncOp func) {
209 func.walk([&](Block *block) {261 func.walk([&](Block *block) {
210 if (isa<func::FuncOp, scf::ForOp, scf::IfOp, scf::WhileOp, scope::ScopeOp>(262 if (isa<func::FuncOp, scf::ForOp, scf::IfOp, scf::WhileOp, scope::ScopeOp>(
211 block->getParentOp())) {263 block->getParentOp())) {
212 block->walk([&](Operation *op) {264 block->walk([&](Operation *op) {
213 if (isOpInBlock(op, block) && isNonVectorizableOp(op)) {265 if (!isOpInBlock(op, block))
214 DenseSet<Operation *> upstreamOps;266 return;
215 DenseSet<Operation *> visitedUpstreamOps;267 DenseSet<Operation *> upstreamOps;
216 findUpstreamFusableOpOf(op, block, upstreamOps, visitedUpstreamOps);268 DenseSet<Operation *> visitedUpstreamOps;
217 DenseSet<Operation *> downstreamOps;269 findUpstreamFusableOpOf(op, block, upstreamOps, visitedUpstreamOps);
218 DenseSet<Operation *> visitedDownstreamOps;270 DenseSet<Operation *> downstreamOps;
219 findDownstreamFusableOpOf(op, block, downstreamOps,271 DenseSet<Operation *> visitedDownstreamOps;
220 visitedDownstreamOps);272 findDownstreamFusableOpOf(op, block, downstreamOps,
221 for (auto upstreamOp : upstreamOps) {273 visitedDownstreamOps);
222 for (auto downstreamOp : downstreamOps) {274 for (auto upstreamOp : upstreamOps) {
223 getInfo(upstreamOp).conflictList.insert(downstreamOp);275 for (auto downstreamOp : downstreamOps) {
224 getInfo(downstreamOp).conflictList.insert(upstreamOp);276 addOpConflict(upstreamOp, downstreamOp, op);
225 }
226 }277 }
278 }
227 279 
280 if (isNonVectorizableOp(op)) {
228 if (enableCrossIfFusion) {281 if (enableCrossIfFusion) {
229 // Only region-bearing ops or explicit sync/copy ops need the282 // Only region-bearing ops or explicit sync/copy ops need the
230 // previous/following scan and the body walk below.283 // previous/following scan and the body walk below.
@@ -263,8 +316,7 @@ void PlanContext::computeConflictLists(func::FuncOp func) {
263 if (op->walk(walker).wasInterrupted()) {316 if (op->walk(walker).wasInterrupted()) {
264 for (auto previousOp : previousOps) {317 for (auto previousOp : previousOps) {
265 for (auto followingOp : followingOps) {318 for (auto followingOp : followingOps) {
266 getInfo(previousOp).conflictList.insert(followingOp);319 addSyncConflict(previousOp, followingOp);
267 getInfo(followingOp).conflictList.insert(previousOp);
268 }320 }
269 }321 }
270 }322 }
@@ -293,8 +345,7 @@ void PlanContext::computeConflictLists(func::FuncOp func) {
293 followingOps);345 followingOps);
294 for (auto previousOp : previousOps) {346 for (auto previousOp : previousOps) {
295 for (auto followingOp : followingOps) {347 for (auto followingOp : followingOps) {
296 getInfo(previousOp).conflictList.insert(followingOp);348 addSyncConflict(previousOp, followingOp);
297 getInfo(followingOp).conflictList.insert(previousOp);
298 }349 }
299 }350 }
300 }351 }
@@ -304,8 +355,27 @@ void PlanContext::computeConflictLists(func::FuncOp func) {
304 });355 });
305}356}
306 357 
358void PlanContext::addSyncConflict(Operation *a, Operation *b) {
359 getInfo(a).syncConflicts.insert(b);
360 getInfo(b).syncConflicts.insert(a);
361}
362void PlanContext::addCopyConflict(Operation *a, Operation *b) {
363 getInfo(a).copyConflicts.insert(b);
364 getInfo(b).copyConflicts.insert(a);
365}
366void PlanContext::addOpConflict(Operation *a, Operation *b, Operation *pivot) {
367 getInfo(a).opConflicts[b].insert(pivot);
368 getInfo(b).opConflicts[a].insert(pivot);
369}
370void PlanContext::addGroupConflict(Operation *a, Operation *b,
371 const FusedNode *pivot) {
372 getInfo(a).groupConflicts[b].insert(pivot);
373 getInfo(b).groupConflicts[a].insert(pivot);
374}
375 
307bool PlanContext::hasConflict(const FusableOpInfo &a, Operation *b) const {376bool PlanContext::hasConflict(const FusableOpInfo &a, Operation *b) const {
308 return a.conflictList.contains(b);377 return a.syncConflicts.contains(b) || a.copyConflicts.contains(b) ||
378 a.opConflicts.contains(b) || a.groupConflicts.contains(b);
309}379}
310bool PlanContext::hasConflict(Operation *a, Operation *b) const {380bool PlanContext::hasConflict(Operation *a, Operation *b) const {
311 return hasConflict(getInfo(a), b);381 return hasConflict(getInfo(a), b);
Mbishengir/lib/Dialect/HFusion/Transforms/AutoVectorizeV2.cpp+0-5
@@ -254,11 +254,6 @@ findBestFusedNodeForProducer(Block *block, Operation *producer,
254 FusableOpInfo &producerInfo = ctx.getInfo(producer);254 FusableOpInfo &producerInfo = ctx.getInfo(producer);
255 if (!bestFusedNode->canFuseProducer(producer))255 if (!bestFusedNode->canFuseProducer(producer))
256 return nullptr;256 return nullptr;
257 // If the closest fuseNode is conflict with the producer, give up fusing.
258 if (llvm::any_of(bestFusedNode->ops(), [&](Operation *fusedOp) {
259 return producerInfo.conflictList.contains(fusedOp);
260 }))
261 return nullptr;
262 int numUsersInBestFusedNode = 0;257 int numUsersInBestFusedNode = 0;
263 for (auto user : DenseSet<Operation *>(producer->getUsers().begin(),258 for (auto user : DenseSet<Operation *>(producer->getUsers().begin(),
264 producer->getUsers().end())) {259 producer->getUsers().end())) {