Skip to content

[Fori_loop|While_loop] Fori loop wrapped from while_loop#6850

Merged
ManfeiBai merged 1 commit intomasterfrom
ManfeiBai-patch-79
Apr 11, 2024
Merged

[Fori_loop|While_loop] Fori loop wrapped from while_loop#6850
ManfeiBai merged 1 commit intomasterfrom
ManfeiBai-patch-79

Conversation

@ManfeiBai
Copy link
Copy Markdown
Collaborator

@ManfeiBai ManfeiBai commented Mar 29, 2024

  • create limited fori_loop(forward) with while_loop to support simple test case
  • insert one s32 unused arguments in cond_xlacomputation, and add CPP test for add unused argument
  • use fake tensor to generate xlacomputation for body/cond to prevent formal IR affect

@ManfeiBai ManfeiBai marked this pull request as ready for review March 29, 2024 19:49
Comment thread test/cpp/test_xla_sharding.cpp Outdated
Comment thread torch_xla/experimental/fori_loop.py Outdated
Comment thread torch_xla/experimental/fori_loop.py Outdated
@ManfeiBai ManfeiBai requested a review from qihqi April 1, 2024 19:00
@ManfeiBai ManfeiBai force-pushed the ManfeiBai-patch-79 branch from 6439a0a to 15fdd9a Compare April 2, 2024 00:03
Comment thread test/test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py Outdated
Comment thread torch_xla/csrc/init_python_bindings.cpp Outdated
Comment thread torch_xla/csrc/init_python_bindings.cpp Outdated
Comment thread torch_xla/experimental/fori_loop.py Outdated
Comment thread torch_xla/experimental/fori_loop.py Outdated
Comment thread torch_xla/experimental/fori_loop.py Outdated
def _xla_while_loop(cond_fn, body_fn, *carried_inputs, additional_inputs):
# untuple carried_inputs from while_loop
carried_inputs = carried_inputs[0]
# fake carried_inputs to split formal code
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JackCaoG could you help me understand this part? why we need this.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

more context: this part create fake tensor from input tensors, so that fake tensor would not include formal IR like input tensors, and then we use clean fake tensor to create body_xlacomputation/cond_xlacomputation

Copy link
Copy Markdown
Contributor

@yeounoh yeounoh left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good, left some comments for minor changes.

@ManfeiBai ManfeiBai requested a review from yeounoh April 9, 2024 01:28
Copy link
Copy Markdown
Contributor

@yeounoh yeounoh left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants