fx node in c++

benchmarks/dynamo/microbenchmarks/fx_microbenchmarks.py

@@ -2,7 +2,7 @@ import timeit
 
 import torch.fx
 
-N = 100000
+N = 10000#0
 K = 1000
 
 

build_variables.bzl

@@ -827,6 +827,7 @@ libtorch_python_core_sources = [
     "torch/csrc/dynamo/guards.cpp",
     "torch/csrc/dynamo/init.cpp",
     "torch/csrc/functorch/init.cpp",
+    "torch/csrc/fx/node.cpp",
     "torch/csrc/mps/Module.cpp",
     "torch/csrc/mtia/Module.cpp",
     "torch/csrc/inductor/aoti_runner/pybind.cpp",

torch/_C/__init__.pyi.in

@@ -2331,3 +2331,7 @@ def _save_pickle(obj: Any) -> bytes: ...
 # Defined in torch/csrc/jit/runtime/static/init.cpp
 def _jit_to_static_module(graph_or_module: Union[Graph,ScriptModule]) -> Any: ...
 def _fuse_to_static_module(graph_or_module: Union[Graph,ScriptModule], min_size: _int) -> Any: ...
+
+# Defined in torch/csrc/fx/node.cpp
+class NodeBase:
+    pass

torch/csrc/Module.cpp

@@ -67,6 +67,7 @@
 #include <torch/csrc/cpu/Module.h>
 #include <torch/csrc/dynamo/init.h>
 #include <torch/csrc/functorch/init.h>
+#include <torch/csrc/fx/node.h>
 #include <torch/csrc/inductor/aoti_runner/pybind.h>
 #include <torch/csrc/jit/python/init.h>
 #include <torch/csrc/jit/python/python_ir.h>
@@ -1602,6 +1603,7 @@ PyObject* initModule() {
   THPDevice_init(module);
   THPStream_init(module);
   THPEvent_init(module);
+  NodeBase_init(module);
   ASSERT_TRUE(THPVariable_initModule(module));
   ASSERT_TRUE(THPFunction_initModule(module));
   ASSERT_TRUE(THPEngine_initModule(module));

torch/csrc/fx/node.cpp

@@ -0,0 +1,98 @@
+#include <torch/csrc/fx/node.h>
+
+#include <structmember.h>
+
+struct NodeBase {
+  PyObject_HEAD int x;
+};
+
+static PyObject* NodeBase_new(
+    PyTypeObject* type,
+    PyObject* args,
+    PyObject* kwds) {
+  PyObject* self = type->tp_alloc(type, 0);
+  if (!self)
+    return nullptr;
+  return self;
+}
+
+static int Nodebase_init(NodeBase* self, PyObject* args, PyObject* kwds) {
+  self->x = 5;
+  return 0;
+}
+
+static void NodeBase_dealloc(NodeBase* self) {
+  // PyObject_GC_UnTrack(self);
+  Py_TYPE(self)->tp_free((PyObject*)self);
+}
+
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
+static struct PyGetSetDef NodeBase_properties[] = {
+    {nullptr} /* Sentinel */
+};
+
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
+static struct PyMemberDef NodeBase_members[] = {
+    {"x", offsetof(NodeBase, x), T_INT, 0, nullptr},
+    {nullptr} /* Sentinel */
+};
+
+static PyTypeObject NodeBaseType = {
+    PyVarObject_HEAD_INIT(nullptr, 0) "torch._C.NodeBase", /* tp_name */
+    sizeof(NodeBase), /* tp_basicsize */
+    0, /* tp_itemsize */
+    (destructor)NodeBase_dealloc, /* tp_dealloc */
+    0, /* tp_vectorcall_offset */
+    nullptr, /* tp_getattr */
+    nullptr, /* tp_setattr */
+    nullptr, /* tp_reserved */
+    nullptr, /* tp_repr */
+    nullptr, /* tp_as_number */
+    nullptr, /* tp_as_sequence */
+    nullptr, /* tp_as_mapping */
+    nullptr, /* tp_hash  */
+    nullptr, /* tp_call */
+    nullptr, /* tp_str */
+    nullptr, /* tp_getattro */
+    nullptr, /* tp_setattro */
+    nullptr, /* tp_as_buffer */
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
+    nullptr, /* tp_doc */
+    nullptr, /* tp_traverse */
+    nullptr, /* tp_clear */
+    nullptr, /* tp_richcompare */
+    0, /* tp_weaklistoffset */
+    nullptr, /* tp_iter */
+    nullptr, /* tp_iternext */
+    nullptr, /* tp_methods */
+    NodeBase_members, /* tp_members */
+    NodeBase_properties, /* tp_getset */
+    nullptr, /* tp_base */
+    nullptr, /* tp_dict */
+    nullptr, /* tp_descr_get */
+    nullptr, /* tp_descr_set */
+    0, /* tp_dictoffset */
+    (initproc)NodeBase_init, /* tp_init */
+    nullptr, /* tp_alloc */
+    NodeBase_new, /* tp_new */
+};
+
+/*
+static PyModuleDef _module = {
+  PyModuleDef_HEAD_INIT,
+  "torch._C.fx",
+  "Module containing C++ implementations",
+  -1,
+  NULL
+};*/
+
+bool NodeBase_init(PyObject* module) {
+  if (PyType_Ready(&NodeBaseType) < 0) {
+    return false;
+  }
+  Py_INCREF(&NodeBaseType);
+  if (PyModule_AddObject(module, "NodeBase", (PyObject*)&NodeBaseType) != 0) {
+    return false;
+  }
+  return true;
+}

torch/csrc/fx/node.h

@@ -0,0 +1,5 @@
+#pragma once
+
+#include <torch/csrc/python_headers.h>
+
+bool NodeBase_init(PyObject* module);

torch/fx/node.py

@@ -11,6 +11,7 @@ import inspect
 import warnings
 from torch.fx.operator_schemas import normalize_function, normalize_module, ArgsKwargsPair
 from .._ops import ops as _ops
+from torch._C import NodeBase
 
 if TYPE_CHECKING:
     from .graph import Graph
@@ -138,8 +139,9 @@ def _format_arg(arg, max_list_len=float('inf')) -> str:
     else:
         return str(arg)
 
+
 @compatibility(is_backward_compatible=True)
-class Node:
+class Node(NodeBase):
     """
     ``Node`` is the data structure that represents individual operations within
     a ``Graph``. For the most part, Nodes represent callsites to various entities,
@@ -197,6 +199,7 @@ class Node:
                 annotation of values in the generated code or for other types
                 of analyses.
         """
+        super().__init__()
         self.graph = graph
         self.name = name  # unique name of value being created
         assert op in ['placeholder', 'call_method', 'call_module', 'call_function', 'get_attr', 'output', 'root']