[Fori_loop|While_loop] Enable fori_loop with add/sub test case #6603

ManfeiBai · 2024-02-23T21:40:38Z

For fori_loop implementation with while_loop, this PR is for lowering body/cond to replace formal placeholder

This is the step two PR, and father PR(#6532), child PR(#6529), source PR(#6563)

some issue fixed:

~~body fn is -, not a torch func, will test later~~ tried torch.sub(a, b), passed too locally
current code has changed many logic of lowering, let's move these logics to a new function without affecting the existing functions
~~input are limited to list/tuple~~ this match torch._higher_order_ops.while_loop required
~~input was trans from list to not list after torch.compile, TODO, add the same logic like torch.compile to use inputs, not like currently create a duplicated tensor in the fori_loop.py file~~

ManfeiBai · 2024-02-26T19:27:23Z

torch_xla/csrc/lowering_context.cpp

+  if (!root_tuple_.empty() & (root_tuple_.size() > 1)) {
+    xla::XlaOp root = xla::Tuple(builder(), root_tuple_);
+    xla = builder()->Build(root);
+  } else if (!root_tuple_.empty() & (root_tuple_.size() == 1)) {


Explain: we need to skip tuple for cond/body computation creation to match the xla::While format check for cond, error log

ManfeiBai · 2024-02-26T19:36:37Z

Hi, @JackCaoG, since this PR would add new function to PyLoweringContext, do we want to request review from aws too?

JackCaoG · 2024-02-26T23:06:58Z

@amithrm FYI

ManfeiBai · 2024-02-26T23:55:25Z

kokoro failure should be fixed on master branch, let's skip it now

torch_xla/csrc/init_python_bindings.cpp

torch_xla/experimental/fori_loop.py

torch_xla/csrc/lowering_context.h

torch_xla/experimental/fori_loop.py

yeounoh

Left some suggestions

yeounoh · 2024-03-04T21:21:28Z

torch_xla/csrc/init_python_bindings.cpp

@@ -1027,7 +1033,9 @@ void BuildLoweringContextSubmodule(py::module* m) {
      .def("hlo_json", &PyLoweringContext::GetHloJsonText)
      .def("parameter_id_tensor_mapping",
           &PyLoweringContext::GetParameterIdTensorMapping)
-      .def("tensor_parameter_id", &PyLoweringContext::GetTensorParameterId);
+      .def("tensor_parameter_id", &PyLoweringContext::GetTensorParameterId)
+      .def("set_name_string", &PyLoweringContext::SetNameString)


Good, thanks!

yeounoh · 2024-03-04T21:27:38Z

torch_xla/csrc/lowering_context.cpp

    xla::XlaOp root = xla::Tuple(builder(), root_tuple_);
    xla = builder()->Build(root);
+  } else if (!root_tuple_.empty() & (root_tuple_.size() == 1)) {


I think we should condition on get_name_string(). Add this check at the top and build for while loop `if get_name_string() == "condctx" or get_name_string() == "bodyctx"; otherwise, you can keep the original build logic.

Have your logic for while loop build in a separate private method, and call it if ``if get_name_string() == "condctx" or get_name_string() == "bodyctx"` is true.

So you can keep BuildXla() simple.

Thanks, make sense and updated in the newest commit, due to the logic for while loop is one simple line code, we run it directly without warping it in a separate private method

yeounoh · 2024-03-04T21:28:36Z

torch_xla/csrc/lowering_context.h

@@ -68,7 +72,8 @@ class LoweringContext : public torch::lazy::LoweringContext {
  xla::XlaOp GetOutputOp(const torch::lazy::Output& output);

  // Build the XLA computation capturing all the operations created with the
-  // embedded XLA builder (returned by the builder() API).
+  // embedded XLA builder (returned by the builder() API) with check whether


Think we can keep the original comment. Add the while specific comment on the new private method, as described above.

Thanks, updated

miladm · 2024-03-04T23:18:28Z

test/test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py

    def cond_fn(x):
-      return x.sum() <= 10
+      ten = torch.ones(1, dtype=torch.int32, device=device)


cosmetic nit: ten can be interpreted as number 10 - suggesting a better name plz

Thanks, updated

ManfeiBai changed the title ~~Update test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py~~ [Do Not Merge] Update test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py Feb 23, 2024

ManfeiBai mentioned this pull request Feb 23, 2024

[Fori_loop] Support torch.while_loop with multi situations #6529

Closed

ManfeiBai added DO_NOT_MERGE_YET For PRs which cannot be merged, despite tests passing fori_loop labels Feb 23, 2024

ManfeiBai changed the title ~~[Do Not Merge] Update test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py~~ [Do Not Merge][Fori_loop|While_loop] Update test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py Feb 24, 2024

ManfeiBai force-pushed the ManfeiBai-patch-73 branch from 7a8f9f5 to c0d3359 Compare February 26, 2024 18:57

ManfeiBai commented Feb 26, 2024

View reviewed changes

ManfeiBai marked this pull request as ready for review February 26, 2024 19:34

ManfeiBai requested review from miladm, qihqi, yeounoh and JackCaoG February 26, 2024 19:35

ManfeiBai changed the title ~~[Do Not Merge][Fori_loop|While_loop] Update test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py~~ [Fori_loop|While_loop] Update test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py Feb 26, 2024

ManfeiBai removed the DO_NOT_MERGE_YET For PRs which cannot be merged, despite tests passing label Feb 27, 2024

yeounoh reviewed Feb 28, 2024

View reviewed changes

torch_xla/csrc/init_python_bindings.cpp Outdated Show resolved Hide resolved

yeounoh reviewed Feb 28, 2024

View reviewed changes

torch_xla/experimental/fori_loop.py Outdated Show resolved Hide resolved

ManfeiBai requested a review from yeounoh February 29, 2024 18:07

qihqi approved these changes Feb 29, 2024

View reviewed changes

yeounoh reviewed Feb 29, 2024

View reviewed changes

torch_xla/csrc/lowering_context.h Outdated Show resolved Hide resolved

yeounoh reviewed Feb 29, 2024

View reviewed changes

torch_xla/experimental/fori_loop.py Outdated Show resolved Hide resolved

yeounoh suggested changes Feb 29, 2024

View reviewed changes

ManfeiBai requested a review from yeounoh February 29, 2024 22:27

ManfeiBai force-pushed the ManfeiBai-patch-73 branch from c35ac53 to c7f09d5 Compare March 4, 2024 17:47

yeounoh reviewed Mar 4, 2024

View reviewed changes

miladm reviewed Mar 4, 2024

View reviewed changes

ManfeiBai and others added 27 commits March 8, 2024 18:44

remove mark_step

abe3067

Update lowering_context.cpp

56f5cb0

Update lowering_context.h

b5714f7

Update init_python_bindings.cpp

e7fa408

Update lowering_context.cpp

66402ef

Update lowering_context.h

6b1a570

Update lowering_context.h

40d39fa

Update fori_loop.py

93638c2

Update fori_loop.py

28a3dd7

refactor logic of check while_loop

e014c1d

update comment

4e6f076

format

dafa8fb

format

1b952cb

format

8ec55c3

add test case

aeb0a58

add test case

02bf5ba

add test case

d51e069

Update test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py

5f3af2b

test case for add and sub

cb26670

format

f37149b

format

a48a3cf

wrap test situation check

f30eaab

wrap test situation check

1633d8e

wrap test situation check

8a3b9fe

format

f6fdc78

Update test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py

5c2a2a0

nit

173ff44

ManfeiBai force-pushed the ManfeiBai-patch-73 branch from 96ce688 to 173ff44 Compare March 8, 2024 18:44

ManfeiBai merged commit 6170df5 into master Mar 8, 2024
18 checks passed

ManfeiBai added the q1_2024_manfei label Mar 22, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Fori_loop|While_loop] Enable fori_loop with add/sub test case #6603

[Fori_loop|While_loop] Enable fori_loop with add/sub test case #6603

ManfeiBai commented Feb 23, 2024 •

edited

Loading

ManfeiBai Feb 26, 2024 •

edited

Loading

ManfeiBai commented Feb 26, 2024 •

edited

Loading

JackCaoG commented Feb 26, 2024

ManfeiBai commented Feb 26, 2024

yeounoh left a comment

yeounoh Mar 4, 2024

yeounoh Mar 4, 2024

ManfeiBai Mar 4, 2024

yeounoh Mar 4, 2024

ManfeiBai Mar 4, 2024

miladm Mar 4, 2024

ManfeiBai Mar 4, 2024

[Fori_loop|While_loop] Enable fori_loop with add/sub test case #6603

[Fori_loop|While_loop] Enable fori_loop with add/sub test case #6603

Conversation

ManfeiBai commented Feb 23, 2024 • edited Loading

ManfeiBai Feb 26, 2024 • edited Loading

Choose a reason for hiding this comment

ManfeiBai commented Feb 26, 2024 • edited Loading

JackCaoG commented Feb 26, 2024

ManfeiBai commented Feb 26, 2024

yeounoh left a comment

Choose a reason for hiding this comment

yeounoh Mar 4, 2024

Choose a reason for hiding this comment

yeounoh Mar 4, 2024

Choose a reason for hiding this comment

ManfeiBai Mar 4, 2024

Choose a reason for hiding this comment

yeounoh Mar 4, 2024

Choose a reason for hiding this comment

ManfeiBai Mar 4, 2024

Choose a reason for hiding this comment

miladm Mar 4, 2024

Choose a reason for hiding this comment

ManfeiBai Mar 4, 2024

Choose a reason for hiding this comment

ManfeiBai commented Feb 23, 2024 •

edited

Loading

ManfeiBai Feb 26, 2024 •

edited

Loading

ManfeiBai commented Feb 26, 2024 •

edited

Loading