POC of dynamicShape+functionalization integration (#4611)

* Support dynamism on "transpose", unsqueeze_copy, and view_copy_symint.

Support dynamism on "transpose" operation. (#4606)

* added failing test_t_copy test

* made transpose work

* disable the backward pass.

* fix linter

Add dynamic shape test for unsqueeze_copy (#4608)

* added failing test

* test correctness.

Support dynamism on view_copy_symint (#4629)

* Disable failing XLA tests

* Re-enable dynamo tests (#4454)

* Turn on keep going

* Re-enable more dynamic shape tests (#4558)

* fix size node ne op kind to be size_ne

* reenable test_nonzero_cast

* fix linter

* disable new failing test after rebase temporarily.

* fix linter

* [Functionalization] Fix ScatterReduce (#4576)

Summary:
ScatterReduce::reduce_ uses a unsafe type c10::string_view to store the
string. Replace it with std::string and re-enable all previous failing tests.

Test Plan:
TPU_LIBRARY_PATH=/home/ptxla/.local/lib/python3.8/site-packages/libtpu/libtpu.so PJRT_DEVICE=CPU test/cpp/build/test_ptxla --gtest_filter=AtenXlaTensorTest.TestScatterReduce*

* Revert "Turn on keep going"

This reverts commit 224bdf3.

* Turn on keep going

* Revert "Turn on keep going"

This reverts commit 224bdf3.

* Turn on keep going

* Revert "Turn on keep going"

This reverts commit 224bdf3.

* code compiles

* wrote new test TestDynamicShapes.test_sizeMod and it succeeds.

* the first view_copy_symint test test_view_copy_symint_with_dyn_input_shape passed

* added two more tests

* run linter

* added a new failing test with negative shape

* fix all tests

* clean up

* clean up

* fix a small test failure.

* fix pr comments

* updated the tests for dynamic input and static input shape

* not support dynamic input and static input shape.

* fix linter

* fix pr comments

* fix build error

* resolve a rebase conflict

* run linter

* add one more test

* add one more test and it passes.

* fix linter

* remove experiemental test.

---------

Co-authored-by: Wonjoo Lee <wonjoo@google.com>
Co-authored-by: Jiewen Tan <jwtan@google.com>

fix a merge error

fix other merge error during rebase.

* Track type of SymNode for XLASymNodeImpl::mod (#4732)

Author: vanbasten23

test/test_dynamic_shape_models.py

@@ -1,4 +1,5 @@
 import argparse
+import os
 import sys
 
 parser = argparse.ArgumentParser(add_help=False)
@@ -12,6 +13,9 @@
 import torch_xla.core.xla_model as xm
 import torch_xla.debug.metrics as met
 
+# It enables us to run python implementations of CompositeAutogradImplicit ops.
+# CompositeAutogradImplicit means we don't have an explicit backward formula for an op instead an op is composed of a bunch of ops that do have backward formulas and combines this formulas is equivalent to differentiating the op explicitly.
+pd = torch._C._EnablePythonDispatcher()
 xla_dev = xm.xla_device()
 
 
@@ -92,9 +96,6 @@ def test_forward_pass_dynamic_input_compile_once(self):
                            met.metric_data('CompileTime')[0],
                            'number of compilation should not increase.')
 
-  @unittest.skip(
-      "disable it due to https://github.com/pytorch/xla/pull/4322#issuecomment-1374312614."
-  )
   def test_backward_pass_with_dynamic_input(self):
     num_features = 2
     num_test_samples = 5
@@ -156,5 +157,8 @@ def create_dynamic_test_data(self,
 
 
 if __name__ == '__main__':
+  assert os.environ['XLA_EXPERIMENTAL'] != ''
   test = unittest.main(verbosity=FLAGS.verbosity, exit=False)
+  # DISABLE PYTHON DISPATCHER FLAG
+  del pd
   sys.exit(0 if test.result.wasSuccessful() else 1)

test/test_dynamic_shapes.py

@@ -144,6 +144,22 @@ def test_empty_symint(self):
     self.assertIsInstance(t2.shape[1], int)
     self.assertEqual(t2.shape[1], 2)
 
+  def test_t_copy(self):
+    t1 = torch.tensor([[1, 0, 0, 5, 0, 6], [1, 3, 2, 0, 0, 1]], device=dev)
+    t2 = torch.nonzero(t1)
+    # t2.shape=torch.Size([<=12, 2]) with real size [7, 2]
+    self.assertEqual(str(t2.shape[0]), '<=12')
+    self.assertEqual(str(t2.shape[1]), '2')
+
+    t2_t = torch.t(t2)
+
+    self.assertIsInstance(t2_t.shape[0], int)
+    self.assertIsInstance(t2_t.shape[1], torch.SymInt)
+    self.assertEqual(str(t2_t.shape[0]), '2')
+    self.assertEqual(str(t2_t.shape[1]), '<=12')
+    self.assertEqual(t2_t.shape[0], 2)
+    self.assertEqual(t2_t.shape[1], 7)
+
   def test_nonzero_shape(self):
     x = torch.tensor((0, 1, 2, 0, 3, 4), device=xm.xla_device())
     x_dim0_shape = torch_xla._XLAC._get_xla_tensor_dimension_size(
@@ -175,13 +191,173 @@ def test_expand_symint_correctness(self):
     t2 = torch.zeros([size1, size2], device=dev)
     t2[3][0] = 1
     t2[3][1] = 1
-    # t2 has size [<=10, 2]
+    # t3 has size [<=10, 2]
     t3 = torch.nonzero(t2)
     t4 = torch.ones([size1, size2], device=dev)
     expand_out_xla = t4.expand(t3.shape[0], size1, size2)
     self.assertEqual(t3.shape[0], 2)
     self.assertEqual(expand_out_aten.cpu(), expand_out_xla.cpu())
 
+  def test_unsqueeze_copy_dynamism(self):
+    t1 = torch.tensor([[1, 0, 0, 5, 0, 6], [1, 3, 2, 0, 0, 1]], device=dev)
+    t2 = torch.nonzero(t1)
+    # t2.shape=torch.Size([<=12, 2]) with real size [7, 2]
+
+    t2_unsqueeze = torch.unsqueeze(t2, 0)
+
+    self.assertEqual(len(t2_unsqueeze.size()), 3)
+    self.assertIsInstance(t2_unsqueeze.shape[0], int)
+    self.assertIsInstance(t2_unsqueeze.shape[1], torch.SymInt)
+    self.assertIsInstance(t2_unsqueeze.shape[2], int)
+    self.assertEqual(str(t2_unsqueeze.shape[0]), '1')
+    self.assertEqual(str(t2_unsqueeze.shape[1]), '<=12')
+    self.assertEqual(str(t2_unsqueeze.shape[2]), '2')
+    self.assertEqual(t2_unsqueeze.shape[0], 1)
+    self.assertEqual(t2_unsqueeze.shape[1], 7)
+    self.assertEqual(t2_unsqueeze.shape[2], 2)
+
+    # test correctness
+    t3 = torch.tensor([[1, 0, 0, 5, 0, 6], [1, 3, 2, 0, 0, 1]])
+    t4 = torch.nonzero(t3)
+    t4_unsqueeze = torch.unsqueeze(t4, 0)
+    self.assertEqual(t2_unsqueeze.cpu(), t4_unsqueeze.cpu())
+
+  def test_view_copy_symint_with_static_input_dyn_input_shape(self):
+    # If the input tensor and shape are “statically” incompatible, a compilation error is raised.
+    t1 = torch.tensor([1, 0, 3, 5, 0, 6], device=dev)
+    # t2.shape=torch.Size([<=6, 1]) with real size [4, 1]
+    # t2 = [[0], [2], [3], [5]]
+    t2 = torch.nonzero(t1)
+    t3 = torch.randint(10, (2, 2), device=dev)
+    self.assertRaises(RuntimeError, lambda: t3.view(t2.shape[0]))
+
+    # If their “dynamic” values are incompatible, a RuntimeError is raised.
+    t4 = torch.randint(10, (2, 3), device=dev)
+    self.assertRaises(RuntimeError, lambda: t4.view(t2.shape[0]))
+
+    # verify if dynamism is propagated correctly.
+    t5 = torch.tensor([1, 1, 3, 5, 1, 6], device=dev)
+    t6 = torch.nonzero(t5)
+    t7 = torch.ones((2, 3), device=dev)
+    t8 = t7.view(t6.shape[0])
+    self.assertIsInstance(t8.shape[0], torch.SymInt)
+    self.assertEqual(str(t8.shape[0]), '<=6')
+    self.assertEqual(t8.shape[0], 6)
+
+    # verify correctness.
+    t5_aten = torch.tensor([1, 1, 3, 5, 1, 6])
+    t6_aten = torch.nonzero(t5_aten)
+    t7_aten = torch.ones((2, 3))
+    t8_aten = t7_aten.view(t6_aten.shape[0])
+    self.assertEqual(t8.cpu(), t8_aten.cpu())
+
+  def test_view_copy_symint_with_static_input_dyn_input_shape2(self):
+    # If the input tensor and shape are “statically” incompatible, a compilation error is raised.
+    t1 = torch.tensor([[1, 0, 3]], device=dev)
+    # t2.shape=torch.Size([<=3, 2]) with real size [2, 2]
+    # t2 = [[0, 0], [0, 2]]
+    t2 = torch.nonzero(t1)
+    t3 = torch.ones((2, 4), device=dev)
+    # Should fail in pytorch utils.infer_size
+    self.assertRaises(RuntimeError, lambda: t3.view(t2.shape))
+
+    # If their “dynamic” values are incompatible, a RuntimeError is raised.
+    t4 = torch.ones((2, 3), device=dev)
+    # Also fails in pytorch utils.infer_size
+    self.assertRaises(RuntimeError, lambda: t4.view(t2.shape))
+
+    # verify if dynamism is propagated correctly.
+    t5 = torch.tensor([[1, 1, 3]], device=dev)
+    t6 = torch.nonzero(t5)
+    # t6.shape=[<=3, 2] with real size [3, 2]
+    t7 = torch.ones((2, 3), device=dev)
+    t8 = t7.view(t6.shape)
+    self.assertIsInstance(t8.shape[0], torch.SymInt)
+    self.assertEqual(str(t8.shape[0]), '<=3')
+    self.assertEqual(t8.shape[0], 3)
+    self.assertIsInstance(t8.shape[1], int)
+    self.assertEqual(str(t8.shape[1]), '2')
+    self.assertEqual(t8.shape[1], 2)
+
+    # verify correctness.
+    t5_aten = torch.tensor([[1, 1, 3]])
+    t6_aten = torch.nonzero(t5_aten)
+    t7_aten = torch.ones((2, 3))
+    t8_aten = t7_aten.view(t6_aten.shape)
+    self.assertEqual(t8.cpu(), t8_aten.cpu())
+
+  def test_view_copy_symint_with_dyn_input_static_input_shape(self):
+    # If the input tensor is dynamic and input shape is static,
+    # it should fail because we will not likely have this case
+    # in reality so we don't support this feature.
+    t1 = torch.tensor([1, 1, 3, 5, 1, 6], device=dev)
+    # t2.shape=torch.Size([<=6, 1]) with real size [6, 1]
+    t2 = torch.nonzero(t1)
+    self.assertRaises(RuntimeError, lambda: t2.view(2, 3))
+
+  def test_view_copy_symint_with_dyn_input_dyn_input_shape(self):
+    # If the input tensor and shape are “statically” incompatible, a compilation error is raised.
+    t1 = torch.tensor([1, 0, 3, 5, 0, 6], device=dev)
+    # t2.shape=torch.Size([<=6, 1]) with real size [4, 1]
+    # t2 = [[0], [2], [3], [5]]
+    t2 = torch.nonzero(t1)
+    t3 = torch.tensor([1, 0, 3, 5, 0, 6, 7], device=dev)
+    # t4.shape=torch.Size([<=7, 1]) with real size [5, 1]
+    t4 = torch.nonzero(t3)
+    self.assertRaises(RuntimeError, lambda: t2.view(t4.shape[0]))
+
+    # If their “dynamic” values are incompatible, a RuntimeError is raised.
+    t5 = torch.tensor([1, 2, 3, 4, 5, 6, 0], device=dev)
+    # t6.shape=torch.Size([<=7, 1]) with real size [6, 1]
+    t6 = torch.nonzero(t5)
+    # statically compatible but dynamically incompatible.
+    # It will fail in pytorch layer.
+    self.assertRaises(RuntimeError, lambda: t6.view(t4.shape[0]))
+
+    # verify if dynamism is propagated correctly.
+    t7 = torch.tensor([1, 0, 3, 5, 0, 6, 7], device=dev)
+    t8 = torch.nonzero(t7)
+    # t8.shape=torch.Size([<=7, 1]) with real size [5, 1]
+    t9 = t8.view(t4.shape[0])
+    self.assertIsInstance(t9.shape[0], torch.SymInt)
+    self.assertEqual(str(t9.shape[0]), '<=7')
+    self.assertEqual(t9.shape[0], 5)
+
+    # verify correctness.
+    t7_aten = torch.tensor([1, 0, 3, 5, 0, 6, 7])
+    t8_aten = torch.nonzero(t7_aten)
+    # t8_aten.size=[5, 1]
+    t3_aten = torch.tensor([1, 0, 3, 5, 0, 6, 7])
+    t4_aten = torch.nonzero(t3_aten)
+    # t4_aten.size=[5, 1]
+    t9_aten = t8_aten.view(t4_aten.shape[0])
+    self.assertEqual(t9.cpu(), t9_aten.cpu())
+
+  def test_sizeMod(self):
+    met.clear_all()
+
+    size1 = 5
+    size2 = 2
+    t1 = torch.zeros([size1, size2], device=dev)
+    t1[3][0] = 1
+    # t2 has size [<=10, 2] with real size [1, 2]
+    t2 = torch.nonzero(t1)
+    # Create a SizeMod IR node.
+    # t2.shape[1] generates a SizeConstant node.
+    dyn_size = t2.shape[0] % t2.shape[1]
+    self.assertGreater(met.counter_value("xla::size_mod"), 0)
+    # Exercises SizeMod::getDynamicValue.
+    dynamic_size = int(dyn_size)
+    self.assertEqual(dynamic_size, 1)
+    self.assertEqual(str(dyn_size), '<=0')
+
+    # t3 has size [<=10, 2] with real size [1, 2]
+    t3 = torch.nonzero(t1)
+    dyn_size = t2.shape[0] % t3.shape[0]
+    dynamic_size = int(dyn_size)
+    self.assertEqual(dynamic_size, 0)
+    self.assertEqual(str(dyn_size), '<=0')
+
   def test_sizeGe(self):
     met.clear_all()
 

torch_xla/csrc/aten_xla_type.cpp

@@ -3186,12 +3186,19 @@ std::tuple<at::Tensor, at::Tensor> XLANativeFunctions::var_mean(
 }
 
 at::Tensor XLANativeFunctions::view_copy_symint(const at::Tensor& self,
-                                                at::SymIntArrayRef sym_size) {
-  // TODO: support symbolic sizes
-  auto size = C10_AS_INTARRAYREF_SLOW(sym_size);
+                                                at::SymIntArrayRef shape) {
   TORCH_LAZY_FN_COUNTER("xla::");
-  return bridge::AtenFromXlaTensor(tensor_methods::view(
-      bridge::GetXlaTensor(self), XlaHelpers::I64List(size)));
+  c10::optional<at::IntArrayRef> int_shape = c10::asIntArrayRefSlowOpt(shape);
+  bool input_shape_static = int_shape.has_value();
+  XLATensorPtr xla_input = bridge::GetXlaTensor(self);
+  bool input_has_dyn_shape = xla_input->shape().get().is_dynamic();
+
+  XLA_CHECK(!(input_has_dyn_shape && input_shape_static))
+      << "This view op has dynamic input tensor but static input shape. This "
+         "behavior is currently unsupported; if the user believes this must be "
+         "supported, please file a feature request against PyTorch/XLA.";
+  return bridge::AtenFromXlaTensor(
+      tensor_methods::view_symint(xla_input, shape));
 }
 
 at::Tensor XLANativeFunctions::where(const at::Tensor& condition,

torch_xla/csrc/helpers.cpp

@@ -254,27 +254,27 @@ xla::XlaOp XlaHelpers::ReshapeToRank(xla::XlaOp input, int64_t expected_rank,
 absl::optional<XlaHelpers::DynamicReshapeInfo>
 XlaHelpers::GetDynamicReshapeInfo(const xla::Shape& input_shape,
                                   absl::Span<const int64_t> output_sizes) {
-  int64_t input_dynamic_dimension = GetDynamicDimension(input_shape);
-  if (input_dynamic_dimension < 0) {
+  int64_t input_dyndim_idx = GetDynamicDimension(input_shape);
+  if (input_dyndim_idx < 0) {
     return absl::nullopt;
   }
   DynamicReshapeInfo info;
   info.output_shape =
       xla::ShapeUtil::MakeShape(input_shape.element_type(), output_sizes);
   if (info.output_shape.rank() > 0) {
-    int64_t size_at_dyndim = 1;
-    for (int64_t i = 0; i <= input_dynamic_dimension; ++i) {
-      size_at_dyndim *= input_shape.dimensions(i);
+    int64_t size_prod_until_dyndim = 1;
+    for (int64_t i = 0; i <= input_dyndim_idx; ++i) {
+      size_prod_until_dyndim *= input_shape.dimensions(i);
     }
     int64_t dynamic_dimension = -1;
     int64_t out_size = 1;
     for (int64_t i = 0; i < output_sizes.size(); ++i) {
-      XLA_CHECK_LE(out_size, size_at_dyndim / input_shape.dimensions(
-                                                  input_dynamic_dimension))
+      XLA_CHECK_LE(out_size, size_prod_until_dyndim /
+                                 input_shape.dimensions(input_dyndim_idx))
           << "Unable to map dynamic dimension of shape " << input_shape
           << " to output sizes (" << absl::StrJoin(output_sizes, ", ") << ")";
       out_size *= output_sizes[i];
-      if (out_size >= size_at_dyndim) {
+      if (out_size >= size_prod_until_dyndim) {
         dynamic_dimension = i;
         break;
       }

torch_xla/csrc/ops/dynamic_ir.cpp

@@ -290,4 +290,41 @@ XlaOpVector SizeDiv::Lower(LoweringContext* loctx) const {
   return ReturnOp(xla::Div(input1, input2), loctx);
 }
 
+SizeMod::SizeMod(torch::lazy::Value a, torch::lazy::Value b)
+    : XlaNode(
+          torch::lazy::OpKind{c10::Symbol::fromQualString("aten::size_mod")},
+          {a, b},
+          xla::ShapeUtil::MakeShape(GetShapeDimensionType(/*device=*/nullptr),
+                                    {}),
+          1) {
+  const torch::lazy::DimensionNode* dim_node_0 = DimCast(operand(0));
+  const torch::lazy::DimensionNode* dim_node_1 = DimCast(operand(1));
+  // SizeDiv can only be perfomed between two DimensionNode
+  XLA_CHECK(dim_node_0);
+  XLA_CHECK(dim_node_1);
+  // We don't need to hash upper_bound_ and because it is computed
+  // from input shapes and input Node already hash its shape.
+  XLA_CHECK(dim_node_1->getStaticValue() != 0)
+      << "Can't divide a dimension by zero";
+  upper_bound_ = dim_node_0->getStaticValue() % dim_node_1->getStaticValue();
+};
+
+int64_t SizeMod::getDynamicValue() const {
+  const torch::lazy::DimensionNode* dim_node_0 = DimCast(operand(0));
+  const torch::lazy::DimensionNode* dim_node_1 = DimCast(operand(1));
+  XLA_CHECK(dim_node_0);
+  XLA_CHECK(dim_node_1);
+  XLA_CHECK(dim_node_1->getDynamicValue() != 0)
+      << "Can't mod a dynamic dimension by zero";
+  return dim_node_0->getDynamicValue() % dim_node_1->getDynamicValue();
+}
+
+std::string SizeMod::ToString() const { return "aten::size_mod"; }
+
+XlaOpVector SizeMod::Lower(LoweringContext* loctx) const {
+  auto input1 = loctx->GetOutputOp(operand(0));
+  auto input2 = loctx->GetOutputOp(operand(1));
+  return ReturnOp(xla::Rem(input1, input2), loctx);
+}
+
 }  // namespace torch_xla

torch_xla/csrc/ops/dynamic_ir.h

@@ -165,6 +165,19 @@ class SizeDiv : public XlaNode, public torch::lazy::DimensionNode {
   int64_t upper_bound_;
 };
 
+class SizeMod : public XlaNode, public torch::lazy::DimensionNode {
+ public:
+  SizeMod(torch::lazy::Value a, torch::lazy::Value b);
+  int64_t getDynamicValue() const override;
+  int64_t getStaticValue() const override { return upper_bound_; }
+  bool isSymbolic() const override { return true; }
+  std::string ToString() const override;
+  virtual XlaOpVector Lower(LoweringContext* loctx) const override;
+
+ private:
+  int64_t upper_bound_;
+};
+
 const torch::lazy::DimensionNode* DimCast(torch::lazy::Output output);
 const torch::lazy::DimensionNode* DimCast(const torch::lazy::Node* node);
 const std::shared_ptr<torch::lazy::DimensionNode> DimCast(

torch_xla/csrc/ops/permute.cpp

@@ -44,9 +44,12 @@ std::string Permute::ToString() const {
 
 xla::Shape Permute::MakePermuteShape(const xla::Shape& source_shape,
                                      absl::Span<const int64_t> permutation) {
-  return XlaHelpers::GetDynamicReshape(
-      source_shape,
-      XlaHelpers::Permute(permutation, source_shape.dimensions()));
+  auto output_static_dims =
+      XlaHelpers::Permute(permutation, source_shape.dimensions());
+  auto output_dyn_dims =
+      XlaHelpers::Permute(permutation, source_shape.dynamic_dimensions());
+  return xla::ShapeUtil::MakeShape(source_shape.element_type(),
+                                   output_static_dims, output_dyn_dims);
 }
 
 }  // namespace torch_xla

torch_xla/csrc/tensor.cpp

@@ -669,8 +669,14 @@ c10::SymNode XLASymNodeImpl::floordiv(const c10::SymNode& other) {
 }
 
 c10::SymNode XLASymNodeImpl::mod(const c10::SymNode& other) {
-  XLA_CHECK(false) << "XLASymNodeImpl::" << __FUNCTION__
-                   << " has not been implemented.";
+  TORCH_LAZY_FN_COUNTER("xla::size_");
+  torch_xla::XLASymNodeImpl* p_other =
+      dynamic_cast<XLASymNodeImpl*>(other.get());
+  XLA_CHECK(is_int()) << __FUNCTION__ << " with non-int NYI";
+  XLA_CHECK(p_other->is_int()) << __FUNCTION__ << " with non-int NYI";
+  torch::lazy::NodePtr n_mod =
+      torch::lazy::MakeNode<SizeMod>(node(), p_other->node());
+  return c10::make_intrusive<XLASymNodeImpl>(n_mod, PyType::INT);
 }
 
 c10::SymNode XLASymNodeImpl::eq(const c10::SymNode& other) {