Skip execution for extract_compiled_graph

[JackCaoG]

Currently we use xla_sync_multi to force a execution which will warm up the cache for future dynamo runs, this is not ideal because

1. execution can be time consuming
2. xla_sync_multi is considered as a mark_step like sync and will trigger buffer aliasing. However we actually saved some XLA data when we cache the dynamo graph. Those buffer might be aliased during the xla_sync_multi and cause random crash when we actually execute the dynamo graph

confirmed that this change does not regress speed on TPU when running torchbench.

FYI @wconstab @will-cromar

[vanbasten23]

Where can I read and learn the buffer aliasing?

[JackCaoG]

Ah, one thing I realized is we actually need to materializing some of the tensor, at least those one we save for future execution. Let me tweak it a bit.

[JackCaoG]

Where can I read and learn the buffer aliasing?

If you search Aliasing in XLA you should see some very very brief doc.. I mostly just read codes.

[JackCaoG]

@shunting314 @alanwaketan This pr is ready for review.

[shunting314]

xla_sync_multi is considered as a mark_step like sync and will trigger buffer aliasing. However we actually saved some XLA data when we cache the dynamo graph. Those buffer might be aliased during the xla_sync_multi and cause random crash when we actually execute the dynamo graph

Is it easy to have a unit test to show the buffer aliasing related issue?

[JackCaoG]

yea, actually PJRT=CPU + the current resnet18 test crashes(currently ci is using XRT, we are migrating test to PJRT runtime then I find this issue). Do you want to add a smaller unit test?

[shunting314]

TBH, I don't quite understand what the 'buffer aliasing' issue fixed by this PR is. Or more specifically, how does it cause crashes. A simple unit test should help.

But it's not blocking, I think we can merge this PR first as well.

[alanwaketan]

LGTM.

[JackCaoG]

@shunting314 I should be more specified. The buffer aliasing is only enabled during a mark_step or equivant(sync_multi ) The code is in

xla/torch_xla/csrc/xla_graph_executor.cpp

Lines 1104 to 1105 in 574cfda

	if (enable_aliasing && coll.config.sync_ltc_data &&
	coll.config.force_ltc_data && !is_sharded) {

In this pr I set the config.force_ltc_data to False during the warm_up_cache hence it will skip the aliasing. Regarding what Aliasing does, we have a simple test in

xla/test/test_input_output_aliases.py

Lines 26 to 34 in 574cfda

	self.assertEqual(met.metric_data("InputOutputAliasCount")[1], 2.0)
	# check in place op aliasing.
	t3 = t1 + t2
	t1 *= 2.0
	t2 += 2.0
	xm.mark_step()
	self.assertEqual(met.metric_data("InputOutputAliasCount")[1], 4.0)

You see that we perform two inplace operation on t1 and t2. We will have a hlo similar to

hlo(input1, input2) -> (output1, output2, output3)

input in this case is t1 and t2, output is the t3 and new t1 and t2. What aliasing does is instead of allocating new buffer for the output1 and output2, it reuses the input buffer since pytorch/xla can tell compiler this is a inplace operation and we don't need t1 and t2's original buffer after this computation.

The bug we had prior this pr is we actually saved some input for each dynamo graph. However when we do xla_sync_multi to warm up the cache, we incorrectly trigger the aliasing. This means the buffer we saved was donated to the output of the xla_sync_multi which will crash when we tried to use saved input for the real dynamo run.