Registering new compiler backend in Pytorch2.0

AR51 · March 3, 2023, 10:29am

Hi,
I want to know how to register the new backend in pytorch 2.0 at the dynamo stage to get the Aten/prim IR and take control of the graph.
Is there any sample document or same code available in the repository.

AR51 · March 7, 2023, 1:10pm

Can anyone help me on this.

jansel · March 8, 2023, 11:14pm

You can call this function to register a backend:

github.com

pytorch/pytorch/blob/32ffd70644d8593e0e126d4df62df27a61e23619/torch/_dynamo/backends/registry.py#L19


      
          class CompiledFn(Protocol):
              def __call__(self, *args: torch.Tensor) -> Tuple[torch.Tensor, ...]:
                  ...
          
          

          
CompilerFn = Callable[[fx.GraphModule, List[torch.Tensor]], CompiledFn]
          
          
_BACKENDS: Dict[str, CompilerFn] = dict()
          
          

          
def register_backend(
              compiler_fn: Optional[CompilerFn] = None,
              name: Optional[str] = None,
              tags: Sequence[str] = (),
          ):
              """
              Decorator to add a given compiler to the registry to allow calling
              `torch.compile` with string shorthand.  Note: for projects not
              imported by default, it might be easier to pass a function directly
              as a backend and not use a string.

This allows you to use a string alias for the backend=... arg to torch.compile(.., backend="some_name"). You can also just pass your backend directly to torch.compile(..., backend=my_backend), no need to register.

To get the aten/prim graph use:

github.com

pytorch/pytorch/blob/32ffd70644d8593e0e126d4df62df27a61e23619/torch/_dynamo/backends/common.py#L14


      
          import torch
          from torch._dynamo import eval_frame
          from torch._dynamo.utils import counters
          from torch._functorch.aot_autograd import aot_module_simplified
          from torch._subclasses import FakeTensor
          from torch.utils._python_dispatch import _disable_current_modes
          
          
log = logging.getLogger(__name__)
          
          

          
def aot_autograd(**kwargs):
              def compiler_fn(gm: torch.fx.GraphModule, example_inputs):
                  import functorch.compile
          
          
        # Hack to get around circular import problems with aot_eager_decomp_partition
                  if callable(kwargs.get("decompositions")):
                      kwargs["decompositions"] = kwargs["decompositions"]()
          
          
        # TODO: stop monkeypatching here (without even cleaning up, UGH!)
                  functorch.compile.config.use_functionalize = True
                  functorch.compile.config.use_fake_tensor = True

AR51 · March 9, 2023, 1:17pm

Thanks for the reply.
How to dump the primIR graph?

SherlockNoMad · March 10, 2023, 8:11am

You might also find this notebook helpful.

AR51 · March 20, 2023, 10:07am

I got. I was able to get to it. Thanks for the help.
Another question is how we are doing mid-layer integration at the Aten FX graph layer.
I’ve my own backend implemented in the cpp with my own data structures. I want to convert the aten FX graph into my own data structure graph which my backend can emit the binary out and run on the respective hardware. Can you give some pointers on that?? How are we providing control from python to cpp and its integration?

I’ve posted the question in: How do we do mid layer integration after Aten fx graph

Topic		Replies	Views
Add a new compiler backend for jit.trace torchscript	1	490	May 25, 2023
Call for backward compatability to enable users to understand and adapt to pytorch compiler compiler	2	440	November 18, 2023
[RFC] New Python operator registration API	10	415	January 31, 2024
Torch.compile can be debugged now! compiler	2	861	October 24, 2023
I build a decompiler to convert bytecode generated by dynamo into readable source code! compiler	4	485	August 29, 2023

Registering new compiler backend in Pytorch2.0

Related Topics