[Fori_loop|While_loop] Enable while_loop/fori_loop, Add linear/MNIST test case #6867

ManfeiBai · 2024-04-01T21:26:19Z

add linear model test case with fori_loop or while_loop
modify logic of buildforiloop to chaneg cond/body's xla_computation to meet requirement of xla::While: cond's input, body's input/output and init should be the same shape;
modify input arguments shape/order to meet xla::While's requirement mentioned above

next plan:

init_python_bindings.cpp: get body xlacomputation arguments' number first then decide items in additional_inputs_list, maybe implement in python level

qihqi · 2024-04-18T20:30:31Z

torch_xla/experimental/fori_loop.py

-    res_list.insert(0, lower)
-    res_list.insert(0, torch.sub(upper, one_value_i))
-    return res_list
+  if (hasattr(body_fun, 'weight') or hasattr(body_fun, 'bias')):


what is the plan to remove this condition?

actually why do we special casing the weight and bias? Is it only for the linear layer?

yes, its only for the linear layer,

special casing the weight and bias due to different body_fn return: weight/bias was not mentioned in inputs, but need to be returned or added in xlacomputation return arguments

plans to remove this condition is:

A: check additional_inputs(weight/bias) like PyTorch and add them into xlacomputation arguments in CPP level

B: only check whether additional_inputs(weight/bias) exist for code of weight/bias in the next step

JackCaoG · 2024-04-19T00:08:55Z

torch_xla/csrc/init_python_bindings.cpp

+        xla::XlaOp x = xla::Parameter(local_builder, parameter_idx, shape,
+                                      "UnusedArgumentsPlaceholder");


I don't quite understand this part, x is a local variable that will be released after the for loop, is the intention of calling xla::Parameter to init a parameter at position parameter_idx with given shape? If so can you add a comment to make it more clear?

yes, x was added here to add expected shape/type arguments in body/cond xlacomputation's arguments due to unused input arguments missed with built via LTC,

to meet XLA::While requirement: parameter of condition and body, the result of the body, and init must all have the same shape

JackCaoG · 2024-04-19T00:09:16Z

torch_xla/csrc/init_python_bindings.cpp

-        xla::Shape shape =
-            xla::ShapeUtil::MakeShape(xla::PrimitiveType::S32, {1});
+      int64_t parameter_idx =
+          2;  // parameter_idx start from 2 after used upper and lower


can you move the comment above the code? Thanks!

JackCaoG · 2024-04-19T00:09:41Z

torch_xla/csrc/init_python_bindings.cpp

+      }
+    }
+
+    // hard-code modify body xlacomputation input arguments with unusedarguments


nit, space between unused and arguments

JackCaoG · 2024-04-19T00:12:31Z

torch_xla/experimental/fori_loop.py

+      weight = body_fun.weight  # not be used actually, initialized as placeholder xlacomputation requirement
+      bias = body_fun.bias  # not be used actually, initialized as placeholder xlacomputation requirement


nit, can we move the comment above the code?

JackCaoG · 2024-04-19T00:13:00Z

torch_xla/experimental/fori_loop.py

+      return upper.clone(), new_lower.clone(), one_value.clone(), torch.add(
+          one_value, x), input_value.clone(), bias.clone(), weight.clone(
+          ), output_value.clone()


why do we need to clone?

miss .clone() here would cause PyTorch ERROR: torch.while_loop's body_fn might be aliasing the input!, so add .clone() to avoid return input arguments directly

JackCaoG · 2024-04-19T00:18:51Z

test/test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py

+      return upper.clone(), new_lower.clone(), one_value.clone(), torch.add(
+          one_value, x), input_value.clone(), bias.clone(), weight.clone(
+          ), output_value.clone()


remind me why do we need to return weight and bias in this function?

we need to make sure body_fn's xlacomputation's input and output are the same, because input would include weight automatically, so here we return weight and bias from python level to ensure weight and bias are included in ouput too. Add bias to avoid output_value is used as bias in calculation, because bias has the same shape and value as output_value

but we also has plan to lower add weight and bias in xlacomputation arguments to CPP level, let me test locally too, if pass, we could avoid return weight and bias from python level

This is too confusing, we need to think of a better UX. body_fn also should take linear_0 as an input instead of calling it from parent scope.

