Support Tensor.is_alias_of (#3966)

Summary:
Tensor.is_alias_of relies on Storage to perform. However, XLATensorImpl was
not implemented with that in mind. This commit adds a fake storage to XLATensor
as a marker to mark XLATensor that point to the same storage. The reason
why it's not done at XLATensorImpl is that XLATensor maintains the view ops/alias
logic in XLATensor class instead of relying on XLATensorImpl to do the check.

Test Plan:
./test/cpp/build/test_ptxla --gtest_filter=AtenXlaTensorTest.TestViewIsAliasOf

Author: alanwaketan

test/cpp/test_aten_xla_tensor.cpp

@@ -11212,5 +11212,25 @@ TEST_F(AtenXlaTensorTest, TestRoll) {
   ExpectCounterChanged("xla::roll", cpp_test::GetIgnoredCounters());
 }
 
+TEST_F(AtenXlaTensorTest, TestViewIsAliasOf) {
+  torch::Tensor a = torch::empty(4, torch::TensorOptions(torch::kFloat));
+  torch::Tensor b = torch::empty(4, torch::TensorOptions(torch::kFloat));
+
+  ForEachDevice([&](const torch::Device& device) {
+    torch::Tensor xla_a = CopyToDevice(a, device);
+    torch::Tensor xla_b = CopyToDevice(b, device);
+    EXPECT_EQ(!a.is_alias_of(b), !xla_a.is_alias_of(xla_b));
+
+    torch::Tensor c = a.view({2, 2});
+    torch::Tensor xla_c = xla_a.view({2, 2});
+    EXPECT_EQ(a.is_alias_of(c), xla_a.is_alias_of(xla_c));
+
+    torch::Tensor d = c.view({1, 4});
+    torch::Tensor lazy_d = xla_c.view({1, 4});
+    EXPECT_EQ(d.is_alias_of(c), lazy_d.is_alias_of(xla_c));
+    EXPECT_EQ(d.is_alias_of(a), lazy_d.is_alias_of(xla_a));
+  });
+}
+
 }  // namespace cpp_test
 }  // namespace torch_xla

torch_xla/csrc/tensor.cpp

@@ -492,28 +492,32 @@ XLATensorPtr XLATensor::Create(
 
 XLATensor::XLATensor(const at::Tensor& tensor,
                      const torch::lazy::BackendDevice& device)
-    : data_(std::make_shared<Data>(tensor, device)) {}
+    : XLATensor(std::make_shared<Data>(tensor, device)) {}
 
 XLATensor::XLATensor(torch::lazy::BackendDataPtr xla_data,
                      c10::optional<at::ScalarType> logical_element_type)
-    : data_(std::make_shared<Data>(xla_data, xla_data->device(),
-                                   logical_element_type)) {}
+    : XLATensor(std::make_shared<Data>(xla_data, xla_data->device(),
+                                       logical_element_type)) {}
 
 XLATensor::XLATensor(torch::lazy::Value ir_value,
                      const torch::lazy::BackendDevice& device,
                      c10::optional<at::ScalarType> logical_element_type)
-    : data_(std::make_shared<Data>(std::move(ir_value), device,
-                                   logical_element_type)) {
+    : XLATensor(std::make_shared<Data>(std::move(ir_value), device,
+                                       logical_element_type)) {
   TryLimitGraphSize();
 }
 
 XLATensor::XLATensor(std::shared_ptr<View> view,
                      const torch::lazy::BackendDevice& device,
                      c10::optional<at::ScalarType> logical_element_type)
-    : data_(std::make_shared<Data>(std::move(view), device,
-                                   logical_element_type)) {}
+    : XLATensor(std::make_shared<Data>(std::move(view), device,
+                                       logical_element_type)) {}
 
-XLATensor::XLATensor(std::shared_ptr<Data> data) : data_(std::move(data)) {}
+XLATensor::XLATensor(std::shared_ptr<Data> data)
+    : data_(std::move(data)),
+      storage_(c10::Storage(
+          {}, 0,
+          c10::DataPtr(nullptr, backendDeviceToAtenDevice(data_->device)))) {}
 
 XLATensor::Data* XLATensor::data() const {
   XLA_CHECK(data_ != nullptr) << "Trying to access a null cursor";
@@ -905,8 +909,10 @@ std::shared_ptr<View> XLATensor::CreateView(ViewInfo view_info) const {
 }
 
 XLATensorPtr XLATensor::CreateViewTensor(ViewInfo view_info) const {
-  return Create(CreateView(std::move(view_info)), GetDevice(),
-                dtype_optional());
+  auto new_tensor =
+      Create(CreateView(std::move(view_info)), GetDevice(), dtype_optional());
+  new_tensor->storage_ = Storage();
+  return new_tensor;
 }
 
 at::Tensor XLATensor::ToTensor(bool detached) {

torch_xla/csrc/tensor.h

@@ -1195,6 +1195,8 @@ class XLATensor : public c10::intrusive_ptr_target {
                        bool manual);
   void ClearShardingSpec();
 
+  const c10::Storage& Storage() const { return storage_; }
+
  private:
   struct SyncTensorsConfig {
     // Whether we want to force XLA data on the target tensors (hence trimming
@@ -1465,6 +1467,12 @@ class XLATensor : public c10::intrusive_ptr_target {
   bool ShouldSyncIrNode();
 
   std::shared_ptr<Data> data_;
+  // Temporarily used to suport Tensor.is_alias_of().
+  // This is a fake storage that doesn't store anything.
+  // Instead it serves as a marker to mark LazyTensors that
+  // points to the same storage, and thus alias of each other.
+  // FIXME(alanwaketan): Remove this once we have functionalization (bdhirsh).
+  c10::Storage storage_;
 };
 
 }  // namespace torch_xla

torch_xla/csrc/tensor_impl.cpp

@@ -186,10 +186,8 @@ void XLATensorImpl::AtenInitialize() {
   // ATEN specific initialization calls placed below.
 }
 
-const at::Storage& XLATensorImpl::storage() const {
-  XLA_ERROR() << "XLA tensors do not have storage";
-}
+const at::Storage& XLATensorImpl::storage() const { return tensor_->Storage(); }
 
-bool XLATensorImpl::has_storage() const { return false; }
+bool XLATensorImpl::has_storage() const { return tensor_->Storage(); }
 
 }  // namespace torch_xla