Introduce vTensorPtr to prevent reference invalidation and remove get_val() API (#2978)

Summary:
Pull Request resolved: #2978

## Context

Currently when writing operators developers will save a reference to a `vTensor` retrieved from a `ComputeGraph`'s list of `values_` like so:

```
vTensor& vten = graph.get_val(vref).toTensor();
```

However, this is dangerous since if any values are added once the reference has been stored, `values_` which is a `std::vector` may have been resized and therefore have its contents moved, meaning the reference is now invalid.

To protect against this, this changeset introduces the `vTensorPtr` class which is a wrapper around a `vTensor*`.  When constructed, it will increment a counter in the `ComputeGraph` instance, and when destroyed it will decrement the counter. `ComputeGraph` cannot add any values while the counter is not zero.

Since `Value` can be converted to other non-trivial types, this changeset also removes the `get_val` function entirely to guard against unsafe behaviour.
ghstack-source-id: 222224052
exported-using-ghexport

Reviewed By: jorgep31415

Differential Revision: D55984187

fbshipit-source-id: 22c619f651b5b3783c7626263694ca46b9f84723

Author: SS-JIA

backends/vulkan/runtime/VulkanBackend.cpp

@@ -334,18 +334,18 @@ bool maybe_resize_input(
     const size_t input_i,
     exec_aten::Tensor& et_tensor) {
   ValueRef in_tensor_ref = graph->inputs()[input_i].value;
-  vTensor& in_tensor = graph->get_val(in_tensor_ref).toTensor();
+  vTensorPtr in_tensor = graph->get_tensor(in_tensor_ref);
 
   ET_CHECK_MSG(
-      et_tensor.dim() == in_tensor.sizes().size(),
+      et_tensor.dim() == in_tensor->sizes().size(),
       "Cannot resize input tensor: old ndim %zu does not match new ndim %zu",
-      static_cast<size_t>(in_tensor.sizes().size()),
+      static_cast<size_t>(in_tensor->sizes().size()),
       static_cast<size_t>(et_tensor.dim()));
 
   bool should_resize = false;
   std::vector<int64_t> new_sizes(et_tensor.dim());
   for (size_t i = 0; i < et_tensor.dim(); i++) {
-    if (in_tensor.sizes()[i] != et_tensor.sizes()[i]) {
+    if (in_tensor->sizes()[i] != et_tensor.sizes()[i]) {
       should_resize = true;
     }
     new_sizes.at(i) = et_tensor.sizes()[i];
@@ -356,9 +356,9 @@ bool maybe_resize_input(
   }
 
   ET_CHECK_MSG(
-      in_tensor.numel() == et_tensor.numel(),
+      in_tensor->numel() == et_tensor.numel(),
       "Vulkan tensor numel %zu does not match ET tensor numel %zu",
-      static_cast<size_t>(in_tensor.numel()),
+      static_cast<size_t>(in_tensor->numel()),
       static_cast<size_t>(et_tensor.numel()));
 
   return should_resize;
@@ -369,12 +369,12 @@ void maybe_resize_output(
     const size_t output_i,
     exec_aten::Tensor& et_tensor) {
   ValueRef out_tensor_ref = graph->outputs()[output_i].value;
-  vTensor& out_tensor = graph->get_val(out_tensor_ref).toTensor();
+  vTensorPtr out_tensor = graph->get_tensor(out_tensor_ref);
 
   exec_aten::SizesType new_output_size[kTensorDimensionLimit];
-  size_t ndim = out_tensor.sizes().size();
+  size_t ndim = out_tensor->sizes().size();
   for (int i = 0; i < ndim; ++i) {
-    new_output_size[i] = out_tensor.sizes()[i];
+    new_output_size[i] = out_tensor->sizes()[i];
   }
 
   exec_aten::ArrayRef<exec_aten::SizesType> output_size{new_output_size, ndim};

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -17,6 +17,39 @@
 
 namespace vkcompute {
 
+//
+// VTensorPtr
+//
+
+#define VALUE_PTR_CLASS_IMPL(classname, ctype, type_name)                 \
+  classname::classname(ComputeGraph* const graph, const ValueRef idx)     \
+      : graph_(graph), ptr_(&(graph_->values_.at(idx).to##type_name())) { \
+    graph_->values_in_use_++;                                             \
+  }                                                                       \
+  ctype* classname::operator->() const {                                  \
+    return ptr_;                                                          \
+  }                                                                       \
+  ctype& classname::operator*() const {                                   \
+    return *ptr_;                                                         \
+  }                                                                       \
+  classname::~classname() {                                               \
+    graph_->values_in_use_--;                                             \
+  }
+
+VALUE_PTR_CLASS_IMPL(vTensorPtr, vTensor, Tensor)
+VALUE_PTR_CLASS_IMPL(TensorRefPtr, TensorRef, TensorRef)
+VALUE_PTR_CLASS_IMPL(StagingPtr, api::StorageBuffer, Staging)
+VALUE_PTR_CLASS_IMPL(IntListPtr, std::vector<int64_t>, IntList)
+VALUE_PTR_CLASS_IMPL(DoubleListPtr, std::vector<double>, DoubleList)
+VALUE_PTR_CLASS_IMPL(BoolListPtr, std::vector<bool>, BoolList)
+VALUE_PTR_CLASS_IMPL(ValueListPtr, std::vector<ValueRef>, ValueList)
+
+#undef VALUE_PTR_CLASS_IMPL
+
+//
+// ComputeGraph
+//
+
 ComputeGraph::ComputeGraph(GraphConfig config)
     : config_{config},
       prepack_descriptor_counts_{},
@@ -105,6 +138,35 @@ api::GPUMemoryLayout ComputeGraph::suggested_memory_layout(
   return api::kChannelsPacked;
 }
 
+void ComputeGraph::check_no_active_value_ptrs() {
+  VK_CHECK_COND(
+      values_in_use_ == 0,
+      "Make sure that there are no pointers stored from the return values of "
+      "`ComputeGraph::get_*()` functions in scope before adding Values to the "
+      "graph. Modifying the graph's values may cause existing pointers to be "
+      "invalidated.");
+}
+
+std::vector<int64_t> ComputeGraph::get_sizes_of(ValueRef idx) {
+  Value& val = values_.at(idx);
+  if (val.isTensor()) {
+    return val.toTensor().sizes();
+  } else if (val.isTensorRef()) {
+    return val.toTensorRef().sizes;
+  }
+  VK_THROW("Could not get sizes of value with type ", val.type());
+}
+
+api::ScalarType ComputeGraph::get_dtype_of(ValueRef idx) {
+  Value& val = values_.at(idx);
+  if (val.isTensor()) {
+    return val.toTensor().dtype();
+  } else if (val.isTensorRef()) {
+    return val.toTensorRef().dtype;
+  }
+  VK_THROW("Could not get dtype of value with type ", val.type());
+}
+
 ValueRef ComputeGraph::add_tensor(
     const std::vector<int64_t>& sizes,
     const api::ScalarType dtype,
@@ -114,6 +176,7 @@ ValueRef ComputeGraph::add_tensor(
   bool allocate_memory = shared_object_idx < 0;
 
   ValueRef idx(static_cast<int>(values_.size()));
+  check_no_active_value_ptrs();
   values_.emplace_back(vTensor(
       context(), sizes, dtype, storage_type, memory_layout, allocate_memory));
 
@@ -133,18 +196,17 @@ ValueRef ComputeGraph::add_tensor(
 }
 
 ValueRef ComputeGraph::add_tensor_like(
-    const ValueRef vref,
+    const ValueRef idx,
     const api::StorageType storage_type,
     const api::GPUMemoryLayout memory_layout) {
-  TensorRef& tref = get_val(vref).toTensorRef();
-  return add_tensor(tref.sizes, tref.dtype, storage_type, memory_layout);
+  return add_tensor(
+      get_sizes_of(idx), get_dtype_of(idx), storage_type, memory_layout);
 }
 
 ValueRef ComputeGraph::add_tensor_like(
-    const ValueRef vref,
+    const ValueRef idx,
     const api::GPUMemoryLayout memory_layout) {
-  TensorRef& tref = get_val(vref).toTensorRef();
-  return add_tensor(tref.sizes, tref.dtype, memory_layout);
+  return add_tensor(get_sizes_of(idx), get_dtype_of(idx), memory_layout);
 }
 
 ValueRef ComputeGraph::add_tensor(
@@ -160,6 +222,7 @@ ValueRef ComputeGraph::add_tensorref(
     const api::ScalarType dtype,
     const void* const data) {
   ValueRef idx(static_cast<int>(values_.size()));
+  check_no_active_value_ptrs();
   values_.emplace_back(TensorRef(sizes, dtype, data));
   return idx;
 }
@@ -168,24 +231,28 @@ ValueRef ComputeGraph::add_staging(
     const api::ScalarType dtype,
     const size_t numel) {
   ValueRef idx(static_cast<int>(values_.size()));
+  check_no_active_value_ptrs();
   values_.emplace_back(api::StorageBuffer(context(), dtype, numel));
   return idx;
 }
 
 ValueRef ComputeGraph::add_none() {
   ValueRef idx(static_cast<int>(values_.size()));
+  check_no_active_value_ptrs();
   values_.emplace_back();
   return idx;
 }
 
 ValueRef ComputeGraph::add_value_list(std::vector<ValueRef>&& value) {
   ValueRef idx(static_cast<int>(values_.size()));
+  check_no_active_value_ptrs();
   values_.emplace_back(std::move(value));
   return idx;
 }
 
 ValueRef ComputeGraph::add_string(std::string&& str) {
   ValueRef idx(static_cast<int>(values_.size()));
+  check_no_active_value_ptrs();
   values_.emplace_back(std::move(str));
   return idx;
 }
@@ -194,8 +261,9 @@ ValueRef ComputeGraph::set_input_tensor(
     const ValueRef idx,
     const bool use_staging) {
   if (use_staging) {
-    vTensor& tensor = get_val(idx).toTensor();
-    ValueRef staging_idx = add_staging(tensor.dtype(), tensor.gpu_numel());
+    api::ScalarType dtype = get_tensor(idx)->dtype();
+    size_t gpu_numel = get_tensor(idx)->gpu_numel();
+    ValueRef staging_idx = add_staging(dtype, gpu_numel);
     add_staging_to_tensor_node(*this, staging_idx, idx);
     inputs_.push_back({idx, staging_idx});
     return staging_idx;
@@ -208,8 +276,9 @@ ValueRef ComputeGraph::set_output_tensor(
     const ValueRef idx,
     const bool use_staging) {
   if (use_staging) {
-    vTensor& tensor = get_val(idx).toTensor();
-    ValueRef staging_idx = add_staging(tensor.dtype(), tensor.gpu_numel());
+    api::ScalarType dtype = get_tensor(idx)->dtype();
+    size_t gpu_numel = get_tensor(idx)->gpu_numel();
+    ValueRef staging_idx = add_staging(dtype, gpu_numel);
     add_tensor_to_staging_node(*this, idx, staging_idx);
     outputs_.push_back({idx, staging_idx});
     return staging_idx;
@@ -229,20 +298,18 @@ void ComputeGraph::copy_into_staging(
     const ValueRef idx,
     const void* data,
     const size_t numel) {
-  Value& in_val = get_val(idx);
-  api::StorageBuffer& staging = in_val.toStaging();
-  size_t nbytes = numel * api::element_size(staging.dtype());
-  copy_ptr_to_staging(data, staging, nbytes);
+  StagingPtr staging = get_staging(idx);
+  size_t nbytes = numel * api::element_size(staging->dtype());
+  copy_ptr_to_staging(data, *staging, nbytes);
 }
 
 void ComputeGraph::copy_from_staging(
     const ValueRef idx,
     void* data,
     const size_t numel) {
-  Value& out_val = get_val(idx);
-  api::StorageBuffer& staging = out_val.toStaging();
-  size_t nbytes = numel * api::element_size(staging.dtype());
-  copy_staging_to_ptr(staging, data, nbytes);
+  StagingPtr staging = get_staging(idx);
+  size_t nbytes = numel * api::element_size(staging->dtype());
+  copy_staging_to_ptr(*staging, data, nbytes);
 }
 
 void ComputeGraph::prepare() {
@@ -308,7 +375,7 @@ void ComputeGraph::resize_input(
     const int64_t idx,
     const std::vector<int64_t>& new_sizes) {
   IOValueRef io_val = inputs_.at(idx);
-  get_val(io_val.value).toTensor().virtual_resize(new_sizes);
+  get_tensor(io_val.value)->virtual_resize(new_sizes);
 }
 
 void ComputeGraph::propagate_resize() {