When does the inductor code run?

Hello, I am new to torch.compile and want to study the inductor.
I wonder when does the inductor’s optimization code run?

for example, there are main.py code here.

def f(x):
    return torch.sin(x)**2 + torch.cos(x)**2

# first spot of main.py
compiled_f = torch.compile(f, backend='inductor', options={'trace.enabled':True, 'trace.graph_diagram':True})

# second spot of main.py
output_from_compiled_f = compiled_f(some_input)

Thanks to the torch.compile, my sin^2+cos^2 function’s nodes are merged like,

I wonder when does this node-merge(inductor’s optimization logic) happen?

In this situation, I insert some print() in source code of dynamo and inductor for debug purposes.
for example 1,

@ torch/_inductor/compile_fx.py
def compile_fx(
    model_: torch.fx.GraphModule,
    example_inputs_: List[torch.Tensor],
    inner_compile: Callable[..., Any] = compile_fx_inner,
    config_patches: Optional[Dict[str, Any]] = None,
    decompositions: Optional[Dict[OpOverload, Callable[..., Any]]] = None,
):
    """Main entrypoint to a compile given FX graph"""
    print("_____________________________ compile_fx.py ")

for example 2,

@torch/_inductor/lowering.py 
def make_foreach_pointwise(pw_fn, allow_alpha=False):
    print("_____________________________ lowering.py ")

In this situation,

  1. example 2’s debug print happen when first spot of main.py.
  2. example 1’s debug print happen when second spot of main.py.

I am confused that why example 1’s debug print happen when second spot of main.py.
Because i think all the inductor’s optimization could happen when first spot of main.py. (some sort of compile-time of DNN)
So, example 1’s debug print should have happen when the first spot of main.py (i believe this way, but this is not the true).

The location of “example 1’s debug print” is compile_fx which is the ’ “”“Main entrypoint to a compile given FX graph”“” '. This means here is the entrance of inductor logic.

But why inductor (backend compiler) logic is start at the second spot of main.py ? (some sort of run-time of DNN)

My question is summarized as follows.
Q1. When does the inductor’s optimization logic(i.e., node merge for less memory access etc) activated ?
at the first spot of main.py or at the second spot of main.py?

Q2. What is the sequence of the inductor’s execution?
I traced the source code of torch. And i got following execution sequence information.

(1) torch/_dynamo/eval_frame.py : def optimize →
(2) torch/_dynamo/eval_frame.py : def _optimize_catch_errors →
(3) torch/_dynamo/eval_frame.py : class OptimizeContext(_TorchDynamoContext): →
(4) torch/_dynamo/eval_frame.py : class _TorchDynamoContext: 's def call → ???

but i am lost in here. where should i start again ? (for example, torch/init’s def compile)
where can i find the caller of inductor optimization code?

For Q1, the Dynamo tutorials at torch.compiler — PyTorch main documentation (Dynamo overview and Dynamo deepdive) should help explaining what’s going on there.

For Q2, I would recommend having a look at the folders torch/_dynamo and torch/_inductor, identify where the codegen is happening, and drop a breakpoint at one of those points. As discussed in the tutorials above, using TORCH_LOGS=+help should help a lot with all this.

1 Like

@Lezcano Thank you for the reply.

I think your references give general entrance of dynamo. But i need specific answer for inductor.

More specifically, your reference(https://pytorch.org/docs/main/torch.compiler_dynamo_overview.html) states that

Dynamo makes it easy to experiment with different compiler backends to make PyTorch code faster with a single line decorator torch._dynamo.optimize() which is wrapped for convenience by torch.compile()

This means that by executing this line (compiled_f = torch.compile(f, backend=‘inductor’)), we can utilize inductor backend compiler code. right?

But what i am confused is that (as i previously stated at Q2) how torch._dynamo.optimize() of at torch/_dynamo/eval_frame.py calls inductor’s code(some codes in the torch/_inductor).

I could see codes in torch/_inductor. But i can not see how these codes are called from the torch.compile() which is the start point.

Q3. Can you elaborate more specifically?
Q4. does the node-merge optimization of inductor happen at “compiled_f = torch.compile(f, backend=‘inductor’)” not at “output_from_compiled_f = compiled_f(some_input)”?

Thank you.

The two tutorials above fully answer Q1 and Q4.

For Q2 and Q3 again, I suggest you drop a few breakpoints in different files within the torch/_inductor folder and look at the backtraces. In particular, try to locate where the codegen of inductor happens (as that’s the last stage) and have a browse at the functions that call into it. As discussed, above, try also with the TORCH_LOGS env variable, as this helps understand the whole process. Have you tried this? If so, what did you try and what worked and what didn’t?

which part of the tutorial answer the Q1 and Q4?
i have read them, but i might miss somthing.

i will try TORCH_LOGS, thanks.

What i have tried was insert print codes on def&class of

  1. torch/init.py
  2. torch/_dynamo/eval_frame.py
  3. torch/_inductor/compile_fx.py
  4. torch/_inductor/lowering.py
  5. torch/_inductor/graph.py

i could see that during torch.compile() lowering.py is activated. but i don’t know which part call this def make_foreach_pointwise of lowering.py .

you said that i should go to torch/_dynamo/codegen. but i am afraid that i again lost there not knowing which part call this codegen part.

you would see my questions are basic and stupid. here is what i had tried(https://docs.google.com/document/d/1xaYOvBkqzQeILRrpAtCVunEd5CVg8PNiR7RcegcqJXc/edit ). this doc is only for me, figuring out inductor’s behavior. it is written in my language, but you can see the gist.

So, if you know the answer, help me.
regardless, i will try codegen thing and TORCH_LOGS.

thank you.

네이버 메일앱에서 보냈습니다.
보낸사람: "Lezcano via PyTorch Dev Discussions"notifications@pytorch1.discoursemail.com
받는사람: pixar0407@naver.com
참조:
날짜: 2024. 5. 15 17:56:15
제목: [PyTorch Dev Discussions] [compiler] When does the inductor code run?

| Lezcano
May 15 |

  • | - |

The two tutorials above fully answer Q1 and Q4.

For Q2 and Q3 again, I suggest you drop a few breakpoints in different files within the torch/_inductor folder and look at the backtraces. In particular, try to locate where the codegen of dynamo happens (as that’s the last stage) and have a browse at the functions that call into it. As discussed, above, try also with the TORCH_LOGS env variable, as this helps understand the whole process. Have you tried this? If so, what did you try and what worked and what didn’t?

Sorry, I meant inductor’s codegen, not dynamo’s. You can see the code we generate. Then tryt to locate the different parts of Inductor that generate this code. From there, look at the backtrace and see what other functions are called before these and explore those.

If you are not able to hit a given point of the code, like make_foreach_pointwise, try a different one. If you don’t hit it it’s because that path is not being executed with your input, so it’s not relevant to your input program.

Regarding the files you have tried, that looks as a start. Keep doing this, try other files, experiment with TORCH_LOGS, and quite soon you’ll start finding your way around inductor :slight_smile: