Trying to locate body of function torch.compile()

torch.compile
1

I see following functions available on torch library:

pwd
/usr/local/lib64/python3.9/site-packages/torch
init.py:1610:def compile(model: Optional[Callable] = None, *,
_inductor/init.py:8:def compile(
_inductor/codecache.py:936: def compile(cls, graph, source_code, cuda):
compiler/init.py:13:def compile(*args, **kwargs):
distributed/_spmd/api.py:498:def compile(
nn/modules/module.py:2536: def compile(self, *args, **kwargs):
onnx/_internal/onnxruntime.py:882: 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/init.py:13: print(“GH: compile()-2”)
grep: _inductor/pycache/init.cpython-39.pyc: binary file matches
compiler/init.py:17: print(“GH: compile()”)
grep: compiler/pycache/init.cpython-39.pyc: binary file matches
distributed/_spmd/api.py:503: print(“GH: compile()-3”)

2

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