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

test/test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py

@@ -20,10 +20,14 @@ def _fake_while_loop(cond_fn, body_fn, operands):
 
 
 def _fake_fori_loop(lower, upper, body_fun, *init_val):
-  (plus_value, init_val) = init_val
-  for i in range((upper - lower)[0]):
-    plus_value, init_val = body_fun(plus_value, init_val)
-  return init_val
+  if len(init_val) > 1:
+    (a, b) = init_val
+    for i in range((upper - lower)[0]):
+      a = body_fun(a, b)
+  else:
+    for i in range((upper - lower)[0]):
+      a = body_fun(*init_val)
+  return a
 
 
 class WhileLoopTest(unittest.TestCase):
@@ -82,25 +86,139 @@ def body_fn(init, limit_value):
     expected = _fake_while_loop(cond_fn, body_fn, (init, limit_value))
     self.assertEqual(expected, res)
 
+  def test_while_loop_tpu_simple_linear(self):
+
+    xm.mark_step()
+    device = xm.xla_device()
+    torch.set_grad_enabled(False)
+
+    linear_0 = torch.nn.Linear(10, 20).to(xm.xla_device())
+
+    def cond_fn(upper, lower, one_value, x, input_value, output_value):
+      return lower[0] < upper[0]
+
+    def body_fn(upper, lower, one_value, x, input_value, output_value):
+      new_lower = torch.add(one_value, lower)
+      output_value = linear_0(input_value)
+      weight = linear_0.weight  # not be used actually, initialized as placeholder xlacomputation requirement
+      bias = linear_0.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(
+          ), output_value.clone()
+
+    upper = torch.tensor([1], dtype=torch.int32, device=device)
+    lower = torch.tensor([0], dtype=torch.int32, device=device)
+    one_value = torch.tensor([1], dtype=torch.int32, device=device)
+    init_val = torch.tensor([1], dtype=torch.int32, device=device)
+    l_in_0 = torch.rand(10, device=xm.xla_device())
+    output_value = torch.zeros([20], dtype=torch.float32, device=device)
+
+    upper__, lower__, one_value__, torch_add_res__, input_value__, bias__, weight__, output_value_real__, = while_loop(
+        cond_fn, body_fn,
+        (upper, lower, one_value, init_val, l_in_0, output_value))
+
+    expected = _fake_fori_loop(lower, upper, linear_0, l_in_0)
+
+    return self.assertTrue(torch.all(torch.eq(expected, output_value_real__)))
+
+  def test_while_loop_tpu_simple_linear_class(self):
+
+    xm.mark_step()
+    device = xm.xla_device()
+    torch.set_grad_enabled(False)
+
+    class SimpleWithLinear(torch.nn.Module):
+
+      def __init__(self):
+        super().__init__()
+        self.linear = torch.nn.Linear(10, 20).to(xm.xla_device())
+
+      def forward(self, upper, lower, one_value, x, input_value, output_value):
+
+        def cond_fn(upper, lower, one_value, x, input_value, output_value):
+          return lower[0] < upper[0]
+
+        def body_fn(upper, lower, one_value, x, input_value, output_value):
+          new_lower = torch.add(one_value, lower)
+          output_value_real = self.linear(input_value)
+          weight = self.linear.weight  # not be used actually, initialized as placeholder xlacomputation requirement
+          bias = self.linear.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(
+              ), output_value_real, weight.clone(), bias.clone()
+
+        return while_loop(
+            cond_fn, body_fn,
+            (upper, lower, one_value, x, input_value, output_value))
+
+    simple_with_linear = SimpleWithLinear()
+    upper = torch.tensor([52], dtype=torch.int32, device=device)
+    lower = torch.tensor([0], dtype=torch.int32, device=device)
+    one_value = torch.tensor([1], dtype=torch.int32, device=device)
+    init_val = torch.tensor([1], dtype=torch.int32, device=device)
+    l_in_0 = torch.rand(10, device=xm.xla_device())
+    output_value = torch.zeros([20], dtype=torch.float32, device=device)
+
+    weight_0 = simple_with_linear.linear.weight
+    bias_0 = simple_with_linear.linear.bias
+
+    aaa = {
+        "simple_with_linear":
+            (simple_with_linear, (upper, lower, one_value, init_val, l_in_0,
+                                  output_value))
+    }
+
+    upper__, lower__, one_value__, torch_add_res__, input_value__, output_value_real__, weight__, bias__ = simple_with_linear(
+        upper, lower, one_value, init_val, l_in_0, output_value)
+
+    # create same weight/bias liear model for compare
+    linear_0 = torch.nn.Linear(10, 20).to(xm.xla_device())
+    linear_0.weight.data = weight__
+    linear_0.bias.data = bias__
+    expected = _fake_fori_loop(lower, upper, linear_0, l_in_0)
+
+    self.assertTrue(torch.all(torch.eq(expected, output_value_real__)))
+    return aaa
+
   def test_fori_loop_tpu_addition(self):
 
     xm.mark_step()
     device = xm.xla_device()
 
     lower = torch.tensor([2], dtype=torch.int32, device=device)
     upper = torch.tensor([52], dtype=torch.int32, device=device)
