Fixed tutorial to be clearer and to recommend other path

Author: janeyx99

advanced_source/cpp_custom_ops.rst

@@ -64,39 +64,76 @@ Using ``cpp_extension`` is as simple as writing the following ``setup.py``:
         ext_modules=[
             cpp_extension.CppExtension(
               "extension_cpp",
-              ["muladd.cpp"]
-              py_limited_api=True)],
+              ["muladd.cpp"],
+              # define Py_LIMITED_API with min version 3.9 to expose only the stable
+              # limited API subset from Python.h
+              extra_compile_args={"cxx": ["-DPy_LIMITED_API=0x03090000"]}, 
+              py_limited_api=True)],  # Build 1 wheel across multiple Python versions
         cmdclass={'build_ext': cpp_extension.BuildExtension},
-        options={"bdist_wheel": {"py_limited_api": "cp39"}}
+        options={"bdist_wheel": {"py_limited_api": "cp39"}}  # 3.9 is minimum supported Python version
   )
 
 If you need to compile CUDA code (for example, ``.cu`` files), then instead use
 `torch.utils.cpp_extension.CUDAExtension <https://pytorch.org/docs/stable/cpp_extension.html#torch.utils.cpp_extension.CUDAExtension>`_.
 Please see `extension-cpp <https://github.com/pytorch/extension-cpp>`_ for an
 example for how this is set up.
 
-Note that you can build a single wheel for multiple CPython versions (similar to
-what you would do for pure python packages) starting with PyTorch 2.6. Specifically,
-if your custom library adheres to the `CPython Stable Limited API
-<https://docs.python.org/3/c-api/stable.html>`_ or avoids CPython entirely, you
-can build one Python agnostic wheel against a minimum supported CPython version
-through setuptools' ``py_limited_api`` flag.
+The above example represents what we refer to as a CPython agnostic wheel, meaning
+we are building a single wheel that can be run across multiple CPython versions (similar
+to pure Python packages). CPython agnosticism is desirable in minimizing the number of wheels your
+custom library needs to support and release. To achieve this, there are three key lines to note.
 
-It is necessary to specify ``py_limited_api=True`` both within ``setup``
-and also as an option to the ``"bdist_wheel"`` command with the minimal supported
-Python version (in this case, 3.9). This ``setup`` would build one wheel that could
-be installed across multiple Python versions ``python>=3.9``.
+The first is the specification of ``Py_LIMITED_API`` in ``extra_compile_args`` to the
+minimum CPython version you would like to support:
 
-.. note::
+.. code-block:: python
+  extra_compile_args={"cxx": ["-DPy_LIMITED_API=0x03090000"]},
+
+Defining the ``Py_LIMITED_API`` flag helps guarantee that the extension is in fact
+only using the `CPython Stable Limited API <https://docs.python.org/3/c-api/stable.html>`_,
+which is a requirement for the building a CPython agnostic wheel. If this requirement
+is not met, it is possible to build a wheel that looks CPython agnostic but will crash,
+or worse, be silently incorrect, in another CPython environment. Take care to avoid
+using unstable CPython APIs, for example APIs from libtorch_python (in particular
+pytorch/python bindings,) and to only use APIs from libtorch (ATen objects, operators
+and the dispatcher). We strongly recommend defining the ``Py_LIMITED_API`` flag to
+ensure the extension is compliant and safe as a CPython agnostic wheel.
+
+The second and third lines inform setuptools that you intend to build a CPython agnostic
+wheel and will influence the naming of the wheel accordingly. It is necessary to specify
+``py_limited_api=True`` as an argument to CppExtension/CUDAExtension and also as an option
+to the ``"bdist_wheel"`` command with the minimal supported CPython version (in this case,
+3.9):
+
+.. code-block:: python
+  setup(name="extension_cpp",
+        ext_modules=[
+            cpp_extension.CppExtension(
+              ...,
+              py_limited_api=True)],  # Build 1 wheel across multiple Python versions
+        ...,
+        options={"bdist_wheel": {"py_limited_api": "cp39"}}  # 3.9 is minimum supported Python version
+  )
+
+This ``setup`` would build one wheel that could be installed across multiple CPython
+versions ``>=3.9``.
+
+If your extension uses CPython APIs outside the stable limited set, then you should build
+a wheel per CPython version instead, like so:
+
+.. code-block:: python
+
+  from setuptools import setup, Extension
+  from torch.utils import cpp_extension
+
+  setup(name="extension_cpp",
+        ext_modules=[
+            cpp_extension.CppExtension(
+              "extension_cpp",
+              ["muladd.cpp"])],
+        cmdclass={'build_ext': cpp_extension.BuildExtension},
+  )
 
-  You must verify independently that the built wheel is truly Python agnostic.
-  Specifying ``py_limited_api`` does not check for any guarantees, so it is possible
-  to build a wheel that looks Python agnostic but will crash, or worse, be silently
-  incorrect, in another Python environment. Take care to avoid using unstable CPython
-  APIs, for example APIs from libtorch_python (in particular pytorch/python bindings,)
-  and to only use APIs from libtorch (ATen objects, operators and the dispatcher).
-  For example, to give access to custom ops from Python, the library should register
-  the ops through the dispatcher (covered below!).
 
 Defining the custom op and adding backend implementations
 ---------------------------------------------------------
@@ -241,15 +278,74 @@ matters (importing in the wrong order will lead to an error).
 
 To use the custom operator with hybrid Python/C++ registrations, we must
 first load the C++ library that holds the custom operator definition
-and then call the ``torch.library`` registration APIs. This can happen in one
-of two ways:
+and then call the ``torch.library`` registration APIs. This can happen in
+three ways:
+
+
+1. The first way to load the C++ library that holds the custom operator definition
+   is to define a dummy Python module for _C. Then, in Python, when you import the
+   module with ``import _C``, the ``.so``s corresponding to the extension will be
+   loaded and the ``TORCH_LIBRARY`` and ``TORCH_LIBRARY_IMPL`` static initializers
+   will run. One can create a dummy Python module with ``PYBIND11_MODULE`` like below,
+   but you will notice that this does not compile with ``Py_LIMITED_API``, because
+   ``pybind11`` does not promise to only use the stable limited CPython API! With
+   the below code, you sadly cannot build a CPython agnostic wheel for your extension!
+   (Foreshadowing: I wonder what the second way is ;)).
+
+.. code-block:: cpp
+  // in, say, not_agnostic/csrc/extension_BAD.cpp
+  #include <pybind11/pybind11.h>
 
+  PYBIND11_MODULE("_C", m) {}
 
-1. In this tutorial, our C++ custom operator is located in a shared library object,
-   and we use ``torch.ops.load_library("/path/to/library.so")`` to load it. This
-   is the blessed path for Python agnosticism, and you will not have a Python C
-   extension module to import. See our `extension_cpp/__init__.py <https://github.com/pytorch/extension-cpp/blob/e4c4eb822889ea67f191071fa627d750e04bf047/extension_cpp/__init__.py>`_
-   for an example:
+.. code-block:: python
+
+  # in, say, extension/__init__.py
+  from . import _C
+
+2. In this tutorial, because we value being able to build a single wheel across multiple
+   CPython versions, we will replace the unstable ``PYBIND11`` call with stable API calls.
+   The below code compiles with ``-DPy_LIMITED_API=0x03090000`` and successfully creates
+   a dummy Python module for our ``_C`` extension so that it can be imported from Python.
+   See `extension_cpp/__init__.py <https://github.com/pytorch/extension-cpp/blob/master/extension_cpp/__init__.py>`_
+   and `extension_cpp/csrc/muladd.cpp  <https://github.com/pytorch/extension-cpp/blob/master/extension_cpp/csrc/muladd.cpp>`_
+   for more details:
+
+.. code-block:: cpp
+  #include <Python.h>
+
+  extern "C" {
+    /* Creates a dummy empty _C module that can be imported from Python.
+      The import from Python will load the .so consisting of this file
+      in this extension, so that the TORCH_LIBRARY static initializers
+      below are run. */
+    PyObject* PyInit__C(void)
+    {
+        static struct PyModuleDef module_def = {
+            PyModuleDef_HEAD_INIT,
+            "_C",   /* name of module */
+            NULL,   /* module documentation, may be NULL */
+            -1,     /* size of per-interpreter state of the module,
+                      or -1 if the module keeps state in global variables. */
+            NULL,   /* methods */
+        };
+        return PyModule_Create(&module_def);
+    }
+  }
+
+.. code-block:: python
+
+  # in, say, extension/__init__.py
+  from . import _C
+
+3. If you want to avoid ``Python.h`` entirely in your C++ custom operator, you may
+   use ``torch.ops.load_library("/path/to/library.so")`` in Python to load the ``.so``
+   file(s) compiled from the extension. Note that, with this method, there is no ``_C``
+   Python module created for the extension so you cannot call ``import _C`` from Python.
+   Instead of relying on the import statement to trigger the custom operators to be
+   registered, ``torch.ops.load_library("/path/to/library.so")`` will do the trick.
+   The challenge then is shifted towards understanding where the ``.so`` files are
+   located so that you can load them, which is not always trivial:
 
 .. code-block:: python
 
@@ -265,22 +361,6 @@ of two ways:
   from . import ops
 
 
-2. You may also see other custom extensions importing the Python C extension module.
-   The module would be created in C++ and then imported in Python, like the code below.
-   This code is not guaranteed to use the stable limited CPython API and would block
-   your extension from building a Python-agnostic wheel! AVOID the following:
-
-.. code-block:: cpp
-
-  // in, say, not_agnostic/csrc/extension_BAD.cpp
-  PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {}
-
-.. code-block:: python
-
-  # in, say, extension_BAD/__init__.py
-  from . import _C
-
-
 Adding training (autograd) support for an operator
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 Use ``torch.library.register_autograd`` to add training support for an operator. Prefer