[inductor][cpu]lennard_jones, pyhpc_equation_of_state and pyhpc_isoneutral_mixing performance regression

[zxd1997066]

🐛 Describe the bug

fp32 static shape default wrapper regression in 2024-03-10

suite	name	thread	batch_size_new	speed_up_new	inductor_new	eager_new	compilation_latency_new	batch_size_old	speed_up_old	inductor_old	eager_old	compilation_latency_old	Ratio Speedup(New/old)	Eager Ratio(old/new)	Inductor Ratio(old/new)	Compilation_latency_Ratio(old/new)
torchbench	lennard_jones	single	1	1.442252	4.8045e-05	6.929299734e-05	5.983722	1.0	1.833953	3.8604e-05	7.079792161199999e-05	4.163693	0.79	1.02	0.8	0.7
torchbench	pyhpc_equation_of_state	single	1	18.71236	5.3968e-05	0.00100986864448	8.857691	1.0	23.188102	4.51e-05	0.0010457834002	6.913662	0.81	1.04	0.84	0.78

fp32 dynamic shape default wrapper regression in 2024-03-10

suite	name	thread	batch_size_new	speed_up_new	inductor_new	eager_new	compilation_latency_new	batch_size_old	speed_up_old	inductor_old	eager_old	compilation_latency_old	Ratio Speedup(New/old)	Eager Ratio(old/new)	Inductor Ratio(old/new)	Compilation_latency_Ratio(old/new)
torchbench	lennard_jones	single	1	1.431658	4.8693e-05	6.9711722994e-05	6.040358	1.0	1.775623	3.9242e-05	6.9678997766e-05	4.229882	0.81	1.0	0.81	0.7
torchbench	pyhpc_equation_of_state	single	1	18.8031	5.4025e-05	0.0010158374775	8.885931	1.0	23.240634	4.4543e-05	0.001035207560262	6.99299	0.81	1.02	0.82	0.79

fp32 static shape cpp wrapper regression in 2024-03-10

suite	name	thread	batch_size_new	speed_up_new	inductor_new	eager_new	compilation_latency_new	batch_size_old	speed_up_old	inductor_old	eager_old	compilation_latency_old	Ratio Speedup(New/old)	Eager Ratio(old/new)	Inductor Ratio(old/new)	Compilation_latency_Ratio(old/new)
torchbench	lennard_jones	single	1	1.810503	3.8294e-05	6.9331401882e-05	14.183242	1.0	2.379847	2.8666e-05	6.822069410199999e-05	12.376875	0.76	0.98	0.75	0.87
torchbench	pyhpc_equation_of_state	single	1	23.197532	4.4514e-05	0.001032614939448	17.09731	1.0	29.022945	3.5698e-05	0.00103606109061	15.172577	0.8	1.0	0.8	0.89

fp32 dynamic shape cpp wrapper regression in 2024-03-10

suite	name	thread	batch_size_new	speed_up_new	inductor_new	eager_new	compilation_latency_new	batch_size_old	speed_up_old	inductor_old	eager_old	compilation_latency_old	Ratio Speedup(New/old)	Eager Ratio(old/new)	Inductor Ratio(old/new)	Compilation_latency_Ratio(old/new)
torchbench	lennard_jones	single	1	1.76331	3.8149000000000004e-05	6.726851319e-05	14.189571	1.0	2.273742	2.8757000000000003e-05	6.5385998694e-05	12.387977	0.78	0.97	0.75	0.87
torchbench	pyhpc_equation_of_state	single	1	22.763	4.464e-05	0.00101614032	17.139039	1.0	31.649198	3.4811e-05	0.001101740231578	15.198502	0.72	1.08	0.78	0.89
torchbench	pyhpc_isoneutral_mixing	single	1	47.525991	5.6787e-05	0.0026988584509169996	20.419413	1.0	54.58171	4.9508e-05	0.0027022312986800003	18.502454	0.87	1.0	0.87	0.91

SW info

name	target_branch	target_commit	refer_branch	refer_commit
torchbench	main	1ef0a39e	main	ff42d907
torch	main	41286f1	main	581fe26
torchvision	main	0.18.0a0+2c127da	main	0.18.0a0+2c127da
torchtext	main	0.16.0a0+b0ebddc	main	0.16.0a0+b0ebddc
torchaudio	main	2.2.0a0+87aeb55	main	2.2.0a0+87aeb55
torchdata	main	0.7.1a0+0790338	main	0.7.1a0+0790338
dynamo_benchmarks	main	nightly	main	nightly

