Add getDynamicValue for SizeMul and SizeDiv nodes. (#4042)

Author: vanbasten23

test/cpp/test_ir.cpp

@@ -250,6 +250,7 @@ TEST_F(IrTest, TestSizeMulNode) {
       std::dynamic_pointer_cast<torch::lazy::DimensionNode>(size_node_mul);
 
   EXPECT_EQ(dim_node_mul->getStaticValue(), 12);
+  EXPECT_EQ(dim_node_mul->getDynamicValue(), 12);
 
   ForEachDevice([&](const torch::lazy::BackendDevice& device) {
     // Lower the SizeAddNode and execute the GetDimensionSize.
@@ -258,6 +259,31 @@ TEST_F(IrTest, TestSizeMulNode) {
   });
 }
 
+TEST_F(IrTest, TestSizeMulNodeDynamic) {
+  int64_t num_non_zero_element = 1;
+  int64_t num_row = 10;
+  int64_t num_col = 10;
+  torch::lazy::NodePtr nonzero_node =
+      CreateNonZeroNode2d(num_non_zero_element, num_row, num_col);
+  torch::lazy::Value node_with_dynamism = torch::lazy::Value(nonzero_node, 0);
+
+  // static value = 100, dynamic value = 1
+  torch::lazy::NodePtr size_node_nonzero_0 =
+      torch::lazy::MakeNode<SizeNode>(node_with_dynamism, 0);
+  // static value = 2, dynamic value = 2
+  torch::lazy::NodePtr size_node_nonzero_1 =
+      torch::lazy::MakeNode<SizeNode>(node_with_dynamism, 1);
+
+  torch::lazy::NodePtr node_mul = torch::lazy::MakeNode<SizeMul>(
+      torch::lazy::Value(size_node_nonzero_0, 0),
+      torch::lazy::Value(size_node_nonzero_1, 0));
+
+  std::shared_ptr<torch::lazy::DimensionNode> dim_node_mul =
+      std::dynamic_pointer_cast<torch::lazy::DimensionNode>(node_mul);
+  EXPECT_EQ(dim_node_mul->getStaticValue(), 200);
+  EXPECT_EQ(dim_node_mul->getDynamicValue(), 2);
+}
+
 TEST_F(IrTest, TestSizeDivNode) {
   torch::lazy::NodePtr scalar_node =
       ScalarOp(1.0, xla::ShapeUtil::MakeShape(xla::F32, {12, 5}));
@@ -271,6 +297,7 @@ TEST_F(IrTest, TestSizeDivNode) {
       std::dynamic_pointer_cast<torch::lazy::DimensionNode>(size_node_div);
 
   EXPECT_EQ(dim_node_div->getStaticValue(), 2);
+  EXPECT_EQ(dim_node_div->getDynamicValue(), 2);
 
   ForEachDevice([&](const torch::lazy::BackendDevice& device) {
     // Lower the SizeAddNode and execute the GetDimensionSize.
@@ -279,5 +306,30 @@ TEST_F(IrTest, TestSizeDivNode) {
   });
 }
 
+TEST_F(IrTest, TestSizeDivNodeDynamic) {
+  int64_t num_non_zero_element = 1;
+  int64_t num_row = 10;
+  int64_t num_col = 10;
+  torch::lazy::NodePtr nonzero_node =
+      CreateNonZeroNode2d(num_non_zero_element, num_row, num_col);
+  torch::lazy::Value node_with_dynamism = torch::lazy::Value(nonzero_node, 0);
+
+  // static value = 100, dynamic value = 1
+  torch::lazy::NodePtr size_node_nonzero_0 =
+      torch::lazy::MakeNode<SizeNode>(node_with_dynamism, 0);
+  // static value = 2, dynamic value = 2
+  torch::lazy::NodePtr size_node_nonzero_1 =
+      torch::lazy::MakeNode<SizeNode>(node_with_dynamism, 1);
+
+  torch::lazy::NodePtr node_div = torch::lazy::MakeNode<SizeDiv>(
+      torch::lazy::Value(size_node_nonzero_0, 0),
+      torch::lazy::Value(size_node_nonzero_1, 0));
+
+  std::shared_ptr<torch::lazy::DimensionNode> dim_node_div =
+      std::dynamic_pointer_cast<torch::lazy::DimensionNode>(node_div);
+  EXPECT_EQ(dim_node_div->getStaticValue(), 50);
+  EXPECT_EQ(dim_node_div->getDynamicValue(), 0);
+}
+
 }  // namespace cpp_test
 }  // namespace torch_xla

torch_xla/csrc/ops/dynamic_ir.cpp

@@ -98,6 +98,14 @@ SizeMul::SizeMul(torch::lazy::Value a, torch::lazy::Value b)
   upper_bound_ = dim_node_0->getStaticValue() * dim_node_1->getStaticValue();
 };
 
+int64_t SizeMul::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);
+  return dim_node_0->getDynamicValue() * dim_node_1->getDynamicValue();
+}
+
 std::string SizeMul::ToString() const { return "SizeMul"; }
 
 XlaOpVector SizeMul::Lower(LoweringContext* loctx) const {
@@ -121,6 +129,14 @@ SizeDiv::SizeDiv(torch::lazy::Value a, torch::lazy::Value b)
   upper_bound_ = dim_node_0->getStaticValue() / dim_node_1->getStaticValue();
 };
 
+int64_t SizeDiv::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);
+  return dim_node_0->getDynamicValue() / dim_node_1->getDynamicValue();
+}
+
 std::string SizeDiv::ToString() const { return "SizeDiv"; }
 
 XlaOpVector SizeDiv::Lower(LoweringContext* loctx) const {

torch_xla/csrc/ops/dynamic_ir.h

@@ -68,6 +68,7 @@ class SizeAdd : public XlaNode, public torch::lazy::DimensionNode {
 class SizeMul : public XlaNode, public torch::lazy::DimensionNode {
  public:
   SizeMul(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;
@@ -80,6 +81,7 @@ class SizeMul : public XlaNode, public torch::lazy::DimensionNode {
 class SizeDiv : public XlaNode, public torch::lazy::DimensionNode {
  public:
   SizeDiv(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;