-    plus_value = torch.tensor([1], dtype=torch.int32, device=device)
+    one_value = torch.tensor([1], dtype=torch.int32, device=device)
+    init_val = torch.tensor([1], dtype=torch.int32, device=device)
+
+    def body_fun(a, b):
+      return torch.add(a, b)
+
+    upper_, new_lower_, one_value_, add_res_x_, res_ = fori_loop(
+        upper, lower, body_fun, one_value, init_val)
+    expected = _fake_fori_loop(lower, upper, body_fun, init_val, one_value)
+    self.assertEqual(expected, res_)
+
+  def test_fori_loop_tpu_simple_linear(self):
+
+    xm.mark_step()
+    device = xm.xla_device()
+    torch.set_grad_enabled(False)
+
+    upper = torch.tensor([52], dtype=torch.int32, device=device)
+    lower = torch.tensor([0], dtype=torch.int32, device=device)
     init_val = torch.tensor([1], dtype=torch.int32, device=device)
+    l_in_0 = torch.randn(10, device=xm.xla_device())
+
+    linear_0 = torch.nn.Linear(10, 20).to(xm.xla_device())
+
+    upper_, lower_, one_value_, add_res_x_, l_in_i_plus_1_, weight_, bias_, l_out_ = fori_loop(
+        upper, lower, linear_0, init_val, l_in_0)
 
-    def body_fun(*argus):
-      plus_value, init_val = argus
-      return plus_value, torch.add(plus_value, init_val)
+    expected = _fake_fori_loop(lower, upper, linear_0, l_in_0)
 
-    _, _, _, actual = fori_loop(upper, lower, body_fun, plus_value, init_val)
-    expected = _fake_fori_loop(lower, upper, body_fun, plus_value, init_val)
-    self.assertEqual(expected, actual)
+    self.assertTrue(torch.all(torch.eq(expected, l_out_)))
 
 
 if __name__ == '__main__':
   test = unittest.main()
-  sys.exit(0 if test.result.wasSuccessful() else 1)
+  sys.exit(0 if test.result.wasSuccessful() else 1)

torch_xla/csrc/init_python_bindings.cpp

