Import FX graph node name

lib/Dialect/Torch/Transforms/GlobalizeObjectGraph.cpp

@@ -524,6 +524,8 @@ static LogicalResult rewriteMonomorphizedFuncClone(
       auto newOp = OpBuilder(op).create<GlobalSlotGetOp>(
           op.getLoc(), op.getType(),
           objectGraphInfo.getGlobalSlotFor(affectedSlot).getSymName());
+      if (auto fxOutputName = op->getAttr("FXOutputName"))
+        newOp->setAttr("FXOutputName", fxOutputName);
       op.replaceAllUsesWith(&*newOp);
     }
     toErase.push_back(op);

lib/Dialect/Torch/Transforms/InlineGlobalSlots.cpp

@@ -375,8 +375,11 @@ 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);
+        if (auto fxOutputName = op->getAttr("FXOutputName"))
+          clonedOp->setAttr("FXOutputName", fxOutputName);
+      }
       auto inlinedInitialValue = mapping.lookup(initialValue);
       inlinedInitialValue = Torch::adjustStaticInformation(
           builder, op.getLoc(), inlinedInitialValue, op.getType(),

python/torch_mlir/__init__.py

@@ -328,7 +328,8 @@ def compile(model: torch.nn.Module,
             backend_legal_ops: Optional[Sequence[str]] = None,
             extra_library: Iterable[Callable] = [],
             verbose: bool = False,
-            use_make_fx: bool = False):
+            use_make_fx: bool = False,
+            add_fx_outputname: bool = False):
     """Convert a PyTorch model to MLIR.
 
     Args:
@@ -437,6 +438,7 @@ def compile(model: torch.nn.Module,
     mb = ModuleBuilder()
     import_options = ImportOptions()
     import_options.ignoreExistingTensorShapesAndDtypes = ignore_traced_shapes
+    import_options.addFxOutputName = add_fx_outputname
     try:
         original_stderr = sys.stderr
         sys.stderr = StringIO()

python/torch_mlir/dialects/torch/importer/jit_ir/csrc/import_options.h

@@ -33,6 +33,10 @@ struct ImportOptions {
   // In that case, the appropriate shape information is provided via the type
   // bound annotations on the function arguments instead.
   bool ignoreExistingTensorShapesAndDtypes = false;
+
+  // Whether to add the FXOutputName attribute from the debug name of the jit
+  // node.
+  bool addFxOutputName = false;
 };
 } // namespace torch_mlir
 

python/torch_mlir/dialects/torch/importer/jit_ir/csrc/import_options_pybind.cpp

@@ -20,5 +20,7 @@ void torch_mlir::initImportOptionsBindings(py::module &m) {
       .def_readwrite("assumeTensorsHaveValueSemantics",
                      &ImportOptions::assumeTensorsHaveValueSemantics)
       .def_readwrite("ignoreExistingTensorShapesAndDtypes",
-                     &ImportOptions::ignoreExistingTensorShapesAndDtypes);
+                     &ImportOptions::ignoreExistingTensorShapesAndDtypes)
+      .def_readwrite("addFxOutputName",
+                     &ImportOptions::addFxOutputName);
 }

python/torch_mlir/dialects/torch/importer/jit_ir/csrc/node_importer.cpp

@@ -84,13 +84,29 @@ 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 (importOptions.addFxOutputName && 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<MlirValue> mappedInputs = lookupMappedValues(node->inputs());
     MlirOperation operation = createMlirOperationAtEnd(
         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;
   }