colin
/
onnx-mlir
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[MLIR] Legalize output tensors' names (#368) 

* [MLIR] legalize output tensor's name

* Update frontend_dialect_transformer.cpp
This commit is contained in:
TUNG LEDUC 2019-11-15 01:11:05 +09:00 committed by Tian Jin
parent 63596e723f
commit 0d644fab92
1 changed files with 6 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
src/builder/frontend_dialect_transformer.cpp
View File

@ -238,16 +238,19 @@ class FrontendGenImpl {
} }
void ImportOutputTensor(onnx::ValueInfoProto& output) { void ImportOutputTensor(onnx::ValueInfoProto& output) {
if (frontend_symbols_.ContainKey(legalize_name(output.name()))) { auto output_tensor_legalized_name = legalize_name(output.name());
if (frontend_symbols_.ContainKey(output_tensor_legalized_name)) {
mlir::OperationState result( mlir::OperationState result(
UnknownLoc(), "frontend.output " + output.name()); UnknownLoc(), "frontend.output " + output_tensor_legalized_name);
mlir::Type elementType = mlir::Type elementType =
TypeConvert(output.type().tensor_type().elem_type()); TypeConvert(output.type().tensor_type().elem_type());
result.addTypes(mlir::UnrankedTensorType::get(elementType)); result.addTypes(mlir::UnrankedTensorType::get(elementType));
result.addOperands(frontend_symbols_.GetTensorByOnnxName(output.name())); result.addOperands(frontend_symbols_.GetTensorByOnnxName(
output_tensor_legalized_name));
builder_.createOperation(result); builder_.createOperation(result);
} else { } else {
// TODO: Why not in the symbol table? something is wrong // TODO: Why not in the symbol table? something is wrong
assert(false && "output name not found");
} }
} }