From 983ce1aa109b3aed1f0de441195e0e1aea49f30e Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Thu, 27 Jul 2023 16:27:23 +0200 Subject: [PATCH] Import node debugName (FX graph node name) --- .../Torch/Transforms/GlobalizeObjectGraph.cpp | 1 + .../Torch/Transforms/InlineGlobalSlots.cpp | 6 ++++-- .../importer/jit_ir/csrc/node_importer.cpp | 18 ++++++++++++++++++ 3 files changed, 23 insertions(+), 2 deletions(-) diff --git a/lib/Dialect/Torch/Transforms/GlobalizeObjectGraph.cpp b/lib/Dialect/Torch/Transforms/GlobalizeObjectGraph.cpp index da8be9b17e0b..74ecae3ac8d9 100644 --- a/lib/Dialect/Torch/Transforms/GlobalizeObjectGraph.cpp +++ b/lib/Dialect/Torch/Transforms/GlobalizeObjectGraph.cpp @@ -524,6 +524,7 @@ static LogicalResult rewriteMonomorphizedFuncClone( auto newOp = OpBuilder(op).create( op.getLoc(), op.getType(), objectGraphInfo.getGlobalSlotFor(affectedSlot).getSymName()); + newOp->setAttr("FXOutputName", op->getAttr("FXOutputName")); op.replaceAllUsesWith(&*newOp); } toErase.push_back(op); diff --git a/lib/Dialect/Torch/Transforms/InlineGlobalSlots.cpp b/lib/Dialect/Torch/Transforms/InlineGlobalSlots.cpp index 76b57fe8c9a3..1d7900ab2b7a 100644 --- a/lib/Dialect/Torch/Transforms/InlineGlobalSlots.cpp +++ b/lib/Dialect/Torch/Transforms/InlineGlobalSlots.cpp @@ -375,8 +375,10 @@ class InlineGlobalSlotsPass getBackwardSliceIncludingRoot(initialValue); IRMapping mapping; OpBuilder builder(op); - for (Operation *opInSlice : slice) - builder.clone(*opInSlice, mapping); + for (Operation *opInSlice : slice) { + auto clonedOp = builder.clone(*opInSlice, mapping); + clonedOp->setAttr("FXOutputName", op->getAttr("FXOutputName")); + } auto inlinedInitialValue = mapping.lookup(initialValue); inlinedInitialValue = Torch::adjustStaticInformation( builder, op.getLoc(), inlinedInitialValue, op.getType(), diff --git a/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/node_importer.cpp b/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/node_importer.cpp index 15cffedbe834..5905662c5fae 100644 --- a/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/node_importer.cpp +++ b/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/node_importer.cpp @@ -84,6 +84,21 @@ void NodeImporter::importNode(Node *node, MlirBlock appendToBlock, MlirLocation loc = getMlirLocationFromNode(context, node); auto kind = node->kind(); + + auto outs = node->outputs(); + auto output = outs.size() == 1 ? outs[0] : nullptr; + + auto addFxOutputNameAttr = [&](MlirOperation& operation) { + if (output && output->hasDebugName()) { + std::string name = output->debugName(); + size_t len = name.size(); + if (len > 2 && name[len-2] == '.' && name[len-1] == '1') + name = name.substr(0, len-2); + auto strAttr = mlirStringAttrGet(context, toMlirStringRef(name)); + mlirOperationSetAttributeByName(operation, toMlirStringRef("FXOutputName"), strAttr); + } + }; + auto createAndMapTrivialNode = [&](Node *node, const std::string &opName, InputsTransformFn t) { std::vector mappedInputs = lookupMappedValues(node->inputs()); @@ -91,6 +106,7 @@ void NodeImporter::importNode(Node *node, MlirBlock appendToBlock, appendToBlock, opName, loc, getMlirTypesFromValues(loc, node->outputs(), importOptions), t ? t(mappedInputs) : mappedInputs); + addFxOutputNameAttr(operation); mapResults(node, operation); }; @@ -102,6 +118,7 @@ void NodeImporter::importNode(Node *node, MlirBlock appendToBlock, getMlirTypesFromValues(loc, node->outputs(), importOptions), lookupMappedValues(node->inputs()), toMlirNamedAttribute(attrName.c_str(), attr)); + addFxOutputNameAttr(operation); mapResults(node, operation); }; @@ -112,6 +129,7 @@ void NodeImporter::importNode(Node *node, MlirBlock appendToBlock, appendToBlock, loc, node->schema(), getMlirTypesFromValues(loc, node->outputs(), importOptions), lookupMappedValues(node->inputs())); + addFxOutputNameAttr(operation); mapResults(node, operation); return; }