so with great interest, I heard @ezyang ‘s comments about torch.compile as a toolkit at PTC.
I’ve been trying to access/manipulate some of the things that dynamo produces. So there is a lot of material on what to do with the fx traces, but to my mind, there are two more important bits it produces:
there is a cache entry of type `torch._C._dynamo.eval_frame._CacheEntry`
the guards, which I can view from the `cache_entry.guard_manager`, but I haven’t found a way to manipulate easily,
the code to actually collect the fx inputs and do whatever with the outputs. My understanding is that the `cache_entry.code` part is executed instead of the frame of the compiled model.
I would appreciate a hint / pointer to docs, source code etc. or even a quick hint if it’s not currently possible to manipulate those.
If you could provide more details as to exactly how you want to manipulate guards or dynamo-generated bytecode, then I could provide more concrete guidance.
We don’t have a lot of public API to manipulate guards since it’s often tricky to manipulate guards will still maintaining correctness. There are options such as guard_filter_fn but this is an unsafe API with no BC guarantees (see https://docs.pytorch.org/docs/stable/generated/torch.compile.html).
so basically, my aim is to take a cache entry and generate a new one (copy with slight modifications), in particular for the guards and possibly also how to get the values fed into the fx graph (i.e. the .code). I would then love to manually run those two bits after modification.
I’ll take a closer look a the code, but I suspect it’s not quite there yet.
So you want to generate a new cache entry with custom guards and custom optimized bytecode, and you would like to register that back to the Dynamo cache? Could you provide more details as to what custom guard behavior you would like and what custom bytecode you want to run?