[test] Againwhile loop lowering with simplecalculation dispatch in torch

.github/workflows/tpu_ci.yml

@@ -46,3 +46,4 @@ jobs:
           python3 -u test/test_autocast.py
           python3 -u test/dynamo/test_dynamo.py
           python3 -u test/spmd/test_spmd_debugging.py
+          python3 -u test/test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py

test/run_tests.sh

@@ -188,6 +188,7 @@ function run_xla_op_tests1 {
 function run_xla_op_tests2 {
   run_downcast_bf16 "$CDIR/test_data_type.py"
   run_test "$CDIR/pjrt/test_dtypes.py"
+  run_test "$CDIR/test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py"
   run_test "$CDIR/test_autocast.py"  # TODO(yeounoh) this is expensive on GPU
 }
 

test/test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py

@@ -0,0 +1,55 @@
+import os
+import unittest
+from typing import Callable, Dict, List
+
+import torch
+import torch_xla
+import torch_xla.experimental.fori_loop
+from torch._higher_order_ops.while_loop import while_loop
+import torch_xla.core.xla_model as xm
+import torch_xla.core.xla_builder as xb
+
+
+def _fake_while_loop(cond_fn, body_fn, operands):
+  # print("fake func operands: ", operands)
+  # print("fake func *operands: ", *operands)
+  # while cond_fn(*operands):
+  #   operands = body_fn(*operands)
+  # return operands
+  while cond_fn(operands):
+    operands = body_fn(operands)
+  return operands
+
+
+class WhileLoopTest(unittest.TestCase):
+
+  def test_while_loop_tpu(self):
+
+    device = xm.xla_device()
+    # ten = torch.ones(1, dtype=torch.int32, device=device)
+    # ten[0] = 10
+
+    def cond_fn(x): # x = (xi,)
+      ten = torch.ones(1, dtype=torch.int32, device=device)
+      # ten = torch.tensor([5], dtype=torch.int32, device=device)
+      return x[0] >= ten[0] # ==x[0] # torch.equal(x[0], ten) # x[0] <= ten # 30
+
+    def body_fn(x): # x = (xi,)
+      # onei = torch.tensor(10, dtype=torch.int32, device=device)
+      return (torch.sub(x[0], 1),) # onei,)
+
+    # device = xm.xla_device()
+    # xi = torch.ones(1, dtype=torch.int32, device=device)
+    xi = torch.tensor([5], dtype=torch.int32, device=device)
+    # xi[0] = 5
+    # xi.to(device)
+    # yi = torch.ones(1, dtype=torch.int32, device=device)
+    # yi = torch.tensor([1], dtype=torch.int32, device=device)
+    res = while_loop(cond_fn, body_fn, (xi,))
+    expected = _fake_while_loop(cond_fn, body_fn, xi)
+    self.assertEqual(expected, res)
+
+
+if __name__ == '__main__':
+  test = unittest.main()
+  sys.exit(0 if test.result.wasSuccessful() else 1)

test/tpu/xla_test_job.yaml

@@ -61,6 +61,7 @@ spec:
       python3 /src/pytorch/xla/test/spmd/test_spmd_debugging.py
       python3 /src/pytorch/xla/test/pjrt/test_dtypes.py
       python3 /src/pytorch/xla/test/pjrt/test_dynamic_plugin_tpu.py
+      python3 /src/pytorch/xla/test/test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py
     volumeMounts:
     - mountPath: /dev/shm
       name: dshm

torch_xla/core/xla_builder.py

@@ -762,6 +762,9 @@ def mkop(name, ops, **kwargs):
   builder = kwargs.get('builder', None)
   if builder is None:
     assert ops
+    # if not isinstance(ops, (list, tuple)):
+    #   builder = torch_xla._XLAC._xla_op_builder(ops)
+    # else:
     builder = torch_xla._XLAC._xla_op_builder(ops[0])
   return Op(torch_xla._XLAC._xla_op_create(builder, name, ops, kwargs))
 

torch_xla/core/xla_op_registry.py

@@ -41,6 +41,8 @@ def __call__(self, *args, **kwargs):
         self._computations[key] = computation
         if xu.getenv_as('XLA_OP_PRINT_COMPUTATIONS', bool, False):
           print(xb.get_computation_hlo(computation), file=sys.stderr)
+    print("777777777 args: ", args)
+    print("777777777 type args: ", type(args))
     result = torch_xla._XLAC._xla_user_computation(self._opname, args,
                                                    computation)
     return result[0] if len(result) == 1 else result

torch_xla/csrc/helpers.cpp

@@ -937,6 +937,8 @@ xla::StatusOr<xla::XlaComputation> XlaHelpers::WrapXlaComputation(
                        /*param_index=*/xla::ShapeIndex({input_index}));
   }
 
+  // xla::XlaOp a = xla::GetTupleElement(orig_result, 0);
+
   return builder.Build(orig_result);
 }
 

torch_xla/csrc/init_python_bindings.cpp

