mirror of https://github.com/llvm/torch-mlir
57 lines
2.1 KiB
C++
57 lines
2.1 KiB
C++
|
//===- function_importer.cpp ----------------------------------------------===//
|
||
|
|
//
|
||
|
|
// This file is licensed under a pytorch-style license
|
||
|
|
// See frontends/pytorch/LICENSE for license information.
|
||
|
|
//
|
||
|
|
//===----------------------------------------------------------------------===//
|
||
|
|
|
||
|
|
#include "function_importer.h"
|
||
|
|
|
||
|
|
#include <unordered_map>
|
||
|
|
|
||
|
|
#include "mlir_utils.h"
|
||
|
|
#include "torch_to_mlir_utils.h"
|
||
|
|
|
||
|
|
#include "mlir-c/BuiltinAttributes.h"
|
||
|
|
#include "mlir-c/BuiltinTypes.h"
|
||
|
|
#include "mlir-c/Diagnostics.h"
|
||
|
|
|
||
|
|
namespace py = pybind11;
|
||
|
|
using namespace torch_mlir;
|
||
|
|
|
||
|
|
MlirOperation
|
||
|
|
torch_mlir::importJitFunctionAsFuncOp(MlirContext context,
|
||
|
|
torch::jit::Function *function) {
|
||
|
|
// Useful for debugging:
|
||
|
|
// graph->dump();
|
||
|
|
MlirLocation loc = mlirLocationUnknownGet(context);
|
||
|
|
MlirType functionType =
|
||
|
|
getFunctionTypeFromSchema(context, function->getSchema());
|
||
|
|
// Use the function's qualified name from the compilation unit.
|
||
|
|
// This is a stable linkage name that matches Python module lookup
|
||
|
|
// conventions (see compilation unit import in IValueImporter for more details
|
||
|
|
// on qualified names).
|
||
|
|
MlirAttribute symNameAttr = mlirStringAttrGet(
|
||
|
|
context, toMlirStringRef(function->qualname().qualifiedName()));
|
||
|
|
MlirOperation func = createMlirOperation(
|
||
|
|
"func", loc, mlirRegionCreate(),
|
||
|
|
toMlirNamedAttribute("type", mlirTypeAttrGet(functionType)),
|
||
|
|
toMlirNamedAttribute("sym_name", symNameAttr));
|
||
|
|
MlirRegion bodyRegion = mlirOperationGetRegion(func, 0);
|
||
|
|
std::vector<MlirType> resultTypes;
|
||
|
|
for (int i = 0, e = mlirFunctionTypeGetNumResults(functionType); i != e;
|
||
|
|
i++) {
|
||
|
|
resultTypes.push_back(mlirFunctionTypeGetResult(functionType, i));
|
||
|
|
}
|
||
|
|
auto createTerminator = [&](c10::ArrayRef<MlirValue> yieldedValues,
|
||
|
|
MlirBlock appendToBlock) {
|
||
|
|
createMlirOperationAtEnd(
|
||
|
|
appendToBlock, "std.return", loc,
|
||
|
|
derefineValues(yieldedValues, resultTypes, loc, appendToBlock));
|
||
|
|
};
|
||
|
|
MlirBlock block =
|
||
|
|
importBlock(context, function->graph()->block(), createTerminator);
|
||
|
|
mlirRegionAppendOwnedBlock(bodyRegion, block);
|
||
|
|
return func;
|
||
|
|
}
|