Can backend compilers modify guards for torch Dynamo frame

hardware-backends
1

Hi,
Is there any interface in torch Dynamo, for backend compilers to add/modify guards for a frame?
Regards,
Sujoy

2

Hey! Great question :slight_smile:

The exact mechanism does not exist today, but we have the elements needed to make this really easy to add. We recently rewrote the lifecycle of a frame to encompass a backend, and have a 1:1 lifecycle of a TracingContext, just for this.

I wrote a little quick 15 min, 10 line, example of how you could modify dynamo to allow arbitrary string guards :wink:

Now, this PR should be rejected because, aside from missing tests and docs and full integration with snapshotting - its a little too loose for my tastes. That being said, the mechanisms demonstrated all the same ones you would use.

Lmk if that makes sense, happy to elaborate further.

3

Thank you, yes something like this would be definitely useful in case the backend wants to specialize the graph further by adding more guards.

I am wondering if there would be a case where the backend may want to remove a particular guard added by Dynamo already. For example, the backend may convert a constant value to a tensor and DMA it every time, in order to avoid graph recompilation when the constant value changes. However, constants are embedded in the FX graph from Dynamo. Is there a way to modify the FX graph to convert the constant into an input tensor, if the backend wants to do so?

4

I am wondering if there would be a case where the backend may want to remove a particular guard added by Dynamo already.

Probably :slight_smile: We lack a good mechanism for this today, but maybe could add it if really pressed for it.

Is there a way to modify the FX graph to convert the constant into an input tensor, if the backend wants to do so?

You can do it in your custom backend you pass to torch._dynamo.optimize, or for export, just modify the output fx graph.

5

Yes, we can write a custom backend pass to modify the FX graph, to avoid a recompilation at the backend. However, won’t it still trigger a recompile at Dynamo level? Is there a way to avoid the recompilation at Dynamo itself?