Skip to content

Commit 8689052

Browse files
authored
[Flow] Dump affinity info in DumpDispatchGraph pass. (#20888)
No tests are added because they don't exist. It dumps more information, so we can do basic affinity analysis outside the main tree. E.g., people can use the tool from https://github.com/hanhanW/iree-toolbox to verify if the input program has reasonable affinity assignment or not. --------- Signed-off-by: hanhanW <[email protected]>
1 parent fa019c9 commit 8689052

File tree

1 file changed

+9
-0
lines changed

1 file changed

+9
-0
lines changed

compiler/src/iree/compiler/Dialect/Flow/Transforms/DumpDispatchGraph.cpp

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,9 @@
1212

1313
#include <utility>
1414

15+
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
1516
#include "iree/compiler/Dialect/Flow/Transforms/Passes.h"
17+
#include "iree/compiler/Dialect/Stream/IR/StreamInterfaces.h"
1618
#include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
1719
#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
1820
#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
@@ -478,6 +480,13 @@ class GraphPrinter {
478480
} else {
479481
os << op->getName() << "\n";
480482
}
483+
if (auto affinityOp = dyn_cast<IREE::Stream::AffinityOpInterface>(op)) {
484+
os << affinityOp.getAffinityAttr() << "\n";
485+
} else if (auto transferOp = dyn_cast<IREE::Flow::TensorTransferOp>(op)) {
486+
os << transferOp.getTarget() << "\n";
487+
} else if (auto barrierOp = dyn_cast<IREE::Flow::TensorBarrierOp>(op)) {
488+
os << barrierOp.getTarget() << "\n";
489+
}
481490

482491
if (printResultTypes) {
483492
std::string buf;

0 commit comments

Comments
 (0)