Repro:
inductor_single_run.sh
bash inductor_single_run.sh thread inference performance torchbench model float32/amp first static/dynamic default/cpp
torchbench-lennard_jones-inference-float32-static-default-single-performance-drop_guilty_commit.log
Suspected guilty commit: a7e93c3
cc @ezyang @msaroufim @bdhirsh @anijain2305 @zou3519 @chauhang @WeizhuoZhang-intel @chuanqi129

[leslie-fang-intel]

May we take a look first between the difference of generated code before and after the regression @zhuhaozhe?

[zhuhaozhe]

Hi, @angelayi.
We observed some regressions after this PR #120296.
I have compared the generated code for lennard_jones but I did not find any difference.

Then I collected the chrome trace and compared them.
The chrome_trace indicate this model is pretty small (20us level for 1 iter) and most event under Torch-compiled Region are not captured.

However, there is ~4us regression after this PR.
Before:

After:

I tried to locate the ~4us overhead and the last point I have found (with 4us regression) is

pytorch/torch/csrc/dynamo/eval_frame.c

Lines 182 to 198 in 1f4d4d3

	inline static PyObject* eval_frame_default(
	PyThreadState* tstate,
	THP_EVAL_API_FRAME_OBJECT* frame,
	int throw_flag) {
	#if PY_VERSION_HEX >= 0x03090000
	if (tstate == NULL) {
	tstate = PyThreadState_GET();
	}
	if (previous_eval_frame) {
	return previous_eval_frame(tstate, frame, throw_flag);
	}
	else {
	return _PyEval_EvalFrameDefault(tstate, frame, throw_flag);
	}
	#else
	return _PyEval_EvalFrameDefault(frame, throw_flag);
	#endif

.

I cannot see how your PR will impact here.
Do you have any insights on it? ~4 us overhead per iter while running an inductor optimized model.

[angelayi]

Hi @zhuhaozhe, sorry for the regression! I'm not exactly sure what could've caused the regression, but one thing I could see in my PR that might be impactful is that I inserted an additional check on this line, which gets run on every node in the torch compiled graph. Maybe you could try commenting out this check and see if that restores the 4us?

Or maybe detecting the functional tensor mode here is costly? Maybe you can try removing that line, and then replacing this line with with disable_above, FunctionalTensorMode(aot_config.pre_dispatch):?

cc @bdhirsh if you have any ideas 😅

[zhuhaozhe]

Hi, @angelayi.
I may located the 4us overhead here:

--- a/torch/_functorch/_aot_autograd/runtime_wrappers.py
+++ b/torch/_functorch/_aot_autograd/runtime_wrappers.py
@@ -76,7 +76,7 @@ def create_runtime_wrapper(
 
     def runtime_wrapper(*args):
         # Pass in effect tokens (See Note [Side-Effectful Tokens in AOTAutograd])
-        args = (*[torch.tensor([])] * num_tokens, *args)
+        # args = (*[torch.tensor([])] * num_tokens, *args)
 
         if trace_joint:
             args_ = list(args)

Coming from torch.tensor([])? We may able to avoid it by checking whether num_tokens ==0.

After comment out: (abs latency=0.027103ms)

dev,name,batch_size,speedup,abs_latency,compilation_latency,compression_ratio,eager_peak_mem,dynamo_peak_mem,calls_captured,unique_graphs,graph_breaks,unique_graph_breaks,autograd_captures,autograd_compiles
cpu,lennard_jones,1,1.941443,0.027103,2.771598,0.876977,40.813773,46.539162,9,1,0,0,0,0

Before comment out (abs latency=0.032131ms):

dev,name,batch_size,speedup,abs_latency,compilation_latency,compression_ratio,eager_peak_mem,dynamo_peak_mem,calls_captured,unique_graphs,graph_breaks,unique_graph_breaks,autograd_captures,autograd_compiles
cpu,lennard_jones,1,1.661396,0.032131,2.792909,0.876954,40.808038,46.533837,9,1,0,0,0,0

[zhuhaozhe]

I tried

--- a/torch/_functorch/_aot_autograd/runtime_wrappers.py
+++ b/torch/_functorch/_aot_autograd/runtime_wrappers.py
@@ -76,7 +76,8 @@ def create_runtime_wrapper(
 
     def runtime_wrapper(*args):
         # Pass in effect tokens (See Note [Side-Effectful Tokens in AOTAutograd])
-        args = (*[torch.tensor([])] * num_tokens, *args)
+        if num_tokens > 0:
+            args = (*[torch.tensor([])] * num_tokens, *args)

And the regression is gone.

[angelayi]

Oh interesting! thanks so much for the investigation! Feel free to send a PR :D

[zhuhaozhe]

Thanks @angelayi, this issue will be fixed by #122346.

[jansel]

Closing this issue as #122346 is landed. Please reopen if not fixed.