What is fx_wrapper (WrapperFxCodegen) for, and how to hook a custom FX consumer?

torch.compile
1

Hi all,

I’m trying to understand Inductor’s fx_wrapper backend (torch._inductor.config.fx_wrapper / WrapperFxCodegen) but can’t find any docs for it, so have a few short questions:

  1. What is fx_wrapper intended for? And how is it going? Seems it was added in version 2.9, but not any released notes as cpp_wrapper in 2.5, will it be a long-term feature?
  2. If the gm is just run eagerly (gm.forward), as it currently be, what’s the advantage over PythonWrapperCodegen and the C++ wrapper? I think it will produce more overhead than cpp warpper or python wrapper.
    def compile_graph(self, gm: GraphModule) -> Callable[..., Any]:
        """
        Converts the graph module into a runnable function. The default implementation
        is simply an interpreter calling kernels in eager mode. Derived backends can
        override this to do further compilation.
        """
        return gm.forward
  1. According to the comments, I think it might be designed to support some custom fx graph consumer so it can benefit from inductor as well. Then what’s the correct way? I can override the function compile_graph to create a new WrapperFxCodegen, but there is no way to register the Derived fx warpper codegen alone, only register_backend_for_device to register everything. So I need to intrusively replace the registered device_codegens var in torch._inductor.codegen.common, is this the correct way?
class MyFxWrapper(WrapperFxCodegen):
    def compile_graph(self, gm):
        fn = MyRuntime(gm)
        return fn

from torch._inductor.codegen.common import device_codegens
device_codegens['cuda'].fx_wrapper_codegen = MyFxWrapper

Thanks!