Registering new compiler backend in Pytorch2.0

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.

Can anyone help me on this.

You can call this function to register a backend:

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:

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

You might also find this notebook helpful.

1 Like

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