torch-mlir/lib/Dialect/ATen/LivenessReport.cpp

254 lines
6.5 KiB
C++

//===- LivenessReport.cpp ---------------------------------------*- C++ -*-===//
//
// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
#include "npcomp/Dialect/ATen/ATenDialect.h"
#include "mlir/Analysis/Liveness.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Module.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/JSON.h"
#include "npcomp/Dialect/ATen/LivenessReport.h"
#include <iostream>
#include <sstream>
#include <vector>
#define DEBUG_TYPE "liveness-report"
using namespace mlir;
namespace {
uint64_t getTensorVolume(const ShapedType ty) {
if (!ty.hasRank())
return 1;
uint64_t volume = 1;
for (auto &d : ty.getShape())
volume *= d;
return volume;
}
uint64_t getTensorVolume(const Type ty) {
if (auto t = ty.dyn_cast<ShapedType>()) {
return getTensorVolume(t);
} else {
return 1;
}
}
} // namespace
namespace mlir {
namespace NPCOMP {
namespace aten {
std::string LivenessReport::generateTextReport() {
resolveLiveness();
std::string output;
for (auto &p : livenessIntervals) {
Value v = p.first;
std::vector<Operation *> &oplist = p.second;
llvm::outs() << "// begin\n";
v.print(llvm::outs());
llvm::outs() << "\n";
for (auto *o : oplist) {
o->print(llvm::outs());
llvm::outs() << "\n";
}
llvm::outs() << "// end \n";
}
return output;
}
std::string LivenessReport::emitJSONReport() {
resolveLiveness();
llvm::json::Object top;
auto context = module.getContext();
auto loc = mlir::UnknownLoc::get(context);
auto graph = module.lookupSymbol<mlir::FuncOp>("graph");
std::map<Operation *, std::vector<Value>> liveAt;
graph.walk([&](Operation *op) {
for (Value result : op->getResults()) {
for (auto &p : livenessIntervals) {
Value v = p.first;
if (v == result) {
std::vector<Operation *> &oplist = p.second;
for (auto *o : oplist)
liveAt[o].push_back(result);
}
}
}
});
auto argList = graph.getBody().getBlocks().front().getArguments();
for (Value &arg : argList) {
for (auto &p : livenessIntervals) {
Value v = p.first;
if (v == arg) {
std::vector<Operation *> &oplist = p.second;
for (auto *o : oplist)
liveAt[o].push_back(arg);
}
}
}
graph.walk([&](Operation *op) {
llvm::json::Object layerDetail;
auto attr = op->getAttrOfType<StringAttr>("layer_name");
if (!attr)
return;
std::vector<Value> &vlist = liveAt[op];
int64_t parameterVol = 0;
int64_t returnVol = 0;
for (auto v : vlist) {
int64_t vol = getTensorVolume(v.getType());
if (v.getDefiningOp()) {
if (auto a =
v.getDefiningOp()->getAttrOfType<StringAttr>("layer_name")) {
auto definingOp = v.getDefiningOp();
auto ld = layerDetail.getInteger(a.getValue().str());
if (ld)
layerDetail[a.getValue().str()] = *ld + vol;
else
layerDetail[a.getValue().str()] = vol;
} else {
llvm_unreachable("unknown type");
}
} else if (std::find(argList.begin(), argList.end(), v) !=
argList.end()) {
parameterVol += vol;
} else {
llvm_unreachable("unknown type");
}
auto ret = cast<ReturnOp>(op->getBlock()->getTerminator());
for (auto oper : ret.getOperands()) {
if (oper == v) {
returnVol += vol;
break;
}
}
}
if (parameterVol) {
layerDetail["parameters"] = parameterVol;
}
if (returnVol) {
layerDetail["returns"] = returnVol;
}
top[attr.getValue().str()] = llvm::json::Value(std::move(layerDetail));
});
llvm::json::Value topv(std::move(top));
std::string ret;
llvm::raw_string_ostream ss(ret);
ss << llvm::formatv("{0:2}", topv) << "\n";
return ss.str();
}
void LivenessReport::resolveLiveness() {
auto context = module.getContext();
auto loc = mlir::UnknownLoc::get(context);
// check that a function called "graph" exists
auto graph = module.lookupSymbol<mlir::FuncOp>("graph");
if (!graph) {
emitError(mlir::UnknownLoc::get(module.getContext()),
"LivenessReport failed: can't find a graph function\n");
return;
}
// put each aten operation into its own basic block,
// so that we can use standard liveness
Region &bodyRegion = graph.getBody();
Block *entryBB = &bodyRegion.getBlocks().front();
Block *BB = entryBB;
std::vector<Block *> new_blocks;
while (true) {
std::vector<Operation *> ops;
for (Operation &op : BB->getOperations())
ops.push_back(&op);
// skip over constant ops
int idx = 0;
while (dyn_cast<mlir::NPCOMP::aten::ConstantOp>(ops[idx]))
idx++;
if (dyn_cast<ReturnOp>(ops[idx]))
break;
Block *newBB = BB->splitBlock(ops[idx + 1]);
new_blocks.push_back(newBB);
mlir::OpBuilder builder = mlir::OpBuilder::atBlockBegin(BB);
builder.create<BranchOp>(loc, newBB);
BB = newBB;
}
// dump transformed function
// graph.dump();
// run MLIR Liveness analysis
auto liveness = Liveness(graph);
for (BlockArgument &arg :
graph.getBody().getBlocks().front().getArguments()) {
auto liveOps = liveness.resolveLiveness(arg);
for (Operation *o : liveOps) {
livenessIntervals[arg].push_back(o);
}
}
graph.walk([&](Operation *op) {
auto attr = op->getAttrOfType<StringAttr>("layer_name");
if (!attr)
return;
for (Value v : op->getResults()) {
auto liveOps = liveness.resolveLiveness(v);
for (Operation *o : liveOps)
if (auto a = o->getAttrOfType<StringAttr>("layer_name"))
livenessIntervals[v].push_back(o);
}
});
// undo the BB insert
auto *deadBr = bodyRegion.getBlocks().front().getTerminator();
for (Block *b : new_blocks) {
auto *br = b->getTerminator();
std::vector<Operation *> ops;
for (auto &op : b->getOperations())
ops.push_back(&op);
for (auto *op : ops) {
if (op == br && !dyn_cast<ReturnOp>(op)) {
op->erase();
continue;
}
op->moveBefore(deadBr);
}
}
deadBr->erase();
for (Block *b : new_blocks)
b->erase();
// graph.dump();
}
} // namespace aten
} // namespace NPCOMP
} // namespace mlir