LTC phase 2 clean up (#3521)

* LTC Phase2 cleanup

* Remove Use as it is not really used anywhere

Author: JackCaoG

test/cpp/test_ir.cpp

@@ -17,24 +17,6 @@ TEST(IrTest, TestScalarCreate) {
   ASSERT_TRUE(scalar != nullptr);
 }
 
-TEST(IrTest, TestReplace) {
-  torch::lazy::NodePtr scalar1 = ir::ops::ScalarOp(1.0, xla::F32);
-  torch::lazy::NodePtr scalar2 = ir::ops::ScalarOp(2.0, xla::F32);
-  torch::lazy::NodePtr add = ir::Value(scalar1, 0) + ir::Value(scalar2, 0);
-
-  EXPECT_EQ(dynamic_cast<ir::Node*>(scalar1.get())->uses().size(), 1);
-  EXPECT_EQ(dynamic_cast<ir::Node*>(scalar2.get())->uses().size(), 1);
-
-  torch::lazy::NodePtr scalar3 = ir::ops::ScalarOp(3.0, xla::F32);
-  dynamic_cast<ir::Node*>(scalar1.get())->ReplaceAllUsesWith(scalar3);
-
-  EXPECT_EQ(dynamic_cast<ir::Node*>(scalar1.get())->uses().size(), 0);
-  EXPECT_EQ(dynamic_cast<ir::Node*>(scalar3.get())->uses().size(), 1);
-
-  dynamic_cast<ir::Node*>(add.get())->ReplaceOperand(0, scalar1);
-  EXPECT_EQ(dynamic_cast<ir::Node*>(scalar1.get())->uses().size(), 1);
-}
-
 TEST(IrTest, TestHash) {
   torch::lazy::NodePtr scalar1 = ir::ops::ScalarOp(1.0, xla::F32);
   torch::lazy::NodePtr scalar2 = ir::ops::ScalarOp(2.0, xla::F32);

torch_xla/csrc/XLANativeFunctions_include.h

@@ -0,0 +1,4 @@
+#include "tensorflow/compiler/xla/xla_client/debug_macros.h"
+#include "tensorflow/compiler/xla/xla_client/metrics.h"
+#include "torch_xla/csrc/aten_cpu_fallback.h"
+#include "torch_xla/csrc/aten_xla_bridge.h"

torch_xla/csrc/aten_xla_bridge.cpp

@@ -105,6 +105,14 @@ XLATensor GetOrCreateXlaTensor(const c10::optional<at::Tensor>& tensor,
   return xtensor ? *xtensor : XLATensor::Create(*tensor, device);
 }
 
+XLATensor GetXlaTensorOrCreateForWrappedNumber(const at::Tensor& tensor,
+                                               const Device& device) {
+  return (tensor.unsafeGetTensorImpl()->is_wrapped_number() ||
+          (tensor.dim() == 0 && tensor.numel() == 1))
+             ? GetOrCreateXlaTensor(tensor, device)
+             : GetXlaTensor(tensor);
+}
+
 std::vector<XLATensor> GetOrCreateXlaTensors(
     absl::Span<const at::Tensor> tensors, const Device& device) {
   std::vector<XLATensor> xla_tensors;

torch_xla/csrc/aten_xla_bridge.h

@@ -34,6 +34,9 @@ XLATensor GetOrCreateXlaTensor(const at::Tensor& tensor, const Device& device);
 
 XLATensor GetOrCreateXlaTensor(const c10::optional<at::Tensor>& tensor,
                                const Device& device);
+// TODO: change to upstream BackendDevice
+XLATensor GetXlaTensorOrCreateForWrappedNumber(const at::Tensor& tensor,
+                                               const Device& device);
 
 std::vector<XLATensor> GetOrCreateXlaTensors(
     absl::Span<const at::Tensor> tensors, const Device& device);

torch_xla/csrc/aten_xla_type.cpp

@@ -1,4 +1,3 @@
-#include <ATen/Context.h>
 #include <ATen/ExpandUtils.h>
 #include <ATen/Operators.h>
 #include <ATen/native/BinaryOps.h>
@@ -25,7 +24,6 @@
 #include "torch_xla/csrc/tensor_impl.h"
 #include "torch_xla/csrc/tensor_util.h"
 #include "torch_xla/csrc/torch_util.h"
-#include "torch_xla/csrc/version.h"
 
 // [Implementation Guidelines]
 // - If you want to call a at::func which doesn't have a kernel registered

torch_xla/csrc/ir.cpp

@@ -85,23 +85,6 @@ torch::lazy::hash_t GetOperandHashes(const OpList& operands,
 
 }  // namespace
 
-bool Use::operator<(const Use& rhs) const {
-  if (node->op() != rhs.node->op()) {
-    return node->op() < rhs.node->op();
-  }
-  if (operand_index != rhs.operand_index) {
-    return operand_index < rhs.operand_index;
-  }
-  return index < rhs.index;
-}
-
-std::string Use::ToString() const {
-  std::stringstream ss;
-  ss << node->ToString() << ", operand_index=" << operand_index
-     << ", index=" << index;
-  return ss.str();
-}
-
 const xla::Shape& Value::xla_shape() const {
   Node* casted = dynamic_cast<Node*>(node.get());
   return casted->xla_shape(index);
@@ -139,12 +122,7 @@ Node::Node(torch::lazy::OpKind op, xla::Shape shape, size_t num_outputs,
       node_hash_(GetOpHash(op, xla_shape_, hash_seed)),
       dag_hash_(node_hash_) {}
 
-Node::~Node() {
-  for (size_t i = 0; i < operands_as_outputs_.size(); ++i) {
-    Node* casted = dynamic_cast<Node*>(operands_[i].get());
-    casted->RemoveUse(Use(this, i, operands_as_outputs_[i].index));
-  }
-}
+Node::~Node() {}
 
 const xla::Shape& Node::xla_shape(size_t output_index) const {
   if (xla_shape_.IsTuple()) {
@@ -159,31 +137,16 @@ void Node::AddOperand(torch::lazy::NodePtr node, size_t index) {
   operands_.push_back(std::move(node));
   operands_as_outputs_.push_back(
       torch::lazy::Output(operands_.back().get(), index));
-  Node* casted = dynamic_cast<Node*>(operands_.back().get());
-  casted->AddUse(Use(this, operands_.size() - 1, index));
 }
 
 void Node::ReplaceOperand(size_t operand_no, torch::lazy::NodePtr node,
                           size_t index) {
   XLA_CHECK_LT(index, node->num_outputs());
-  Node* casted = dynamic_cast<Node*>(node.get());
   torch::lazy::Output* output = &operands_as_outputs_.at(operand_no);
-  Node* casted_to_remove = dynamic_cast<Node*>(operands_[operand_no].get());
-  casted_to_remove->RemoveUse(Use(this, operand_no, output->index));
-  casted->AddUse(Use(this, operand_no, index));
   *output = torch::lazy::Output(node.get(), index);
   operands_[operand_no] = std::move(node);
 }
 
-void Node::ReplaceAllUsesWith(torch::lazy::NodePtr node, size_t index) {
-  // A call to ReplaceOperand() will end up calling RemoveUse() into the
-  // current node, so snapshot the current uses and iterate over them.
-  std::vector<Use> current_uses(uses_.begin(), uses_.end());
-  for (auto& use : current_uses) {
-    use.node->ReplaceOperand(use.operand_index, node, index);
-  }
-}
-
 XlaOpVector Node::ReturnOp(xla::XlaOp op, LoweringContext* loctx) const {
   XLA_CHECK_EQ(num_outputs(), 1);
   loctx->AssignOutputOp(torch::lazy::Output(this), op);

torch_xla/csrc/ir.h

@@ -28,31 +28,6 @@ class LoweringContext;
 
 using XlaOpVector = tensorflow::gtl::InlinedVector<xla::XlaOp, 1>;
 
-// Represents a use of the output of a given node.
-// If use U is within node N, it means that node U.node is using the output
-// U.index of the node N.
-struct Use {
-  Use() = default;
-  Use(Node* node, size_t operand_index, size_t index)
-      : node(node), operand_index(operand_index), index(index) {}
-
-  bool operator<(const Use& rhs) const;
-
-  std::string ToString() const;
-
-  // The node using the output of the node this use belongs to.
-  Node* node = nullptr;
-  // The operand index, within node's operands, which this use refers to.
-  size_t operand_index = 0;
-  // The index within output the user node refers to.
-  size_t index = 0;
-};
-
-inline std::ostream& operator<<(std::ostream& stream, const Use& use) {
-  stream << use.ToString();
-  return stream;
-}
-
 template <typename T>
 using OutputMap =
     std::unordered_map<torch::lazy::Output, T, torch::lazy::Output::Hasher>;
@@ -109,13 +84,9 @@ class Node : public torch::lazy::Node {
   // multi-output node, output_index must be zero.
   const xla::Shape& xla_shape(size_t output_index) const;
 
-  const std::set<Use>& uses() const { return uses_; }
-
   void ReplaceOperand(size_t operand_no, torch::lazy::NodePtr node,
                       size_t index = 0);
 
-  void ReplaceAllUsesWith(torch::lazy::NodePtr node, size_t index = 0);
-
   virtual torch::lazy::NodePtr Clone(OpList operands) const;
 
   virtual XlaOpVector Lower(LoweringContext* loctx) const;
@@ -135,10 +106,6 @@ class Node : public torch::lazy::Node {
   // Adds node's index output number as operand.
   void AddOperand(torch::lazy::NodePtr node, size_t index = 0);
 
-  void AddUse(Use use) { uses_.insert(std::move(use)); }
-
-  void RemoveUse(const Use& use) { uses_.erase(use); }
-
   xla::Shape GetOpShape(const std::function<xla::Shape()>& shape_fn) const;
 
   static torch::lazy::hash_t GetOpHash(torch::lazy::OpKind op,
@@ -148,8 +115,6 @@ class Node : public torch::lazy::Node {
   static std::vector<torch::lazy::SourceLocation> GetFrameInfo();
 
   xla::Shape xla_shape_;
-  // We use a set for uses, as we want deterministic use sequencing.
-  std::set<Use> uses_;
   torch::lazy::hash_t node_hash_;
   torch::lazy::hash_t dag_hash_;
 };