I think instead of manually return the each parameter in the module, we should just return module's named_paramter . User also shouldn't need to manually order the return parameter(this will be super confusing to the user), we should do it in our layer.

miladm · 2024-05-06T22:17:52Z

do we have an ETA for this PR?

JackCaoG · 2024-05-07T12:57:45Z

I felt like the way we determine the HLO input parameter order is not ideal, it only works on simple examples. In this pr Manfei is trying to support Linear and she need to do some hacks to make the order correct. I am trying to figure out if there is a more general way to determine the parameter order, hopefully I can find some time today and tmr to draft a poc out..

JackCaoG · 2024-05-07T15:00:28Z

test/test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py

+      return upper.clone(), new_lower.clone(), one_value.clone(), torch.add(
+          one_value, x), input_value.clone(), bias.clone(), weight.clone(


why are we returning a torch.add(one_value, x) here?

It was used here to confirm calculation run expected times as a timer

I think our test case is too complicated, we should aim to support what pytorch support, similar to https://github.com/pytorch/pytorch/blob/8573d9551a7694b9313310412867ac3b6b751f26/test/functorch/test_control_flow.py#L137-L150.

ManfeiBai · 2024-05-07T18:32:28Z

I felt like the way we determine the HLO input parameter order is not ideal, it only works on simple examples. In this pr Manfei is trying to support Linear and she need to do some hacks to make the order correct. I am trying to figure out if there is a more general way to determine the parameter order, hopefully I can find some time today and tmr to draft a poc out..

Thanks, @JackCaoG, for comparison, this WIP PR is also trying post-order like XLAGraphExecutor::Compile according to your idea and suggestion: #7031, then add python interface for while_loop to generate xlacomputation for cond/body

JackCaoG · 2024-05-07T23:27:00Z

I have some high level ideas

as long as there is no inplace update to any parameters(which will not be true I guess since we need to decrement the iterator), get_hlo(input1, input2, input3, output) should have parameter in the same order. This is because in order to compute the new input1, we need input1 as result. The compilation of this approach is if somehow input1 is updated, for example

fn(input1, input2, input3, output):
  input1 += input2
  ...

get_hlo(input1, input2, input3, output) will likely result in the parameter order of input2, input1, input3, output. The key is always put all inputs in front of the output.

The same rule should apply for the named_parameters from the module , if they are being passed as additional_input from the while loop, we just need to stich them to inputs list(not from user code level, in our while loop implementation level). Hopefully this will just work..

ManfeiBai · 2024-05-08T19:09:25Z

get_hlo

Thanks, @JackCaoG,

for this idea, do we mean get_hlo is code and code?

I have some high level ideas

as long as there is no inplace update to any parameters(which will not be true I guess since we need to decrement the iterator), get_hlo(input1, input2, input3, output) should have parameter in the same order. This is because in order to compute the new input1, we need input1 as result. The compilation of this approach is if somehow input1 is updated, for example
fn(input1, input2, input3, output):
  input1 += input2
  ...
get_hlo(input1, input2, input3, output) will likely result in the parameter order of input2, input1, input3, output. The key is always put all inputs in front of the output.

The same rule should apply for the named_parameters from the module , if they are being passed as additional_input from the while loop, we just need to stich them to inputs list(not from user code level, in our while loop implementation level). Hopefully this will just work..

Thanks, @JackCaoG

for order change of input1 and input2, tried with test of add and linear: test code and test log based on current PR's branch, we saw the wrong result might due to the order change reason too when we change order of lower and upper

checked HLO of passed linear test: HLO, and failed linear test after switch lower and upper: HLO, looks like cond's xlacomputation didn't change after we switch order or inputs, so we do need to fix order of input and output like your suggestion, both in fori_loop's and while_loop's implementation

for order of input/output, we hasn't saw this error happen when we switch order of body's input_value and output_value yet, do we have more situation that missed here?

ManfeiBai · 2024-05-31T21:54:52Z

thanks for review and comments, since failed to push due to permission, would continue track in #7157

ManfeiBai force-pushed the fori_loop_simple_case_test branch 3 times, most recently from 4e0dcaa to 50193f3 Compare April 9, 2024 20:17

ManfeiBai force-pushed the fori_loop_simple_case_test branch 8 times, most recently from 92d17fd to a79a609 Compare April 16, 2024 17:34

ManfeiBai added 19 commits April 17, 2024 21:52

update

c4fe122

update

80d2003

update

71718e1

update

e8e18f7

update

9741e8d

update

26e30a3

update

c6f8259

update

5b9378c

update

6408879

update

7c408a9

update

9c68528

update

141c704

update

d5e5537

update

cfbe475

update

b11f015

update

9cd8ca0

update

89a2b34

update

9e420ad

update

a536a3e

ManfeiBai added 3 commits April 17, 2024 22:47

format

da35561

format

431ab66

format

6244d21

ManfeiBai requested review from GleasonK, miladm, qihqi, yeounoh and JackCaoG April 17, 2024 22:53

ManfeiBai added 2 commits April 17, 2024 23:06

format

04ca72d

format

33fa1fb

ManfeiBai marked this pull request as ready for review April 18, 2024 00:31

qihqi reviewed Apr 18, 2024

View reviewed changes

JackCaoG reviewed Apr 19, 2024

View reviewed changes

JackCaoG reviewed May 7, 2024

View reviewed changes

ManfeiBai changed the title ~~[Fori_loop|While_loop] Modify XlaComputation and add linear model test case~~ [Fori_loop|While_loop] Enable while_loop/fori_loop, Add linear/MNIST test case May 31, 2024

[rebase] rebase fori_loop_simple_case_test (#7165)

332bd40

ManfeiBai requested review from will-cromar and lsy323 as code owners May 31, 2024 21:48

ManfeiBai closed this May 31, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Fori_loop|While_loop] Enable while_loop/fori_loop, Add linear/MNIST test case #6867

[Fori_loop|While_loop] Enable while_loop/fori_loop, Add linear/MNIST test case #6867

ManfeiBai commented Apr 1, 2024 •

edited

Loading

qihqi Apr 18, 2024

JackCaoG Apr 19, 2024

ManfeiBai Apr 19, 2024 •

edited

Loading

JackCaoG Apr 19, 2024

ManfeiBai Apr 19, 2024 •

edited

Loading

JackCaoG Apr 19, 2024

JackCaoG Apr 19, 2024

JackCaoG Apr 19, 2024

JackCaoG Apr 19, 2024

ManfeiBai Apr 19, 2024

JackCaoG Apr 19, 2024

ManfeiBai Apr 19, 2024 •

edited

Loading

JackCaoG Apr 19, 2024

miladm commented May 6, 2024

JackCaoG commented May 7, 2024

JackCaoG May 7, 2024

ManfeiBai May 7, 2024

JackCaoG May 7, 2024

ManfeiBai commented May 7, 2024

JackCaoG commented May 7, 2024

ManfeiBai commented May 8, 2024 •

edited

Loading

ManfeiBai commented May 31, 2024 •

edited

Loading

		xla::XlaOp x = xla::Parameter(local_builder, parameter_idx, shape,
		"UnusedArgumentsPlaceholder");

		weight = body_fun.weight # not be used actually, initialized as placeholder xlacomputation requirement
		bias = body_fun.bias # not be used actually, initialized as placeholder xlacomputation requirement

		return upper.clone(), new_lower.clone(), one_value.clone(), torch.add(
		one_value, x), input_value.clone(), bias.clone(), weight.clone(

[Fori_loop|While_loop] Enable while_loop/fori_loop, Add linear/MNIST test case #6867

[Fori_loop|While_loop] Enable while_loop/fori_loop, Add linear/MNIST test case #6867

Conversation

ManfeiBai commented Apr 1, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ManfeiBai Apr 19, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ManfeiBai Apr 19, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ManfeiBai Apr 19, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

miladm commented May 6, 2024

JackCaoG commented May 7, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ManfeiBai commented May 7, 2024

JackCaoG commented May 7, 2024

ManfeiBai commented May 8, 2024 • edited Loading

ManfeiBai commented May 31, 2024 • edited Loading

ManfeiBai commented Apr 1, 2024 •

edited

Loading

ManfeiBai Apr 19, 2024 •

edited

Loading

ManfeiBai Apr 19, 2024 •

edited

Loading

ManfeiBai Apr 19, 2024 •

edited

Loading

ManfeiBai commented May 8, 2024 •

edited

Loading

ManfeiBai commented May 31, 2024 •

edited

Loading