Trying to locate body of function torch.compile()

I see following functions available on torch library:

/usr/local/lib64/python3.9/site-packages/torch compile(model: Optional[Callable] = None, *,
_inductor/ compile(
_inductor/ def compile(cls, graph, source_code, cuda):
compiler/ compile(*args, **kwargs):
distributed/_spmd/ compile(
nn/modules/ def compile(self, *args, **kwargs):
onnx/_internal/ def compile(self, graph_module: torch.fx.GraphModule, args) → torch.fx.GraphModule:

For each one, I instrumented code to insert debug statement to see which one is called by torch.compile() but nothing prints out. Am I missing something? Or is it because it is cached somewhere?

_inductor/ print(“GH: compile()-2”)
grep: _inductor/pycache/init.cpython-39.pyc: binary file matches
compiler/ print(“GH: compile()”)
grep: compiler/pycache/init.cpython-39.pyc: binary file matches
distributed/_spmd/ print(“GH: compile()-3”)

torch.compile won’t actually compile anything until the first time you invoke your compiled function/module. This is because torch.compile will (potentially) specialize on your particular sample inputs to generate efficient code.

If you want to e.g. put a breakpoint in inductor and look at the graph that inductor is compiling, you’d need to actually invoke the compiled function:

def f(x):
    return x.sin() + x.cos()

x = torch.ones(4)
compiled_f = torch.compile(f)
# run compiled_f, which will actually invoke all of the logic in torch.compile:
out = compiled_f(x)
1 Like