@@ -913,14 +913,31 @@ class PyLoweringContext {
   // Builds a HLO graph given a set of output tensors, and add unused parameters
   // needed in xlacomputation.
   void BuildForiLoop(std::vector<at::Tensor> tensors,
-                     std::vector<at::Tensor> input_arguments = {}) {
+                     std::vector<at::Tensor> additional_inputs_list = {}) {
+    // hard-code modify cond xlacomputation input arguments with unusedarguments
+    // for xla::while requriement
     if (GetNameString() == "condctx") {
       xla::XlaBuilder* local_builder = lowering_ctx.builder();
-      // hard-code parameter_idx to 2 to skip existing upper/lower arguments
-      int64_t parameter_idx = 2;
-      for (at::Tensor input_argument : input_arguments) {
-        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
+      for (auto& additional_input_tensor : additional_inputs_list) {
+        XLATensorPtr xtensor = bridge::GetXlaTensor(additional_input_tensor);
+        xla::Shape shape = xtensor->shape().get();
+        xla::XlaOp x = xla::Parameter(local_builder, parameter_idx, shape,
+                                      "UnusedArgumentsPlaceholder");
+        parameter_idx += 1;
+      }
+    }
+
+    // hard-code modify body xlacomputation input arguments with unusedarguments
+    // for xla::while requriement
+    if (GetNameString() == "bodyctx") {
+      xla::XlaBuilder* local_builder = lowering_ctx.builder();
+      // TODO(@manfei): treat hard code parameter_idx value
+      int64_t parameter_idx = 7;
+      for (auto& additional_input_tensor : additional_inputs_list) {
+        XLATensorPtr xtensor = bridge::GetXlaTensor(additional_input_tensor);
+        xla::Shape shape = xtensor->shape().get();
         xla::XlaOp x = xla::Parameter(local_builder, parameter_idx, shape,
                                       "UnusedArgumentsPlaceholder");
         parameter_idx += 1;

torch_xla/experimental/fori_loop.py

@@ -10,66 +10,95 @@
 from torch._ops import HigherOrderOperator
 import torch._higher_order_ops.while_loop
 from torch._higher_order_ops.while_loop import while_loop_op
+from torch._higher_order_ops.while_loop import while_loop as torch_while_loop
 
 
-def fori_loop(lower, upper, user_body_func, *init_val):
+# TODO(@manfei): treat *input_value
+def fori_loop(upper, lower, body_fun, init_val, input_value):
 
   device = xm.xla_device()
 
-  def cond_fn(upper, lower, *init_val):
-    return lower[0] < upper[0]
+  one_value = torch.tensor([1], dtype=torch.int32, device=device)
 
-  def body_fn(upper, lower, *init_val):
-    one_value_i = torch.ones(1, dtype=torch.int32, device=device)
-    res_list = list(user_body_func(*init_val))
-    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')):
+    output_value = torch.zeros([20], dtype=torch.float32, device=device)
+
+    def cond_fn(upper, lower, one_value, x, input_value, output_value):
+      return lower[0] < upper[0]
+
+    def body_fn(upper, lower, one_value, x, input_value, output_value):
+      new_lower = torch.add(one_value, lower)
+      output_value = body_fun(input_value)
+      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(
+          ), output_value.clone()
+
+    res = torch_while_loop(
+        cond_fn, body_fn,
+        (upper, lower, one_value, init_val, input_value, output_value))
+  else:
+    output_value = torch.tensor([1], dtype=torch.int32, device=device)
+
+    def cond_fn(upper, lower, one_value, x, input_value):
+      return lower[0] < upper[0]
+
+    def body_fn(upper, lower, one_value, x, input_value):
+      new_lower = torch.add(one_value, lower)
+      output_val = body_fun(one_value, input_value)
+      return upper.clone(), new_lower.clone(), one_value.clone(), torch.add(
+          one_value, x), output_val.clone()
+
+    res = torch_while_loop(cond_fn, body_fn,
+                           (upper, lower, one_value, init_val, input_value))
 
-  res = while_loop(cond_fn, body_fn, (lower, upper, *init_val))
   return res
 
 
 @while_loop_op.py_impl(DispatchKey.XLA)
-def while_loop(cond_fn, body_fn, *carried_inputs, additional_inputs=None):
+def while_loop(cond_fn, body_fn, carried_inputs, additional_inputs=None):
   # TODO(@manfei): PyTorch require carried_inputs to be list/tuple, PyTorch/XLA _xla_while_loop only accept *operands, *operands would tuple items again: (a, '')
   # cond_fn&body_fn: callable
   # carried_inputs: (Tuple of possibly nested dict/list/tuple of tensors)
   if additional_inputs is None:
     additional_inputs = tuple()
-  return _xla_while_loop(
-      cond_fn, body_fn, *carried_inputs, additional_inputs=additional_inputs)
+  return _xla_while_loop(cond_fn, body_fn, carried_inputs, additional_inputs)
 
 
-def _xla_while_loop(cond_fn, body_fn, *carried_inputs, additional_inputs):
-  # untuple carried_inputs from while_loop
-  carried_inputs = carried_inputs[0]
+def _xla_while_loop(cond_fn, body_fn, carried_inputs, additional_inputs=None):
   # fake carried_inputs to split formal code
   fake_carried_inputs = []
   for carried_input in carried_inputs:
     device = carried_input.device
     fake_carried_inputs.append(
         torch.randint(10, carried_input.size(),
                       dtype=carried_input.dtype).to(device))
-  fake_carried_inputs = tuple(fake_carried_inputs)
-
-  # trans fake_carried_inputs from list(tensor) to list(xla::op)
-  kwargs = {}
-  if type(fake_carried_inputs) is tuple:
-    shapes = xb.tensor_shape(fake_carried_inputs)
-  else:
-    shapes = xb.tensor_shape((fake_carried_inputs))
-  builder = xb.create_builder('test_while')
-  params = []
-  for shape in shapes:
-    p = xb.mkparam(builder, len(params), shape)
-    params.append(p)
+  for additional_input in additional_inputs:
+    device = additional_input.device
+    fake_carried_inputs.append(
+        torch.randint(
+            10, additional_input.size(),
+            dtype=additional_input.dtype).to(device))
 
-  # generate cond_fn xlacomputation
+  # TODO(@manfei): specify which element is for which argument like a,b,c
   cond_result = cond_fn(*fake_carried_inputs)
   cond_ctx = torch_xla._XLAC.lowering.LoweringContext()
   cond_ctx.set_name_string("condctx")
-  cond_ctx.buildforiloop([cond_result], list(fake_carried_inputs[2:]))
+
+  # TODO(@manfei): treat hard-code cond xlacomputation change: currently switch output_value and weight position if additional_inputs(weight/bias) exists
+  additional_inputs_list_cond = list(
+      fake_carried_inputs[2:]
+  )  # all missed arguments except upper/lower due to PyTorch/XLA trace from output tensor
+  if additional_inputs:
+    tmp_bias = additional_inputs_list_cond[
+        -3]  # not used, change order doesn't affect logic
+    del additional_inputs_list_cond[
+        -3]  # not used, change order doesn't affect logic
+    additional_inputs_list_cond.append(
+        tmp_bias)  # not used, change order doesn't affect logic
+
+  cond_ctx.buildforiloop([cond_result], additional_inputs_list_cond)
   cond_hlo = cond_ctx.hlo()
   cond_computation = xb.computation_from_module_proto("condcomputation",
                                                       cond_hlo)
@@ -78,11 +107,38 @@ def _xla_while_loop(cond_fn, body_fn, *carried_inputs, additional_inputs):
   body_result = body_fn(*fake_carried_inputs)
   body_ctx = torch_xla._XLAC.lowering.LoweringContext()
   body_ctx.set_name_string("bodyctx")
-  body_ctx.buildforiloop(list(body_result), [])
+
+  # TODO(@manfei): treat hard-code body xlacomputation change: currently add non-changed output_value argument if additional_inputs(weight/bias) exists
+  if additional_inputs:
+    additional_inputs_list_body = [fake_carried_inputs[-3]]
+  else:
+    additional_inputs_list_body = []
+
+  # TODO(@manfei): treat hard-code parameters: additional_inputs_list_body
+  body_ctx.buildforiloop(list(body_result), additional_inputs_list_body)
   body_hlo = body_ctx.hlo()
   body_computation = xb.computation_from_module_proto("bodycomputation",
                                                       body_hlo)
 
+  # trans fake_carried_inputs from list(tensor) to list(xla::op), which part could change init of xla::while
+  total_inputs = carried_inputs + additional_inputs
+  kwargs = {}
+  if type(total_inputs) is tuple:
+    shapes = xb.tensor_shape(total_inputs)
+  else:
+    shapes = xb.tensor_shape((total_inputs))
+  builder = xb.create_builder('test_while')
+  params = []
+  for shape in shapes:
+    p = xb.mkparam(builder, len(params), shape)
+    params.append(p)
+
+  # TODO(@manfei): treat hard-code input arguments, currently switch bias and output_value if additional_inputs(weight/bias) exists
+  if additional_inputs:
+    tmp_bias = params[-3]
+    del params[-3]
+    params.append(tmp_bias)
+
   # generate while xlacomputation
   input_tuple = xb.Op.tuple(tuple(params))
   w = xb.mkop(
@@ -94,6 +150,6 @@ def _xla_while_loop(cond_fn, body_fn, *carried_inputs, additional_inputs):
 
   # gain final result with generated while xlacomputation
   result = torch_xla._XLAC._xla_user_computation('xla::_op_test_while',
-                                                 (carried_inputs), computation)
+                                                 (total_inputs), computation)
 
-  return result
+  return result