Replace XlaValue with upstream lazy::Value (#3572)

* Replace XlaValue with upstream lazy::Value

* Update tests

* Remove comments

* Add index to shape getter

Author: wonjoo-wj

scripts/gen_lazy_tensor.py

@@ -79,6 +79,5 @@ def node_base_ctor_call(self, schema: LazyIrSchema) -> str:
       create_tensor="XLATensor::Create",
       create_aten_from_ltc_tensor="torch_xla::bridge::AtenFromXlaTensor",
       tuple_aten_from_ltc_tensors="torch::lazy::TupleAtenFromLtcTensors",
-      lazy_value_class="torch_xla::XlaValue",
       lazy_tensor_ptr="torch_xla::XLATensorPtr",
       get_device_fn="torch_xla::bridge::GetXlaDevice")

test/cpp/cpp_test_util.cpp

@@ -250,14 +250,14 @@ std::string GetTensorHloGraph(at::Tensor tensor) {
   return DumpUtil::ToHlo({xtensor.GetIrValue()}, xtensor.GetDevice());
 }
 
-XlaValue GetTensorIrValue(const at::Tensor& tensor,
-                          const torch::lazy::BackendDevice& device) {
+torch::lazy::Value GetTensorIrValue(const at::Tensor& tensor,
+                                    const torch::lazy::BackendDevice& device) {
   xla::ComputationClient::DataPtr data = TensorToXlaData(tensor, device);
   return torch::lazy::MakeNode<DeviceData>(std::move(data));
 }
 
 std::vector<xla::ComputationClient::DataPtr> Execute(
-    absl::Span<const XlaValue> roots,
+    absl::Span<const torch::lazy::Value> roots,
     const torch::lazy::BackendDevice& device) {
   LoweringContext lowering_ctx("Execute", device);
   for (auto node : roots) {
@@ -300,7 +300,7 @@ std::vector<at::Tensor> Fetch(
 }
 
 std::vector<at::Tensor> ExecuteAndFetch(
-    absl::Span<const XlaValue> roots,
+    absl::Span<const torch::lazy::Value> roots,
     const torch::lazy::BackendDevice& device) {
   auto results = Execute(roots, device);
   return Fetch(results);

test/cpp/cpp_test_util.h

@@ -85,17 +85,19 @@ std::string GetTensorDotGraph(at::Tensor tensor);
 
 std::string GetTensorHloGraph(at::Tensor tensor);
 
-XlaValue GetTensorIrValue(const at::Tensor& tensor,
-                          const torch::lazy::BackendDevice& device);
+torch::lazy::Value GetTensorIrValue(const at::Tensor& tensor,
+                                    const torch::lazy::BackendDevice& device);
 
 std::vector<xla::ComputationClient::DataPtr> Execute(
-    absl::Span<const XlaValue> roots, const torch::lazy::BackendDevice& device);
+    absl::Span<const torch::lazy::Value> roots,
+    const torch::lazy::BackendDevice& device);
 
 std::vector<at::Tensor> Fetch(
     absl::Span<const xla::ComputationClient::DataPtr> device_data);
 
 std::vector<at::Tensor> ExecuteAndFetch(
-    absl::Span<const XlaValue> roots, const torch::lazy::BackendDevice& device);
+    absl::Span<const torch::lazy::Value> roots,
+    const torch::lazy::BackendDevice& device);
 
 void AssertBackward(const torch::Tensor& xla_output,
                     const std::vector<torch::Tensor>& xla_inputs,

test/cpp/test_ir.cpp

@@ -20,20 +20,24 @@ TEST(IrTest, TestScalarCreate) {
 TEST(IrTest, TestHash) {
   torch::lazy::NodePtr scalar1 = ScalarOp(1.0, xla::F32);
   torch::lazy::NodePtr scalar2 = ScalarOp(2.0, xla::F32);
-  XlaValue add1 = XlaValue(scalar1, 0) + XlaValue(scalar2, 0);
+  torch::lazy::Value add1 =
+      torch::lazy::Value(scalar1, 0) + torch::lazy::Value(scalar2, 0);
 
   torch::lazy::NodePtr scalar3 = ScalarOp(1.0, xla::F32);
   torch::lazy::NodePtr scalar4 = ScalarOp(2.0, xla::F32);
-  XlaValue add2 = XlaValue(scalar3, 0) + XlaValue(scalar4, 0);
+  torch::lazy::Value add2 =
+      torch::lazy::Value(scalar3, 0) + torch::lazy::Value(scalar4, 0);
 
   torch::lazy::NodePtr scalar5 = ScalarOp(11.0, xla::F32);
   torch::lazy::NodePtr scalar6 = ScalarOp(22.0, xla::F32);
-  XlaValue add3 = XlaValue(scalar5, 0) + XlaValue(scalar6, 0);
+  torch::lazy::Value add3 =
+      torch::lazy::Value(scalar5, 0) + torch::lazy::Value(scalar6, 0);
 
   EXPECT_EQ(add1->hash(), add2->hash());
   EXPECT_NE(add1->hash(), add3->hash());
 
-  XlaValue sub = XlaValue(scalar1, 0) - XlaValue(scalar2, 0);
+  torch::lazy::Value sub =
+      torch::lazy::Value(scalar1, 0) - torch::lazy::Value(scalar2, 0);
 
   EXPECT_NE(add1->hash(), sub->hash());
 }
@@ -43,9 +47,10 @@ TEST(IrTest, TestSelectUnselect) {
     at::Tensor a =
         at::rand({4, 16, 3}, at::TensorOptions(at::kFloat)).abs() + 1.0;
 
-    XlaValue v_a = GetTensorIrValue(a, device);
-    XlaValue v_s = torch::lazy::MakeNode<Select>(v_a, /*dim=*/1, /*start=*/3,
-                                                 /*end=*/14, /*stride=*/3);
+    torch::lazy::Value v_a = GetTensorIrValue(a, device);
+    torch::lazy::Value v_s =
+        torch::lazy::MakeNode<Select>(v_a, /*dim=*/1, /*start=*/3,
+                                      /*end=*/14, /*stride=*/3);
 
     auto results = ExecuteAndFetch({v_s}, device);
     at::Tensor b =
@@ -54,8 +59,8 @@ TEST(IrTest, TestSelectUnselect) {
 
     // Paste zeros back into the selected view.
     at::Tensor z = at::zeros_like(b);
-    XlaValue v_z = GetTensorIrValue(z, device);
-    XlaValue v_u =
+    torch::lazy::Value v_z = GetTensorIrValue(z, device);
+    torch::lazy::Value v_u =
         torch::lazy::MakeNode<Unselect>(v_a, v_z, /*dim=*/1, /*start=*/3,
                                         /*end=*/14, /*stride=*/3);
     results = ExecuteAndFetch({v_u}, device);

test/cpp/test_op_by_op_executor.cpp

@@ -17,9 +17,9 @@ TEST(OpByOpExecutorTest, TestSimpleAdd) {
     at::Tensor b = at::rand({4, 16, 3}, at::TensorOptions(at::kFloat));
     at::Tensor c = a + b;
 
-    XlaValue v_a = GetTensorIrValue(a, device);
-    XlaValue v_b = GetTensorIrValue(b, device);
-    XlaValue v_c = v_a + v_b;
+    torch::lazy::Value v_a = GetTensorIrValue(a, device);
+    torch::lazy::Value v_b = GetTensorIrValue(b, device);
+    torch::lazy::Value v_c = v_a + v_b;
 
     auto results_data =
         OpByOpExecutor::Get()->Execute({v_c}, device.toString(), {});
@@ -35,10 +35,10 @@ TEST(OpByOpExecutorTest, TestStack) {
     at::Tensor b = at::rand({4, 8, 3}, at::TensorOptions(at::kFloat));
     at::Tensor c = at::stack({a, b}, 1);
 
-    XlaValue v_a = GetTensorIrValue(a, device);
-    XlaValue v_b = GetTensorIrValue(b, device);
-    XlaValue v_c =
-        torch::lazy::MakeNode<Stack>(std::vector<XlaValue>({v_a, v_b}), 1);
+    torch::lazy::Value v_a = GetTensorIrValue(a, device);
+    torch::lazy::Value v_b = GetTensorIrValue(b, device);
+    torch::lazy::Value v_c = torch::lazy::MakeNode<Stack>(
+        std::vector<torch::lazy::Value>({v_a, v_b}), 1);
 
     auto results_data =
         OpByOpExecutor::Get()->Execute({v_c}, device.toString(), {});
@@ -54,10 +54,10 @@ TEST(OpByOpExecutorTest, TestAsyncStack) {
     at::Tensor b = at::rand({4, 8, 3}, at::TensorOptions(at::kFloat));
     at::Tensor c = at::stack({a, b}, 1);
 
-    XlaValue v_a = GetTensorIrValue(a, device);
-    XlaValue v_b = GetTensorIrValue(b, device);
-    XlaValue v_c =
-        torch::lazy::MakeNode<Stack>(std::vector<XlaValue>({v_a, v_b}), 1);
+    torch::lazy::Value v_a = GetTensorIrValue(a, device);
+    torch::lazy::Value v_b = GetTensorIrValue(b, device);
+    torch::lazy::Value v_c = torch::lazy::MakeNode<Stack>(
+        std::vector<torch::lazy::Value>({v_a, v_b}), 1);
 
     auto async =
         OpByOpExecutor::Get()->ExecuteAsync({v_c}, device.toString(), {});

torch_xla/csrc/debug_util.cpp

@@ -55,13 +55,13 @@ std::string DebugUtil::GetTensorsGraphInfo(absl::Span<const XLATensor> tensors,
                                            const std::vector<size_t>* indices,
                                            GraphFormat format) {
   std::vector<const torch::lazy::Node*> root_nodes;
-  std::vector<XlaValue> root_values;
+  std::vector<torch::lazy::Value> root_values;
   std::vector<torch::lazy::hash_t> root_hashes;
   xla::util::Unique<torch::lazy::BackendDevice> unique_device;
   if (indices != nullptr) {
     for (auto index : *indices) {
       const XLATensor& tensor = tensors[index];
-      XlaValue ir_value = tensor.CurrentIrValue();
+      torch::lazy::Value ir_value = tensor.CurrentIrValue();
       if (ir_value) {
         root_nodes.push_back(ir_value.node.get());
         root_hashes.push_back(ir_value.hash());
@@ -71,7 +71,7 @@ std::string DebugUtil::GetTensorsGraphInfo(absl::Span<const XLATensor> tensors,
     }
   } else {
     for (auto& tensor : tensors) {
-      XlaValue ir_value = tensor.CurrentIrValue();
+      torch::lazy::Value ir_value = tensor.CurrentIrValue();
       if (ir_value) {
         root_nodes.push_back(ir_value.node.get());
         root_hashes.push_back(ir_value.hash());