Skip to content

Conversation

@lsy323
Copy link
Collaborator

@lsy323 lsy323 commented Feb 29, 2024

Add FX passes to support dynamism on torch ops that take symbolic dim size as parameters (e.g. `aten.view(x, [sym_size, -1]). The FX passes groups the generation of symbolic dim size and the torch op into a new XLA op. Grouping them together is necessary to lowering these ops with dynamism, because the operations on symbolic size cannot be traced in LTC (context in #6393). Once the source of the symbolic size and the consuming torch ops are captured in a single XLA op, it becomes feasible to lowered to HLO/StableHLO with dynamism semantics.

The FX passes will run automatically if the exported program has symbolic shape input.

The following ops are fused into XLA ops:

  • sym_size.int + aten.expand => xla.dynamic_expand
  • sym_size.int + (mul) + aten.view => xla.dynamic_view

Some torch ops are generating aten.view with symbolic dim size during decomposition in upstream PyTorch, or existing lowering logic in torch_xla. FX passes are introduce to handle these ops as well.

  • aten.native_layer_norm
  • aten.group_norm
  • aten.select
  • aten.unsqueeze

Other changes:

  • Adds dynamism support to rsub
  • Adds dynamism support to mean and var
  • Move torch_xla/test/stablehlo/utils.py to torch_xla/utils/stablehlo_test_utils.py, to fix path not found issue when running tests with pytest.
  • Added unbounded dynamism unit tests for more ops.

Test:

  • Added test script for each of the FX passes, and unbounded dynamism unit tests for each supported op. Correctness of the lowering is done by comparing torch output and tf.saved_model inference result.
@lsy323 lsy323 force-pushed the lsiyuan/export-dynamism-fx-pass branch 4 times, most recently from b7924af to de636ac Compare March 7, 2024 21:52
@lsy323 lsy323 changed the title [WIP] Add some fx passes to support exporting unbounded dynamism Add some fx passes to support exporting unbounded dynamism Mar 8, 2024
@lsy323 lsy323 force-pushed the lsiyuan/export-dynamism-fx-pass branch 2 times, most recently from 6c237e6 to 4b41659 Compare March 20, 2024 00:41
@lsy323 lsy323 changed the title Add some fx passes to support exporting unbounded dynamism Add fx passes to support exporting unbounded dynamism Mar 20, 2024
Siyuan Liu added 8 commits March 20, 2024 02:25
format fix comment for skipped tests cover mul (cherry picked from commit f55abc88ae361e89da675a1aa1e4a19e7a5c762a) cover mul (cherry picked from commit 30abe2be43defc25db8954c525d34f7f3de35292) add missing tests to ci scripts yapf fix scalar type (cherry picked from commit 8526b2091ffafccf6972ecba3c111d1b0869621e) disable addmm test disable mark pattern api in gh ci, due to tf dep enable conv dynamism support addmm enable softmax dynamism update comment for slice add slice support, need converter change update test script take dynamic shape in save model export api verify lowering by adding tfl inference in tests remove debug pritn add assertion of sliced dim in select lowering remove log in conv, remove assertion in select re-enable test add select fx pass add no op slice removal pass add fx passes add tests' support layernorm add vit export scripot fix ep callable enable gelu test add export script support dynamic view with sym dim on dims other than BS add tests for gemma export support unsqueeze support softmax reduction on dynamic dim support unbounded index (unfinished) support dynamic expand add groupnorm add conv1d support, add dynamism (partially) to view add wav2vec2 export script add cumsum test, ne test remove existing tests change from crlf to lf add checks on view move stablehlo test util script remove debugging print add more assertions to fx passes remove test print add docstr to dynamic op make export script more concise
@lsy323 lsy323 force-pushed the lsiyuan/export-dynamism-fx-pass branch from 4b41659 to ce436f8 Compare March 20, 2024 02:29
@lsy323 lsy323 requested a review from qihqi March 20, 2024 16:30
@lsy323 lsy323 added the dynamism Dynamic Shape Features label Mar 20, 2024
@qihqi qihqi merged commit 7e0d3a5 into master Mar 20, 2024
@lsy323 lsy323 deleted the lsiyuan/export-dynamism-fx-pass branch March 20, 2024 16:51
@lsy323 lsy323 changed the title Add fx passes to support exporting unbounded dynamism Add fx passes to support unbounded dynamism in torch op arg Aug 30, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

dynamism Dynamic Shape Features

2 participants