Replace XlaValue with upstream lazy::Value

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")

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());

torch_xla/csrc/init_python_bindings.cpp

@@ -84,15 +84,15 @@ torch::lazy::hash_t GetOperandHashes(const OpList& operands,
 
 }  // namespace
 
-const xla::Shape& XlaValue::xla_shape() const {
-  XlaNode* casted = dynamic_cast<XlaNode*>(node.get());
-  return casted->xla_shape(index);
-}
+// const xla::Shape& XlaValue::xla_shape() const {
+//   XlaNode* casted = dynamic_cast<XlaNode*>(node.get());
+//   return casted->xla_shape(index);
+// }
 
-const xla::Shape& XlaValue::xla_node_shape() const {
-  XlaNode* casted = dynamic_cast<XlaNode*>(node.get());
-  return casted->xla_shape();
-}
+// const xla::Shape& XlaValue::xla_node_shape() const {
+//   XlaNode* casted = dynamic_cast<XlaNode*>(node.get());
+//   return casted->xla_shape();
+// }
 
 XlaNode::XlaNode(torch::lazy::OpKind op, OpList operands,
                  std::vector<torch::lazy::Shape>&& shapes, xla::Shape xla_shape,
@@ -102,7 +102,7 @@ XlaNode::XlaNode(torch::lazy::OpKind op, OpList operands,
       node_hash_(torch::lazy::HashCombine(op.hash(), hash_seed)),
       dag_hash_(GetOperandHashes(operands, node_hash_)) {
   // We have to call AddOperand here since upstream OpList is
-  // an array of torch::lazy::Value while we uses XlaValue.
+  // an array of torch::lazy::Value while we uses torch::lazy::Value.
   for (auto& operand : operands) {
     AddOperand(operand.node, operand.index);
   }
@@ -116,7 +116,7 @@ XlaNode::XlaNode(torch::lazy::OpKind op, OpList operands,
       node_hash_(torch::lazy::HashCombine(op.hash(), hash_seed)),
       dag_hash_(GetOperandHashes(operands, node_hash_)) {
   // We have to call AddOperand here since upstream OpList is
-  // an array of torch::lazy::Value while we uses XlaValue.
+  // an array of torch::lazy::Value while we uses torch::lazy::Value.
   for (auto& operand : operands) {
     AddOperand(operand.node, operand.index);
   }
@@ -131,7 +131,7 @@ XlaNode::XlaNode(torch::lazy::OpKind op, OpList operands,
       node_hash_(torch::lazy::HashCombine(op.hash(), hash_seed)),
       dag_hash_(GetOperandHashes(operands, node_hash_)) {
   // We have to call AddOperand here since upstream OpList is
-  // an array of torch::lazy::Value while we uses XlaValue.
+  // an array of torch::lazy::Value while we uses torch::lazy::Value.
   for (auto& operand : operands) {
     AddOperand(operand.node, operand.index);
   }
@@ -233,4 +233,9 @@ ScopePusher::~ScopePusher() { PopScope(); }
 
 void ScopePusher::ResetScopes() { ResetScopeContext(); }
 
+const xla::Shape& GetXlaShape(const torch::lazy::Value& value) {
+  XlaNode* casted = dynamic_cast<XlaNode*>(value.node.get());
+  return casted->xla_shape();
+}
+
 }  // namespace torch_xla

torch_xla/csrc/ir.cpp

@@ -35,21 +35,24 @@ using OutputMap =
     std::unordered_map<torch::lazy::Output, T, torch::lazy::Output::Hasher>;
 
 // Represents an input/operand for a XlaNode object.
-struct XlaValue : public torch::lazy::Value {
-  XlaValue() = default;
-  XlaValue(torch::lazy::NodePtr node, size_t index = 0)
-      : torch::lazy::Value(std::dynamic_pointer_cast<torch::lazy::Node>(node),
-                           index) {}
-
-  // Retrieves the shape of this value. If the IR XlaNode generating the value
-  // is a multi-output node, the shape returned by this API will not be the full
-  // tuple shape, but only the shape at index referred by this value.
-  // To retrieve the full tuple shape in that case, use the node_shape() API.
-  const xla::Shape& xla_shape() const;
-  const xla::Shape& xla_node_shape() const;
-};
-
-using OpList = absl::Span<const XlaValue>;
+// struct XlaValue : public torch::lazy::Value {
+//   torch::lazy::Value() = default;
+//   torch::lazy::Value(torch::lazy::NodePtr node, size_t index = 0)
+//       :
+//       torch::lazy::Value(std::dynamic_pointer_cast<torch::lazy::Node>(node),
+//                            index) {}
+
+//   // Retrieves the shape of this value. If the IR XlaNode generating the
+//   value
+//   // is a multi-output node, the shape returned by this API will not be the
+//   full
+//   // tuple shape, but only the shape at index referred by this value.
+//   // To retrieve the full tuple shape in that case, use the node_shape() API.
+//   // const xla::Shape& xla_shape() const;
+//   // const xla::Shape& xla_node_shape() const;
+// };
+
+using OpList = absl::Span<const torch::lazy::Value>;
 
 // A node in the graph. Nodes for operations which requires extra data to be
 // stored for lowering, should inherit from this class and add operation
@@ -159,6 +162,8 @@ inline std::ostream& operator<<(std::ostream& stream, const XlaNode& node) {
   return stream;
 }
 
+const xla::Shape& GetXlaShape(const torch::lazy::Value& value);
+
 template <typename T>
 T* NodeCast(const torch::lazy::Node* node, torch::lazy::OpKind op) {
   if (op != node->op()) {

torch_xla/csrc/ir.h

@@ -244,7 +244,7 @@ std::string DumpUtil::PostOrderToText(
   return ss.str();
 }
 
-std::string DumpUtil::ToHlo(absl::Span<const XlaValue> values,
+std::string DumpUtil::ToHlo(absl::Span<const torch::lazy::Value> values,
                             const torch::lazy::BackendDevice& device) {
   LoweringContext lowering_ctx("IrToHlo", device);
   for (auto& ir_value : values) {

torch_xla/csrc/ir_dump_util.cpp

@@ -22,7 +22,7 @@ class DumpUtil {
       absl::Span<const torch::lazy::Node* const> post_order,
       absl::Span<const torch::lazy::Node* const> roots);
 
-  static std::string ToHlo(absl::Span<const XlaValue> values,
+  static std::string ToHlo(absl::Span<const torch::lazy::Value> values,
                            const torch::lazy::BackendDevice& device);
 };
 

torch_xla/csrc/ir_dump_util.h

@@ -58,15 +58,15 @@ std::vector<const torch::lazy::Node*> Util::ComputePostOrder(
   return ComputePostOrder(nodes, &emap);
 }
 
-std::vector<XlaValue> Util::Clone(
-    absl::Span<const XlaValue> values,
+std::vector<torch::lazy::Value> Util::Clone(
+    absl::Span<const torch::lazy::Value> values,
     absl::Span<const torch::lazy::Node* const> post_order) {
   std::unordered_map<const torch::lazy::Node*, torch::lazy::NodePtr> clone_map;
   for (auto node : post_order) {
     if (clone_map.count(node) > 0) {
       continue;
     }
-    std::vector<XlaValue> inputs;
+    std::vector<torch::lazy::Value> inputs;
     for (auto& output : node->operands()) {
       auto it = clone_map.find(output.node);
       XLA_CHECK(it != clone_map.end())
@@ -77,7 +77,7 @@ std::vector<XlaValue> Util::Clone(
     clone_map[node] = casted->Clone(inputs);
   }
 
-  std::vector<XlaValue> cloned;
+  std::vector<torch::lazy::Value> cloned;
   for (auto& value : values) {
     auto it = clone_map.find(value.node.get());
     XLA_CHECK(it != clone_map.end()) << "Bad post-order: " << value->ToString();
@@ -86,7 +86,8 @@ std::vector<XlaValue> Util::Clone(
   return cloned;
 }
 
-std::vector<XlaValue> Util::Clone(absl::Span<const XlaValue> values) {
+std::vector<torch::lazy::Value> Util::Clone(
+    absl::Span<const torch::lazy::Value> values) {
   std::vector<const torch::lazy::Node*> nodes;
   for (auto& value : values) {
     nodes.push_back(value.node.get());

torch_xla/csrc/ir_util.cpp

@@ -31,11 +31,12 @@ class Util {
       absl::Span<const torch::lazy::Node* const> nodes);
 
   // Clones the IR graph whose roots are passed in the values parameter.
-  static std::vector<XlaValue> Clone(absl::Span<const XlaValue> values);
+  static std::vector<torch::lazy::Value> Clone(
+      absl::Span<const torch::lazy::Value> values);
 
   // Same as the above, but the post-order is passed as parameter.
-  static std::vector<XlaValue> Clone(
-      absl::Span<const XlaValue> values,
+  static std::vector<torch::lazy::Value> Clone(
+      absl::Span<const torch::lazy::Value> values,
       absl::Span<const torch::lazy::Node* const> post_order);
 
   // Retrieves the number of nodes within the graph whose sink are passed in the

torch_xla/csrc/ir_util.h

@@ -275,7 +275,7 @@ NmsResult BuildNms(xla::XlaOp boxes, xla::XlaOp scores,
   xla::XlaOp ones_included =
       xla::Select(included, xla::Broadcast(one_s32, {num_boxes}),
                   xla::Broadcast(zero_s32, {num_boxes}));
-  // num_valid is scalar. XlaValue should be bound by output_size.
+  // num_valid is scalar. torch::lazy::Value should be bound by output_size.
   xla::XlaOp num_valid_total = xla::Reduce(
       ones_included,
       /*init_value=*/zero_s32,