@@ -17,6 +17,7 @@
 #include <thread>
 #include <unordered_map>
 #include <vector>
+#include <iostream>
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/strings/str_cat.h"
@@ -671,7 +672,17 @@ py::object XlaNms(const at::Tensor& boxes, const at::Tensor& scores,
 std::vector<at::Tensor> XlaUserComputation(
     const std::string& opname, const std::vector<at::Tensor>& inputs,
     runtime::ComputationClient::ComputationPtr computation) {
+  std::cout << " !!!$$$: " << std::endl;
+  for (int i = 0; i < inputs.size(); i++) {
+    std::cout << inputs[i] << "; " << std::endl;
+    std::cout << inputs[i].type() << "; type !!!" << std::endl;
+    // inputs[i] = (inputs[i]);
+  }
   std::vector<XLATensorPtr> xinputs = GetXlaTensors(inputs, /*want_all=*/true);
+  // std::cout << " !!!$$$###: " << std::endl;
+  // for (int i = 0; i < xinputs.size(); i++) {
+  //   std::cout << DumpUtil::ToText(xinputs[i]->CurrentIrValue().node.get()) << "; ";
+  // }
   std::vector<XLATensorPtr> xresults =
       tensor_methods::user_computation(opname, xinputs, std::move(computation));
   std::vector<at::Tensor> results;
@@ -685,7 +696,14 @@ std::vector<at::Tensor> XlaUserComputation(
 
 runtime::ComputationClient::ComputationPtr CreateComputation(
     const std::string& name, xla::XlaOp root) {
+  std::cout << "w's build func name: " << name << std::endl;
+  // std::cout << "w's build builder name: " << root.builder().name_ << std::endl;
+  // https://github.com/openxla/xla/blob/762bde36adf22792e91c38fe87cabe5af05bfadc/xla/client/xla_builder.cc#L710
   xla::XlaComputation computation = ConsumeValue(root.builder()->Build(root));
+  // std::vector<std::pair<int64_t, int64_t>> input_output_alias_pair;
+  // xla::ProgramShape program_shape = ConsumeValue(computation.GetProgramShape());
+  // computation = ConsumeValue(XlaHelpers::WrapXlaComputation(
+  //   computation, program_shape.parameters(), input_output_alias_pair));
   return std::make_shared<runtime::ComputationClient::Computation>(
       name, std::move(computation));
 }
@@ -876,13 +894,19 @@ void BuildProfilerSubmodule(py::module* m) {
 
 class PyLoweringContext {
  public:
+  PyLoweringContext(const std::string& name) : PyLoweringContext(name, bridge::GetCurrentDevice()) {}
+
+  PyLoweringContext(const std::string& name, torch::lazy::BackendDevice device)
+      : lowering_ctx(name, device) {}
+
   PyLoweringContext() : PyLoweringContext(bridge::GetCurrentDevice()) {}
 
   PyLoweringContext(torch::lazy::BackendDevice device)
       : lowering_ctx("PyLoweringContext", device) {}
 
   // Builds a HLO graph given a set of output tensors.
   void Build(std::vector<at::Tensor> tensors) {
+    // std::cout<< "let's see how many timed this was called? !!!" << GetNameString() << std::endl;
     // Get the backing XLA tensors from the output torch tensor handles
     std::vector<XLATensorPtr> xtensors =
         GetXlaTensors(tensors, /*want_all=*/true);
@@ -894,13 +918,31 @@ class PyLoweringContext {
       ir_values.push_back(value);
     }
 
+    // // check computation name
+    // XLA_ERROR() << computation.proto().name();
+
     // Lower the graph using the output IR values
     for (auto& ir_value : ir_values) {
       xla::XlaOp root = lowering_ctx.GetOutputOp(
           torch::lazy::Output(ir_value.node.get(), ir_value.index));
+      // if (computation.proto().name()=='condctx') {
+      // xla::XlaOp a = xla::GetTupleElement(root, 0); // they are not tupled here
+      // }
       lowering_ctx.AddResult(root);
     }
     computation = ConsumeValue(lowering_ctx.BuildXla());
+
+    std::vector<std::pair<int64_t, int64_t>> input_output_alias_pair;
+    xla::ProgramShape program_shape = ConsumeValue(computation.GetProgramShape());
+    bool should_wrap_parameter = (program_shape.parameters_size() >= 2); // true;
+    if (should_wrap_parameter) {
+      computation = ConsumeValue(XlaHelpers::WrapXlaComputation(
+        computation, program_shape.parameters(), input_output_alias_pair));
+    }
+
+    // // unwrap (pred[])
+    // xla::XlaBuilder builder(computation.proto().name());
+    // xla::XlaOp orig_result = xla::Call(&builder, computation, inner_params);
   }
 
   // Get a mapping from the HLO input parameters to the backing Tensor values.
@@ -983,6 +1025,22 @@ class PyLoweringContext {
     return result;
   }
 
+  void SetNameString(const std::string& name) {
+    lowering_ctx.setnamestring(name);
+  }
+
+  std::string GetNameString() {
+    return lowering_ctx.getnamestring();
+  }
+
+  // LoweringContext GetLoweringCtx() {
+  //   return lowering_ctx;
+  // }
+
+  // LoweringContext SetLoweringCtxName(const std::string name) {
+  //   lowering_ctx.builder().name_ = name;
+  // }
+
  private:
   LoweringContext lowering_ctx;
   xla::XlaComputation computation;
@@ -1027,7 +1085,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("setnamestring", &PyLoweringContext::SetNameString)
+      .def("getnamestring", &PyLoweringContext::GetNameString);
 }
 
 void InitXlaModuleBindings(py::module m) {

torch_xla/csrc/lowering_context.cpp

@@ -154,9 +154,34 @@ void LoweringContext::SetResult(size_t index, xla::XlaOp op) {
 
 xla::StatusOr<xla::XlaComputation> LoweringContext::BuildXla() {
   xla::StatusOr<xla::XlaComputation> xla;
-  if (!root_tuple_.empty()) {
+  // if (builder_.name() == 'bodyctx') {
+  // XLA_ERROR() << builder_.name();
+  // }
+  std::cout << "???" << getnamestring();
+  if (!root_tuple_.empty() & (root_tuple_.size()>1)) {
     xla::XlaOp root = xla::Tuple(builder(), root_tuple_);
+    // xla::XlaOp a = xla::GetTupleElement(root, 0);
     xla = builder()->Build(root);
+  } else if (!root_tuple_.empty() & (root_tuple_.size()==1)) {
+    // xla::XlaOp root = xla::Tuple(builder(), root_tuple_);
+    // xla::XlaOp a = xla::GetTupleElement(root, 0);
+    // const xla::Shape& root_shape = ShapeHelper::ShapeOfXlaOp(root_tuple_.at(0));
+    // xla::XlaBuilder cb("predone");
+    // xla::Shape xla_scalar_shape = xla::ShapeUtil::MakeShape(element_type, {});
+    // xla::XlaOp p0 = xla::Parameter(&cb, 0, xla_scalar_shape, "p0");
+    // Below are to untuple the parameters
+    // xla = builder()->Build(root_tuple_.at(0)); // root);
+    const std::string condctx = "condctx";
+    const std::string bodyctx = "bodyctx";
+    // std::cout << "???" << builder()->name();
+    const std::string currentname = getnamestring();
+    if ((currentname == condctx) or (currentname == bodyctx)) { // == "condctx") {
+      xla = builder()->Build(root_tuple_.at(0)); // root);
+    }
+    else {
+      xla::XlaOp root = xla::Tuple(builder(), root_tuple_);
+      xla = builder()->Build(root);
+    }
   } else {
     xla = builder()->Build();
   }

torch_xla/csrc/lowering_context.h

@@ -5,6 +5,7 @@
 #include <torch/csrc/lazy/backend/lowering_context.h>
 #include <torch/csrc/lazy/core/ir_util.h>
 
+#include <iostream>
 #include <memory>
 #include <string>
 #include <string_view>
@@ -34,6 +35,10 @@ class LoweringContext : public torch::lazy::LoweringContext {
 
   xla::XlaBuilder* builder() { return &builder_; }
 
+  void setnamestring(const std::string& name) { name_ = name; std::cout << "LoweringContext~~~??>>: " << name_ << std::endl;}
+
+  const std::string& getnamestring() { return name_; }
+
   StackFrameIndexBuilder* stack_frame_index_builder() {
     return stack_frame_index_builder_.get();
   }
@@ -121,6 +126,7 @@ class LoweringContext : public torch::lazy::LoweringContext {
       parameters_map_;
   std::vector<xla::XlaOp> root_tuple_;
   OutputMap<xla::XlaOp> emitted_outputs_;
+  std::string name_;
 
   std::shared_ptr<StackFrameIndexBuilder> stack_frame_index_builder_;
 };  // namespace torch_xla

torch_xla/csrc/tensor_methods.cpp

@@ -629,9 +629,14 @@ std::vector<XLATensorPtr> user_computation(
     runtime::ComputationClient::ComputationPtr computation) {
   XLA_CHECK(!inputs.empty());
   std::vector<torch::lazy::Value> input_values;
+  std::vector<const torch::lazy::Node*> root_nodes;
   for (auto& input : inputs) {
-    input_values.push_back(input->GetIrValue());
+    torch::lazy::Value ir_value = input->GetIrValue();
+    input_values.push_back(ir_value);
+    root_nodes.push_back(ir_value.node.get());
   }
+  std::string graph_str = DumpUtil::ToText(root_nodes);
+  std::cout << "inputs' torch::lazy::Node are ##@#@#@#@#: " << graph_str << std::endl;
   torch::lazy::NodePtr node = torch::lazy::MakeNode<UserComputation>(
       torch::lazy::OpKind::Get(opname), input_values, std::move(computation));
   // Cast can be one of the user computation and we don't want to inherit the