Extend TensorImpl with BackendMeta (pytorch#97429)

BackendMeta offers a binary interface for the backend to attach arbitrary data to TensorImpl. TensorImpl has exactly one "slot" for backend metadata, however backend is free to compose any structure that is opaque to the framework beyond iheriting standard BackendMeta base.

Change-Id: I670fcdd16dd1c2b00f7eaa1cbc5b5dfea59a6221

Fixes #ISSUE_NUMBER

Pull Request resolved: pytorch#97429
Approved by: https://github.com/ezyang

Author: ppiskorski

c10/core/TensorImpl.h

@@ -19,6 +19,7 @@
 #include <c10/util/Logging.h>
 #include <c10/util/Optional.h>
 #include <c10/util/accumulate.h>
+#include <c10/util/intrusive_ptr.h>
 #include <c10/util/irange.h>
 #include <c10/util/python_stub.h>
 #include <c10/util/safe_numerics.h>
@@ -217,6 +218,18 @@ is_channels_last_3d
 is_non_overlapping_and_dense
 #endif
 
+/**
+ * This structure is intended to hold additional metadata of the specific device
+ *backend
+ **/
+struct C10_API BackendMeta : intrusive_ptr_target {
+  virtual ~BackendMeta(){};
+  virtual intrusive_ptr<BackendMeta> clone(
+      const intrusive_ptr<BackendMeta>& ptr) const {
+    return ptr;
+  }
+};
+
 struct C10_API ExtraMeta {
   SymDimVector sizes_ = {0};
   SymDimVector strides_ = {1};
@@ -229,6 +242,7 @@ struct C10_API ExtraMeta {
   SymBool is_channels_last_3d_{false};
   SymBool is_non_overlapping_and_dense_{true};
   std::unique_ptr<c10::NamedTensorMetaInterface> named_tensor_meta_ = nullptr;
+  intrusive_ptr<c10::BackendMeta> backend_meta_;
 
   ExtraMeta() = default;
 
@@ -243,7 +257,8 @@ struct C10_API ExtraMeta {
       SymBool is_channels_last,
       SymBool is_channels_last_3d,
       SymBool is_non_overlapping_and_dense,
-      std::unique_ptr<c10::NamedTensorMetaInterface> named_tensor_meta)
+      std::unique_ptr<c10::NamedTensorMetaInterface> named_tensor_meta,
+      intrusive_ptr<c10::BackendMeta> backend_meta)
       : sizes_(std::move(sizes)),
         strides_(std::move(strides)),
         numel_(std::move(numel)),
@@ -255,7 +270,8 @@ struct C10_API ExtraMeta {
         is_channels_last_(std::move(is_channels_last)),
         is_channels_last_3d_(std::move(is_channels_last_3d)),
         is_non_overlapping_and_dense_(std::move(is_non_overlapping_and_dense)),
-        named_tensor_meta_(std::move(named_tensor_meta)) {}
+        named_tensor_meta_(std::move(named_tensor_meta)),
+        backend_meta_(backend_meta) {}
 
   std::unique_ptr<ExtraMeta> clone() const {
     return std::make_unique<ExtraMeta>(
@@ -269,7 +285,8 @@ struct C10_API ExtraMeta {
         is_channels_last_,
         is_channels_last_3d_,
         is_non_overlapping_and_dense_,
-        named_tensor_meta_ ? named_tensor_meta_->clone() : nullptr);
+        named_tensor_meta_ ? named_tensor_meta_->clone() : nullptr,
+        backend_meta_ ? backend_meta_->clone(backend_meta_) : nullptr);
   }
 };
 
@@ -1576,6 +1593,27 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target {
     return data_type_.itemsize();
   }
 
+  void set_backend_meta(intrusive_ptr<c10::BackendMeta> backend_meta) {
+    if (!extra_meta_) {
+      extra_meta_ = std::make_unique<ExtraMeta>();
+    }
+    extra_meta_->backend_meta_ = std::move(backend_meta);
+  }
+
+  c10::BackendMeta* get_backend_meta() {
+    if (!extra_meta_) {
+      return nullptr;
+    }
+    return extra_meta_->backend_meta_.get();
+  }
+
+  intrusive_ptr<c10::BackendMeta> get_backend_meta_intrusive_ptr() const {
+    if (!extra_meta_) {
+      return nullptr;
+    }
+    return extra_meta_->backend_meta_;
+  }
+
  protected:
   /**
    * Returns the human-readable name of the actual type of this object (e.g.,

test/cpp/api/tensor.cpp

@@ -1217,3 +1217,44 @@ TEST(TensorTest, ReshapeAlias) {
   torch::_reshape_alias((z * z), {9}, {1}).mean().backward();
   ASSERT_TRUE(torch::equal(y.grad(), z.grad()));
 }
+
+TEST(TensorTest, BackendMetadata) {
+  // Tests ability to assign custom backend metadata to tensor.
+
+  struct CustomBackendMetadata : public c10::BackendMeta {
+    mutable bool cloned_{false}; // for testing this field will mutate when
+                                 // clone() is called by shallow_copy_from.
+    c10::intrusive_ptr<c10::BackendMeta> clone(
+        const c10::intrusive_ptr<c10::BackendMeta>& ptr) const override {
+      cloned_ = true;
+      return c10::BackendMeta::clone(ptr);
+    }
+  };
+
+  at::Tensor y;
+  c10::intrusive_ptr<c10::BackendMeta> tmeta{};
+  CustomBackendMetadata* custom_tmeta{nullptr};
+
+  {
+    auto x = torch::ones({3, 3});
+    auto impl{x.unsafeGetTensorImpl()};
+    ASSERT_TRUE(impl != nullptr);
+
+    tmeta = impl->get_backend_meta_intrusive_ptr();
+    ASSERT_TRUE(tmeta == nullptr);
+    c10::intrusive_ptr<c10::BackendMeta> new_tmeta{
+        std::unique_ptr<c10::BackendMeta>(new CustomBackendMetadata())};
+    impl->set_backend_meta(new_tmeta);
+    tmeta = impl->get_backend_meta_intrusive_ptr();
+    ASSERT_TRUE(tmeta == new_tmeta);
+    custom_tmeta = dynamic_cast<CustomBackendMetadata*>(tmeta.get());
+    ASSERT_TRUE(custom_tmeta != nullptr);
+    ASSERT_TRUE(custom_tmeta->cloned_ == false);
+    y.unsafeGetTensorImpl()->shallow_copy_from(x.getIntrusivePtr());
+  }
+
+  ASSERT_TRUE(
+      tmeta == y.unsafeGetTensorImpl()->get_backend_meta_intrusive_ptr());
+  ASSERT_TRUE(tmeta.get() == y.unsafeGetTensorImpl()->get_backend_meta());
+  ASSERT_TRUE(custom_tmeta->cloned_ == true);
+}