torch-mlir/frontends/pytorch/csrc/builder/node_importer.cpp

329 lines
12 KiB
C++

//===- node_importer.cpp --------------------------------------------------===//
//
// This file is licensed under a pytorch-style license
// See frontends/pytorch/LICENSE for license information.
//
//===----------------------------------------------------------------------===//
#include "node_importer.h"
#include <unordered_map>
#include "mlir_utils.h"
#include "op_builder.h"
#include "mlir-c/BuiltinAttributes.h"
#include "mlir-c/BuiltinTypes.h"
#include "mlir-c/Diagnostics.h"
#include "npcomp-c/Types.h"
namespace py = pybind11;
using namespace torch_mlir;
using Value = torch::jit::Value;
using Block = torch::jit::Block;
using Node = torch::jit::Node;
namespace {
class NodeImporter {
public:
NodeImporter(MlirContext context) : context(context) {}
void importNode(Node *node, MlirBlock appendToBlock);
MlirBlock importBlock(Block *jitBlock, CreateTerminatorFn createTerminator);
private:
void importPrimNode(Node *node, MlirBlock appendToBlock);
void importKernelCall(Node *node, MlirBlock appendToBlock);
MlirBlock createBlockFor(Block *jitBlock);
void mapValue(Value *jitValue, MlirValue value);
void mapResults(Node *node, MlirOperation operation);
MlirValue lookupMappedValue(Value *jitValue);
std::vector<MlirValue> lookupMappedValues(c10::ArrayRef<Value *> values);
MlirContext context;
std::unordered_map<Value *, MlirValue> valueMap;
};
} // namespace
void NodeImporter::importPrimNode(Node *node, MlirBlock appendToBlock) {
TypeMapper typeMapper(context);
MlirLocation loc = getMlirLocationFromNode(context, node);
auto kind = node->kind();
auto createAndMapTrivialNode = [&](Node *node, const std::string &opName) {
MlirOperation operation =
createMlirOperationAtEnd(appendToBlock, opName, loc,
getMlirTypesFromValues(loc, node->outputs()),
lookupMappedValues(node->inputs()));
mapResults(node, operation);
};
auto createAndMapNodeWithAttribute = [&](Node *node,
const std::string &opName,
const std::string &attrName,
MlirAttribute attr) {
MlirOperation operation =
createMlirOperationAtEnd(appendToBlock, opName, loc,
getMlirTypesFromValues(loc, node->outputs()),
lookupMappedValues(node->inputs()),
toMlirNamedAttribute(attrName.c_str(), attr));
mapResults(node, operation);
};
switch (kind) {
case c10::prim::TupleIndex:
case c10::prim::TupleUnpack:
case c10::prim::ListUnpack:
case c10::prim::dtype:
case c10::prim::device:
case c10::prim::unchecked_cast:
case c10::prim::Uninitialized:
case c10::prim::RaiseException:
case c10::prim::Print:
case c10::prim::NumToTensor: {
createAndMapTrivialNode(node,
"torch.prim." + std::string(kind.toUnqualString()));
return;
}
case c10::prim::ListConstruct: {
createAndMapTrivialNode(node, "basicpy.build_list");
return;
}
case c10::prim::TupleConstruct: {
createAndMapTrivialNode(node, "basicpy.build_tuple");
return;
}
case c10::prim::GetAttr:
case c10::prim::SetAttr: {
createAndMapNodeWithAttribute(
node, "torch.prim." + std::string(kind.toUnqualString()), "name",
importAttribute(loc, node, c10::attr::name));
return;
}
}
if (kind == c10::prim::Constant) {
auto output = node->output();
MlirOperation op;
OpBuilder builder(context);
if (output->type()->cast<c10::NoneType>()) {
op = builder.createNoneConstant(loc);
} else if (output->type()->cast<c10::BoolType>()) {
op = builder.createBoolConstant(
loc, static_cast<bool>(node->i(c10::attr::value)));
} else if (output->type()->cast<c10::StringType>()) {
// TODO: Are TorchScript strings bytes or str technically?
// For now, model it as bytes to avoid pledging more than we currently
// model (e.g. no unicode, etc.).
op = builder.createBytesConstant(loc, node->s(c10::attr::value));
} else if (auto functionType = output->type()->cast<c10::FunctionType>()) {
torch::jit::Function *function = functionType->function();
const std::string &symName = function->qualname().qualifiedName();
op = createMlirOperation(
"std.constant", loc,
getFunctionTypeFromSchema(context, function->getSchema()),
toMlirNamedAttribute(
"value",
mlirFlatSymbolRefAttrGet(context, toMlirStringRef(symName))));
} else {
MlirAttribute valueAttr = importAttribute(loc, node, c10::attr::value);
op = builder.createStdConstant(loc, valueAttr);
}
mlirBlockAppendOwnedOperation(appendToBlock, op);
mapResults(node, op);
return;
}
if (kind == c10::prim::Loop) {
std::vector<MlirType> resultTypes =
getMlirTypesFromValues(loc, node->outputs());
MlirOperation operation = createMlirOperationAtEnd(
appendToBlock, "torch.prim.Loop", loc, resultTypes,
lookupMappedValues(node->inputs().slice(0, 2)),
derefineValues(lookupMappedValues(node->inputs().slice(2)), resultTypes,
loc, appendToBlock),
mlirRegionCreate());
mapResults(node, operation);
std::vector<MlirType> terminatorOperandTypes = {npcompBoolTypeGet(context)};
terminatorOperandTypes.insert(terminatorOperandTypes.end(),
resultTypes.begin(), resultTypes.end());
auto createTerminator = [&](c10::ArrayRef<MlirValue> yieldedValues,
MlirBlock appendToBlock) {
createMlirOperationAtEnd(appendToBlock, "torch.prim.Loop.condition", loc,
derefineValues(yieldedValues,
terminatorOperandTypes, loc,
appendToBlock));
};
mlirRegionAppendOwnedBlock(
mlirOperationGetRegion(operation, 0),
importBlock(node->blocks()[0], createTerminator));
return;
}
if (kind == c10::prim::If) {
// TorchScript will already have an explicit op to determine truthiness. So
// all we need to do here is launder !basicpy.BoolType to i1 for `scf.if`.
MlirOperation pred = createMlirOperationAtEnd(
appendToBlock, "basicpy.bool_cast", loc, mlirIntegerTypeGet(context, 1),
lookupMappedValue(node->input()));
std::vector<MlirType> resultTypes =
getMlirTypesFromValues(loc, node->outputs());
MlirOperation operation = createMlirOperationAtEnd(
appendToBlock, "scf.if", loc, mlirOperationGetResult(pred, 0),
resultTypes, mlirRegionCreate(), mlirRegionCreate());
mapResults(node, operation);
auto createTerminator =
[&](c10::ArrayRef<MlirValue> yieldedValues, MlirBlock appendToBlock) {
createMlirOperationAtEnd(
appendToBlock, "scf.yield", loc,
derefineValues(yieldedValues, resultTypes, loc, appendToBlock));
};
mlirRegionAppendOwnedBlock(
mlirOperationGetRegion(operation, 0),
importBlock(node->blocks()[0], createTerminator));
mlirRegionAppendOwnedBlock(
mlirOperationGetRegion(operation, 1),
importBlock(node->blocks()[1], createTerminator));
return;
}
if (kind == c10::prim::CallMethod) {
auto classType = node->input(0)->type()->cast<c10::ClassType>();
auto methodName = node->s(c10::attr::name);
torch::jit::Function *function = classType->findMethod(methodName);
torch::jit::Block *calleeEntryBlock = function->graph()->block();
auto expectedTypes = c10::fmap(calleeEntryBlock->inputs(), [&](Value *v) {
return typeMapper.mapFromTorchType(loc, v->type());
});
MlirOperation operation = createMlirOperationAtEnd(
appendToBlock, "torch.prim.CallMethod", loc,
getMlirTypesFromValues(loc, node->outputs()),
derefineValues(lookupMappedValues(node->inputs()), expectedTypes, loc,
appendToBlock),
toMlirNamedAttribute("name",
importAttribute(loc, node, c10::attr::name)));
mapResults(node, operation);
return;
}
if (kind == c10::prim::CallFunction) {
auto functionType = node->input(0)->type()->cast<c10::FunctionType>();
torch::jit::Block *calleeEntryBlock =
functionType->function()->graph()->block();
auto expectedTypes = c10::fmap(calleeEntryBlock->inputs(), [&](Value *v) {
return typeMapper.mapFromTorchType(loc, v->type());
});
MlirOperation operation = createMlirOperationAtEnd(
appendToBlock, "std.call_indirect", loc,
getMlirTypesFromValues(loc, node->outputs()),
lookupMappedValue(node->input(0)),
derefineValues(lookupMappedValues(node->inputs().slice(1)),
expectedTypes, loc, appendToBlock));
mapResults(node, operation);
return;
}
// Unhandled.
{
std::stringstream msg;
msg << "unhandled prim operation: ";
node->print(msg, 0, nullptr);
mlirEmitError(getMlirLocationFromNode(context, node), msg.str().c_str());
throw mlir_diagnostic_emitted();
}
}
void NodeImporter::importKernelCall(Node *node, MlirBlock appendToBlock) {
TypeMapper typeMapper(context);
MlirLocation loc = getMlirLocationFromNode(context, node);
KernelCallBuilder kcb(context, loc, node->kind().toQualString(),
node->schema());
for (MlirValue value : lookupMappedValues(node->inputs())) {
kcb.addOperand(value);
}
for (MlirType type : getMlirTypesFromValues(loc, node->outputs())) {
kcb.addResultType(type);
}
MlirOperation op = kcb.create();
mlirBlockAppendOwnedOperation(appendToBlock, op);
mapResults(node, op);
}
void NodeImporter::importNode(Node *node, MlirBlock appendToBlock) {
if (node->kind().ns() == c10::namespaces::prim) {
importPrimNode(node, appendToBlock);
return;
}
if (node->maybeSchema()) {
importKernelCall(node, appendToBlock);
return;
}
{
std::stringstream msg;
msg << "unhandled: generic operation: ";
node->print(msg, 0, nullptr);
mlirEmitError(getMlirLocationFromNode(context, node), msg.str().c_str());
throw mlir_diagnostic_emitted();
}
}
MlirBlock NodeImporter::importBlock(Block *jitBlock,
CreateTerminatorFn createTerminator) {
MlirBlock block = createBlockFor(jitBlock);
for (Node *node : jitBlock->nodes()) {
importNode(node, block);
}
Node *returnNode = jitBlock->return_node();
createTerminator(lookupMappedValues(returnNode->inputs()), block);
return block;
}
MlirBlock NodeImporter::createBlockFor(Block *jitBlock) {
Node *paramNode = jitBlock->param_node();
MlirLocation loc = getMlirLocationFromNode(context, paramNode);
std::vector<MlirType> blockArgTypes =
getMlirTypesFromValues(loc, paramNode->outputs());
MlirBlock block = mlirBlockCreate(blockArgTypes.size(), blockArgTypes.data());
for (int i = 0, e = mlirBlockGetNumArguments(block); i < e; i++) {
Value *jitValue = paramNode->outputs()[i];
MlirValue value = mlirBlockGetArgument(block, i);
mapValue(jitValue, value);
}
return block;
}
void NodeImporter::mapValue(Value *jitValue, MlirValue value) {
auto it = valueMap.find(jitValue);
(void)it;
assert(it == valueMap.end() && "jitValue has already been mapped");
valueMap[jitValue] = value;
}
void NodeImporter::mapResults(Node *node, MlirOperation operation) {
assert(node->outputs().size() ==
(size_t)mlirOperationGetNumResults(operation));
for (int i = 0, e = node->outputs().size(); i < e; i++) {
mapValue(node->outputs()[i], mlirOperationGetResult(operation, i));
}
}
MlirValue NodeImporter::lookupMappedValue(Value *jitValue) {
auto it = valueMap.find(jitValue);
assert(it != valueMap.end() &&
"trying to get mapping for jitValue that is not mapped yet!");
return it->second;
}
std::vector<MlirValue>
NodeImporter::lookupMappedValues(c10::ArrayRef<Value *> values) {
std::vector<MlirValue> ret;
for (Value *value : values) {
ret.push_back(lookupMappedValue(value));
}
return ret;
}
MlirBlock torch_mlir::importBlock(MlirContext context, Block *jitBlock,
CreateTerminatorFn createTerminator) {
NodeImporter importer(context);
return importer.importBlock(jitBlock, createTerminator);
}