add PIM accelerator
This commit is contained in:
NiccoloN
283
src/PIM/Conversion/SpatialToGraphviz/SpatialToGraphviz.cpp
Normal file
283
src/PIM/Conversion/SpatialToGraphviz/SpatialToGraphviz.cpp
Normal file
@@ -0,0 +1,283 @@
|
||||
#include "mlir/IR/Block.h"
|
||||
#include "mlir/IR/Diagnostics.h"
|
||||
#include "mlir/IR/Value.h"
|
||||
#include "mlir/Support/LLVM.h"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
||||
#include "src/Accelerators/PIM/Pass/PimPasses.hpp"
|
||||
#include "src/Dialect/ONNX/ONNXOps.hpp"
|
||||
|
||||
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
||||
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
|
||||
#include "mlir/Pass/Pass.h"
|
||||
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
||||
|
||||
#include "llvm/Support/Casting.h"
|
||||
#include "llvm/Support/Format.h"
|
||||
|
||||
#define FORMAT_OPERATION(op) \
|
||||
'x' << llvm::format_hex_no_prefix(reinterpret_cast<size_t>(op), 0)
|
||||
#define FORMAT_ARGUMENT(computeOpPointer, argumentNum) \
|
||||
llvm::format("Arg_%p_%u", computeOpPointer, argumentNum)
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace onnx_mlir {
|
||||
|
||||
namespace {
|
||||
|
||||
struct SpatialToGraphvizPass
|
||||
: public PassWrapper<SpatialToGraphvizPass, OperationPass<ModuleOp>> {
|
||||
|
||||
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialToGraphvizPass)
|
||||
|
||||
StringRef getArgument() const override {
|
||||
return "convert-spatial-to-graphviz";
|
||||
}
|
||||
|
||||
StringRef getDescription() const override {
|
||||
return "Lower ONNX ops to Spatial ops.";
|
||||
}
|
||||
|
||||
SpatialToGraphvizPass(raw_ostream &os = llvm::errs()) : os(os) {}
|
||||
SpatialToGraphvizPass(const SpatialToGraphvizPass &pass)
|
||||
: SpatialToGraphvizPass(pass.os) {}
|
||||
void runOnOperation() final;
|
||||
|
||||
private:
|
||||
raw_ostream &os;
|
||||
|
||||
/**
|
||||
* Draws the subgraph for a given spatial::SpatWeightedCompute, including:
|
||||
* 1. Input nodes (block arguments)
|
||||
* 2. Operations
|
||||
* 3. Edges between yield (output) and its users
|
||||
*
|
||||
* @param op The spatial::SpatWeightedCompute to draw the subgraph for.
|
||||
* @param computeNum The number of the compute operation.
|
||||
*/
|
||||
void drawComputeOpSubgraph(spatial::SpatWeightedCompute op, size_t computeNum) {
|
||||
os << "\tsubgraph cluster" << computeNum << " {\n\t\tlabel=\"Compute"
|
||||
<< computeNum << "\";\n"
|
||||
<< "\t\tstyle=filled;\n"
|
||||
<< "\t\tcolor=lightblue;\n";
|
||||
|
||||
Block &block = op.getBody().front();
|
||||
|
||||
// Inputs
|
||||
size_t inputNum = 0;
|
||||
for (BlockArgument &input : block.getArguments()) {
|
||||
|
||||
auto fromOp = FORMAT_ARGUMENT(op.getOperation(), inputNum);
|
||||
|
||||
os << "\t\t" << fromOp << " [label=\"Arg" << inputNum
|
||||
<< "\",shape=box];\n";
|
||||
for (auto userOp : input.getUsers()) {
|
||||
os << "\t\t" << fromOp << " -> " << FORMAT_OPERATION(userOp) << ";\n";
|
||||
}
|
||||
inputNum++;
|
||||
}
|
||||
|
||||
// Iterate operations
|
||||
for (auto &childOp : block.getOperations()) {
|
||||
os << "\t\t" << FORMAT_OPERATION(&childOp) << " [label=\""
|
||||
<< childOp.getName() << "\"];\n";
|
||||
|
||||
drawEdgesFromOpToItsUsers(&childOp);
|
||||
}
|
||||
|
||||
os << "\t}\n";
|
||||
|
||||
// Draw edges from the yield to the users of this computeOp
|
||||
Operation *yieldOp = block.getTerminator();
|
||||
if (!isa<spatial::SpatYieldOp>(yieldOp)) {
|
||||
yieldOp->emitError("Terminator of block must be YieldOp ???");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
|
||||
for (auto computeOpResult : op->getResults()) {
|
||||
for (auto &computeOpUse : computeOpResult.getUses()) {
|
||||
auto toOp = FORMAT_ARGUMENT(
|
||||
computeOpUse.getOwner(), computeOpUse.getOperandNumber());
|
||||
os << "\t" << FORMAT_OPERATION(yieldOp) << " -> " << toOp << ";\n";
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Draws the subgraph for a concatOp.
|
||||
*
|
||||
* This function draws a subgraph for a concatOp. The subgraph consists of a
|
||||
* node for each input of the concatOp, as well as an output node. Edges are
|
||||
* created from the output node to each user of the concatOp.
|
||||
*
|
||||
* @param concatOp The concatOp for which the subgraph is drawn.
|
||||
* @param concatOpNum The number of the concatOp.
|
||||
*/
|
||||
void drawConcatOpSubgraph(Operation *concatOp, size_t concatOpNum) {
|
||||
os << "\tsubgraph clusterconcat" << concatOpNum
|
||||
<< " {\n\t\tlabel=\"ConcatOp" << concatOpNum << "\";\n"
|
||||
<< "\t\tstyle=filled;\n"
|
||||
<< "\t\tcolor=orange;\n";
|
||||
|
||||
// Inputs
|
||||
size_t inputNum = 0;
|
||||
for (Value input : concatOp->getOperands()) {
|
||||
auto fromOp = FORMAT_ARGUMENT(concatOp, inputNum);
|
||||
|
||||
os << "\t\t" << fromOp << " [label=\"Input" << inputNum << "\"];\n";
|
||||
for (auto userOp : input.getUsers()) {
|
||||
os << "\t\t" << fromOp << " -> " << FORMAT_OPERATION(userOp) << ";\n";
|
||||
}
|
||||
inputNum++;
|
||||
}
|
||||
|
||||
// Output
|
||||
os << "\t\t" << FORMAT_OPERATION(concatOp) << " [label=Out];\n";
|
||||
|
||||
os << "\t}\n";
|
||||
|
||||
// Edges from output to users
|
||||
|
||||
for (auto &computeOpUse : concatOp->getResult(0).getUses()) {
|
||||
os << "\t" << FORMAT_OPERATION(concatOp) << " -> "
|
||||
<< FORMAT_ARGUMENT(
|
||||
computeOpUse.getOwner(), computeOpUse.getOperandNumber())
|
||||
<< ";\n";
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Draws the ExtractSliceOp in the graph visualization.
|
||||
*
|
||||
* This function takes a tensor::ExtractSliceOp and adds the corresponding
|
||||
* node and edges to the graph visualization. It creates a node with the
|
||||
* label as the static offsets attribute of the sliceOp, and connects it to
|
||||
* the compute operations that use the result of the sliceOp.
|
||||
*
|
||||
* @param sliceOp The tensor::ExtractSliceOp to be drawn in the graph
|
||||
* visualization.
|
||||
*/
|
||||
void drawExtractSliceOp(tensor::ExtractSliceOp sliceOp) {
|
||||
auto nodeId = FORMAT_ARGUMENT(sliceOp.getOperation(), 0);
|
||||
os << "\t" << nodeId << " [label=\"Slice: ";
|
||||
sliceOp.getStaticOffsetsAttr().print(os);
|
||||
os << "\",color=lawngreen];\n";
|
||||
|
||||
for (auto &computeOpUse : sliceOp.getResult().getUses()) {
|
||||
os << "\t" << nodeId << " -> "
|
||||
<< FORMAT_ARGUMENT(
|
||||
computeOpUse.getOwner(), computeOpUse.getOperandNumber())
|
||||
<< ";\n";
|
||||
}
|
||||
}
|
||||
|
||||
void drawBiasTileOp(tensor::ExtractSliceOp sliceOp) {
|
||||
auto nodeId = FORMAT_ARGUMENT(sliceOp.getOperation(), 0);
|
||||
os << "\t" << nodeId << " [label=\"Bias: ";
|
||||
sliceOp.getStaticOffsetsAttr().print(os);
|
||||
os << "\",color=lightpink];\n";
|
||||
|
||||
for (auto user : sliceOp.getResult().getUsers()) {
|
||||
os << "\t" << nodeId << " -> " << FORMAT_OPERATION(user) << ";\n";
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Draws edges from the given operation to its users.
|
||||
*
|
||||
* @param fromOp The operation from which the edges are drawn.
|
||||
*/
|
||||
void drawEdgesFromOpToItsUsers(mlir::Operation *fromOp) {
|
||||
for (auto result : fromOp->getResults()) {
|
||||
for (auto userOp : result.getUsers()) {
|
||||
os << "\t\t" << FORMAT_OPERATION(fromOp) << " -> "
|
||||
<< FORMAT_OPERATION(userOp) << ";\n";
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Draws input node and edges for the given `funcOp`.
|
||||
*
|
||||
* @param funcOp The `funcOp` for which to draw input nodes and edges.
|
||||
*/
|
||||
void drawInputNodesAndEdges(func::FuncOp &funcOp) {
|
||||
os << "\tinput [label=\"Module Input\",color=green];\n";
|
||||
|
||||
size_t funcOpArgNum = 0;
|
||||
for (BlockArgument &arg : funcOp.getArguments()) {
|
||||
|
||||
for (auto &useOp : arg.getUses()) {
|
||||
os << "\tinput -> "
|
||||
<< FORMAT_ARGUMENT(useOp.getOwner(), useOp.getOperandNumber())
|
||||
<< "[label=" << funcOpArgNum << "];\n";
|
||||
}
|
||||
funcOpArgNum++;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
void SpatialToGraphvizPass::runOnOperation() {
|
||||
ModuleOp module = getOperation();
|
||||
|
||||
// Get the first OP, must be a FuncOp
|
||||
func::FuncOp func = *module.getOps<func::FuncOp>().begin();
|
||||
if (!func) {
|
||||
module->emitError("No FuncOp found in the begin of module");
|
||||
signalPassFailure();
|
||||
}
|
||||
|
||||
os << "digraph G {\n"
|
||||
<< "\tnode [style=filled,color=white];\n";
|
||||
|
||||
size_t computeNum = 0;
|
||||
size_t concatNum = 0;
|
||||
|
||||
// Iterate over the ComputeOps within FuncOp:
|
||||
// 1. Print their subgraph
|
||||
// 2. Print the edges from its inputs to its outputs
|
||||
for (Operation &op : func.getOps()) {
|
||||
if (auto computeOp = dyn_cast<spatial::SpatWeightedCompute>(op)) {
|
||||
drawComputeOpSubgraph(computeOp, computeNum++);
|
||||
} else if (auto concatOp = dyn_cast<tensor::ConcatOp>(op)) {
|
||||
drawConcatOpSubgraph(concatOp, concatNum++);
|
||||
} else if (auto imgConcatOp = dyn_cast<spatial::SpatImgConcatOp>(op)) {
|
||||
drawConcatOpSubgraph(imgConcatOp, concatNum++);
|
||||
} else if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(op)) {
|
||||
auto producerOp = extractSliceOp->getOperand(0).getDefiningOp();
|
||||
if (producerOp) {
|
||||
// Skip extractSliceOp if producer is constant weights (ONNXConstantOp)
|
||||
if (llvm::isa<ONNXConstantOp>(producerOp)) {
|
||||
continue;
|
||||
}
|
||||
// If produced by tosa::ReshapeOp (i.e. it is a bias tile) connect
|
||||
// directly to its user, which is not a ComputeOp argument.
|
||||
if (llvm::isa<tosa::ReshapeOp>(producerOp)) {
|
||||
drawBiasTileOp(extractSliceOp);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
drawExtractSliceOp(extractSliceOp);
|
||||
}
|
||||
}
|
||||
|
||||
// Draw input node, and edges to it users
|
||||
drawInputNodesAndEdges(func);
|
||||
|
||||
// Draw output node (use the return Operation - argument number=0 - as nodeId)
|
||||
auto returnOp = func.getBody().front().getTerminator();
|
||||
os << '\t' << FORMAT_ARGUMENT(returnOp, 0)
|
||||
<< " [label=\"Module Output\",color=green];\n";
|
||||
|
||||
os << "}\n";
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
std::unique_ptr<Pass> createSpatialToGraphvizPass() {
|
||||
return std::make_unique<SpatialToGraphvizPass>();
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir
|
||||
Reference in New Issue
Block a user