[LT] Support Tensor.is_alias_of

Jiewen Tan · pytorchmergebot · commit ab0d9b18e97a · 2022-04-14T07:28:03.000Z
Summary: Tensor.is_alias_of relies on Storage to perform. However, LTCTensorImpl was not implemented with that in mind. This commit adds a fake storage to LazyTensor as a marker to mark LazyTensors that point to the same storage. The reason why it's not done at LTCTensorImpl is that LazyTensor maintains the view ops/alias logic in LazyTensor class instead of relying on TensorImpl to do the check. Test Plan: ./build/bin/test_lazy --gtest_filter=LazyOpsTest.IsAliasOf Pull Request resolved: #75246 Approved by: https://github.com/bdhirsh
diff --git a/c10/core/Allocator.h b/c10/core/Allocator.h
@@ -217,9 +217,9 @@ struct AllocatorRegisterer {
   }
 };
 
-#define REGISTER_ALLOCATOR(t, f)                  \
-  namespace {                                     \
-  static AllocatorRegisterer<t> g_allocator_d(f); \
+#define REGISTER_ALLOCATOR(t, f)                       \
+  namespace {                                          \
+  static c10::AllocatorRegisterer<t> g_allocator_d(f); \
   }
 
 // An interface for reporting thread local memory usage
diff --git a/test/cpp/lazy/test_lazy_ops.cpp b/test/cpp/lazy/test_lazy_ops.cpp
@@ -10721,6 +10721,26 @@ TEST_F(LazyOpsTest, TestLerpScalarOut) {
   ExpectCounterChanged("lazy::lerp", GetIgnoredCounters());
 }
 
+TEST_F(LazyOpsTest, IsAliasOf) {
+  auto a = torch::empty(4, torch::TensorOptions(torch::kFloat).device(DefaultDevice()));
+  auto b = torch::empty(4, torch::TensorOptions(torch::kFloat).device(DefaultDevice()));
+
+  ForEachDevice([&](const torch::Device& device) {
+    auto lazy_a = CopyToDevice(a, device);
+    auto lazy_b = CopyToDevice(b, device);
+    EXPECT_EQ(!a.is_alias_of(b), !lazy_a.is_alias_of(lazy_b));
+
+    auto c = a.view({2, 2});
+    auto lazy_c = lazy_a.view({2, 2});
+    EXPECT_EQ(a.is_alias_of(c), lazy_a.is_alias_of(lazy_c));
+
+    auto d = c.view({1, 4});
+    auto lazy_d = lazy_c.view({1, 4});
+    EXPECT_EQ(d.is_alias_of(c), lazy_d.is_alias_of(lazy_c));
+    EXPECT_EQ(d.is_alias_of(a), lazy_d.is_alias_of(lazy_a));
+  });
+}
+
 #endif // FBCODE_CAFFE2
 
 }  // namespace lazy
diff --git a/torch/csrc/lazy/core/tensor.cpp b/torch/csrc/lazy/core/tensor.cpp
@@ -1,7 +1,7 @@
-#include <c10/util/irange.h>
+#include <torch/csrc/lazy/core/config.h>
 #include <torch/csrc/lazy/core/tensor.h>
 
-#include <torch/csrc/lazy/core/config.h>
+#include <c10/util/irange.h>
 #include <torch/csrc/lazy/core/helpers.h>
 #include <torch/csrc/lazy/core/ir_dump_util.h>
 #include <torch/csrc/lazy/core/lazy_graph_executor.h>
@@ -64,22 +64,24 @@ LazyTensorPtr LazyTensor::Create(std::shared_ptr<Data> data) {
 }
 
 LazyTensor::LazyTensor(const at::Tensor& tensor, const BackendDevice& device)
-    : data_(std::make_shared<Data>(tensor, device)) {}
+    : LazyTensor(std::make_shared<Data>(tensor, device)) {}
 
 LazyTensor::LazyTensor(BackendDataPtr handle)
-    : data_(std::make_shared<Data>(handle, handle->device())) {}
+    : LazyTensor(std::make_shared<Data>(handle, handle->device())) {}
 
 LazyTensor::LazyTensor(Value ir_value, const BackendDevice& device)
-    : data_(std::make_shared<Data>(std::move(ir_value), device)) {
+    : LazyTensor(std::make_shared<Data>(std::move(ir_value), device)) {
   TryLimitGraphSize();
 }
 
 LazyTensor::LazyTensor(
     std::shared_ptr<LazyView> view,
     const BackendDevice& device)
-    : data_(std::make_shared<Data>(std::move(view), device)) {}
+    : LazyTensor(std::make_shared<Data>(std::move(view), device)) {}
 
-LazyTensor::LazyTensor(std::shared_ptr<Data> data) : data_(std::move(data)) {}
+LazyTensor::LazyTensor(std::shared_ptr<Data> data)
+  : data_(std::move(data))
+  , storage_(c10::Storage({}, 0, c10::DataPtr(nullptr, backendDeviceToAtenDevice(data_->device)))) {}
 
 LazyTensor::Data* LazyTensor::data() const {
   TORCH_CHECK(data_ != nullptr, "Trying to access a null cursor");
@@ -346,7 +348,9 @@ std::shared_ptr<LazyView> LazyTensor::CreateView(ViewInfo view_info) const {
 }
 
 LazyTensorPtr LazyTensor::CreateViewTensor(ViewInfo view_info) const {
-  return Create(CreateView(std::move(view_info)), GetDevice());
+  auto new_tensor = Create(CreateView(std::move(view_info)), GetDevice());
+  new_tensor->storage_ = Storage();
+  return new_tensor;
 }
 
 at::Tensor LazyTensor::ToTensor(bool detached) {
diff --git a/torch/csrc/lazy/core/tensor.h b/torch/csrc/lazy/core/tensor.h
@@ -133,6 +133,11 @@ class TORCH_API LazyTensor : public c10::intrusive_ptr_target {
   // Applies the queue of operations in preparation for using the data.
   void ApplyPendingGraph();
 
+  const c10::Storage& Storage() const { return storage_; }
+  // This is currently only used by outlier view ops such as expand that
+  // don't go through CreateViewTensor to support Tensor.is_alias_of.
+  void SetStorage(const c10::Storage& storage) { storage_ = storage; }
+
  private:
   LazyTensor(const at::Tensor& tensor, const BackendDevice& device);
   LazyTensor(Value ir_value, const BackendDevice& device);
@@ -177,6 +182,12 @@ class TORCH_API LazyTensor : public c10::intrusive_ptr_target {
   static int64_t GetNextTensorId();
 
   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_;
 };
 
 // Utils to convert at::Tensor to LazyTensor, and vice versa.
diff --git a/torch/csrc/lazy/core/tensor_impl.cpp b/torch/csrc/lazy/core/tensor_impl.cpp
@@ -1,5 +1,6 @@
 #include <torch/csrc/lazy/core/tensor_impl.h>
 
+#include <c10/core/Allocator.h>
 #include <c10/core/ScalarType.h>
 #include <c10/core/impl/DeviceGuardImplInterface.h>
 #include <c10/macros/Macros.h>
@@ -191,10 +192,6 @@ bool LTCTensorImpl::is_contiguous(c10::MemoryFormat _unused) const {
   return true;
 }
 
-const at::Storage& LTCTensorImpl::storage() const {
-  TORCH_CHECK(false, "Lazy tensors do not have storage");
-}
-
 #endif  // C10_DISABLE_TENSORIMPL_EXTENSIBILITY
 
 }  // namespace lazy
diff --git a/torch/csrc/lazy/core/tensor_impl.h b/torch/csrc/lazy/core/tensor_impl.h
@@ -1,7 +1,6 @@
 #pragma once
 
 #include <ATen/Tensor.h>
-#include <c10/core/Storage.h>
 #include <c10/core/TensorImpl.h>
 
 #include <torch/csrc/lazy/core/tensor.h>
@@ -44,8 +43,8 @@ class TORCH_API LTCTensorImpl final : public c10::TensorImpl {
   int64_t numel() const override;
 
   bool is_contiguous(at::MemoryFormat memory_format) const override;
-  const at::Storage& storage() const override;
-  bool has_storage() const override { return false; }
+  const at::Storage& storage() const override { return tensor_->Storage(); }
+  bool has_storage() const override { return tensor_->Storage(); }
 #endif  // C10_DISABLE_TENSORIMPL_EXTENSIBILITY
 
  private:
diff --git a/torch/csrc/lazy/ts_backend/tensor_aten_ops.cpp b/torch/csrc/lazy/ts_backend/tensor_aten_ops.cpp
@@ -116,10 +116,12 @@ void as_strided_(torch::lazy::LazyTensorPtr& input, std::vector<int64_t> size,
 
 torch::lazy::LazyTensorPtr expand(const torch::lazy::LazyTensorPtr& input, std::vector<int64_t> size) {
   auto input_shape = input->shape();
-  return torch::lazy::LazyTensor::Create(torch::lazy::MakeNode<torch::lazy::Expand>(
+  auto output = torch::lazy::LazyTensor::Create(torch::lazy::MakeNode<torch::lazy::Expand>(
       input->GetIrValue(),
       GetExpandDimensions(input_shape.Get(), std::move(size)),
       /*is_scalar_expand=*/false), input->GetDevice());
+  output->SetStorage(input->Storage());
+  return output;
 }
 
 void fill_(torch::lazy::LazyTensorPtr& input, const at::Scalar& value) {

Original file line number	Diff line number	Diff line change
`@@ -217,9 +217,9 @@ struct AllocatorRegisterer {`
`217`	`217`	`}`
`218`	`218`	`};`
`219`	`219`
`220`		`-#define REGISTER_ALLOCATOR(t, f) \`
`221`		`- namespace { \`
`222`		`- static AllocatorRegisterer<t> g_allocator_d(f); \`
	`220`	`+#define REGISTER_ALLOCATOR(t, f) \`
	`221`	`+ namespace { \`
	`222`	`+ static c10::AllocatorRegisterer<t> g_allocator_d(f); \`
`223`	`223`	`}`
`224`	`224`
`225`	`225`	`// An interface for reporting thread local memory usage`