[Vulkan] Let tensor wrap an external buffer with caller-supplied sizes (pytorch#22031)

Follow-up to pytorch#21930, which added external-image storage.
`vTensorStorage` had no buffer equivalent, so a tensor could only be
built over externally owned storage if that storage was a `VulkanImage`.

This adds an `external_buffer` parameter beside `external_image`, along
with a `vTensorStorage` constructor and a
`ComputeGraph::add_tensor` overload taking a `VulkanBuffer`. The buffer
is held as a `VulkanBuffer` copy, which aliases the handle with
`is_copy_` set, so the view never frees memory it does not own. At most
one of the two may be set.

  ## Files

- `backends/vulkan/runtime/api/containers/Tensor.h` — `external_buffer`
on the `vTensor` constructor, and the `vTensorStorage` buffer
constructor.
- `backends/vulkan/runtime/api/containers/Tensor.cpp` — the buffer
storage constructor, external-storage selection, and the size check.
- `backends/vulkan/runtime/graph/ComputeGraph.h`, `.cpp` — an
`add_tensor` overload taking sizes, dtype, memory layout and a buffer.
  - `backends/vulkan/test/vulkan_compute_api_test.cpp` — 3 new tests.

  ## Testing

`vulkan_compute_api_test` New coverage: wrapping a buffer at `[1, 37, 8,
64]` and keeping the sizes and buffer storage type; rejecting sizes
larger than the buffer; and a view reporting `is_copy_of` the source,
with the source buffer still live after the view is destroyed and
deferred cleanup has run.

cc @SS-JIA @manuelcandales @digantdesai @cbilgin

Author: kiymetakdemir

backends/vulkan/runtime/api/containers/Tensor.cpp

@@ -732,6 +732,23 @@ vTensorStorage::vTensorStorage(
           allocate_memory)),
       last_access_{} {}
 
+namespace {
+
+std::shared_ptr<vTensorStorage> acquire_external_storage(
+    Context* const context,
+    const vkapi::VulkanImage* image,
+    const vkapi::VulkanBuffer* buffer,
+    const vkapi::ScalarType dtype) {
+  VK_CHECK_COND(
+      image == nullptr || buffer == nullptr,
+      "a tensor can wrap an image or a buffer, not both");
+  return image != nullptr
+      ? std::make_shared<vTensorStorage>(context, *image)
+      : std::make_shared<vTensorStorage>(context, *buffer, dtype);
+}
+
+} // namespace
+
 vTensorStorage::vTensorStorage(
     Context* const context,
     const vkapi::VulkanImage& image)
@@ -747,6 +764,20 @@ vTensorStorage::vTensorStorage(
       buffer_(vkapi::VulkanBuffer()),
       last_access_{} {}
 
+vTensorStorage::vTensorStorage(
+    Context* const context,
+    const vkapi::VulkanBuffer& buffer,
+    const vkapi::ScalarType dtype)
+    : context_(context),
+      storage_type_{utils::kBuffer},
+      image_extents_({0, 0, 0}),
+      buffer_length_{
+          static_cast<int64_t>(buffer.mem_size() / vkapi::element_size(dtype))},
+      buffer_offset_{0},
+      image_(vkapi::VulkanImage()),
+      buffer_(buffer),
+      last_access_{} {}
+
 vTensorStorage::~vTensorStorage() {
   flush();
 }
@@ -852,7 +883,8 @@ vTensor::vTensor(
     const utils::GPUMemoryLayout memory_layout,
     const bool allocate_memory,
     const utils::AxisMapLayout axis_map_layout,
-    const vkapi::VulkanImage* external_image)
+    const vkapi::VulkanImage* external_image,
+    const vkapi::VulkanBuffer* external_buffer)
     : dtype_(get_effective_scalar_type(context, dtype, memory_layout)),
       packed_dim_info_(calculate_packed_dim_info(memory_layout, storage_type)),
       // Calculate tensor metadata
@@ -881,8 +913,12 @@ vTensor::vTensor(
       buffer_meta_(),
       // Construct Tensor storage
       storage_(
-          external_image != nullptr
-              ? std::make_shared<vTensorStorage>(context, *external_image)
+          (external_image != nullptr || external_buffer != nullptr)
+              ? acquire_external_storage(
+                    context,
+                    external_image,
+                    external_buffer,
+                    dtype_)
               : std::make_shared<vTensorStorage>(
                     context,
                     storage_type,
@@ -892,15 +928,18 @@ vTensor::vTensor(
                     dtype_,
                     physical_numel_,
                     allocate_memory)) {
-  // An external image was sized by its creator, so nothing guarantees it
-  // matches the metadata derived from `sizes`. The extents feed the
-  // shader-visible logical limits below, so they have to agree.
+  // External storage was sized by its creator, so nothing guarantees it fits
+  // the metadata derived from `sizes`. An image must match exactly, since its
+  // extents feed the shader-visible logical limits below; a buffer need only
+  // be large enough.
   if (external_image != nullptr) {
     VK_CHECK_COND(
         storage_->image_extents_ ==
             calculate_image_extents(
                 dtype_, packed_dim_info_, padded_sizes_, axis_map_),
         "external image extents do not match the requested tensor sizes");
+  } else if (external_buffer != nullptr) {
+    check_sizes(sizes_);
   }
   // uniform_data_ only valid for low dim tensors
   if (sizes.size() <= 4) {

backends/vulkan/runtime/api/containers/Tensor.h

@@ -140,6 +140,11 @@ class vTensorStorage final {
 
   vTensorStorage(Context* const context, const vkapi::VulkanImage& image);
 
+  vTensorStorage(
+      Context* const context,
+      const vkapi::VulkanBuffer& buffer,
+      const vkapi::ScalarType dtype);
+
  public:
   vTensorStorage(vTensorStorage& other) = delete;
   vTensorStorage& operator=(const vTensorStorage& other) = delete;
@@ -211,9 +216,11 @@ class vTensor final {
       const utils::GPUMemoryLayout memory_layout = utils::kChannelsPacked,
       const bool allocate_memory = true,
       const utils::AxisMapLayout axis_map_layout = utils::kDefaultAxisMap,
-      // When non-null, wrap this image instead of allocating. Stored as a
-      // VulkanImage copy, which aliases the handle without owning it.
-      const vkapi::VulkanImage* external_image = nullptr);
+      // When non-null, wrap this resource instead of allocating. Stored as a
+      // copy, which aliases the handle without owning it. At most one may be
+      // set.
+      const vkapi::VulkanImage* external_image = nullptr,
+      const vkapi::VulkanBuffer* external_buffer = nullptr);
 
   vTensor(const vTensor& other) = delete;
 

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -573,6 +573,26 @@ ValueRef ComputeGraph::add_tensor(
   return idx;
 }
 
+ValueRef ComputeGraph::add_tensor(
+    const std::vector<int64_t>& sizes,
+    const vkapi::ScalarType dtype,
+    const utils::GPUMemoryLayout memory_layout,
+    const vkapi::VulkanBuffer& buffer) {
+  ValueRef idx(static_cast<int>(values_.size()));
+  check_no_active_value_ptrs();
+  values_.emplace_back(api::vTensor(
+      context(),
+      sizes,
+      dtype,
+      utils::kBuffer,
+      memory_layout,
+      /*allocate_memory=*/false,
+      utils::kDefaultAxisMap,
+      /*external_image=*/nullptr,
+      &buffer));
+  return idx;
+}
+
 ValueRef ComputeGraph::add_tensor(const vkapi::VulkanImage& image) {
   ValueRef idx(static_cast<int>(values_.size()));
   check_no_active_value_ptrs();

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -790,6 +790,16 @@ class ComputeGraph final {
       const utils::GPUMemoryLayout memory_layout,
       const vkapi::VulkanImage& image);
 
+  /*
+   * Wrap an externally owned buffer as a tensor with the given logical sizes.
+   * The buffer must be large enough for them.
+   */
+  ValueRef add_tensor(
+      const std::vector<int64_t>& sizes,
+      const vkapi::ScalarType dtype,
+      const utils::GPUMemoryLayout memory_layout,
+      const vkapi::VulkanBuffer& buffer);
+
   /*
    * Add a `api::vTensor` value to the graph with the properties of `vref`.
    */

backends/vulkan/test/vulkan_compute_api_test.cpp

@@ -1781,6 +1781,79 @@ TEST_F(VulkanComputeAPITest, test_tensor_over_external_image_does_not_own_it) {
   EXPECT_FALSE(owner.image().is_copy());
 }
 
+static vTensor make_buffer_owner() {
+  return vTensor(
+      context(),
+      kExternalSizes,
+      vkapi::kFloat,
+      utils::kBuffer,
+      utils::kWidthPacked);
+}
+
+TEST_F(VulkanComputeAPITest, test_tensor_over_external_buffer_keeps_sizes) {
+  vTensor owner = make_buffer_owner();
+
+  vTensor view(
+      context(),
+      kExternalSizes,
+      vkapi::kFloat,
+      utils::kBuffer,
+      utils::kWidthPacked,
+      /*allocate_memory = */ false,
+      utils::kDefaultAxisMap,
+      /*external_image = */ nullptr,
+      &owner.buffer());
+
+  EXPECT_TRUE(view.sizes() == kExternalSizes);
+  EXPECT_TRUE(view.numel() == owner.numel());
+  EXPECT_TRUE(view.storage_type() == utils::kBuffer);
+}
+
+TEST_F(
+    VulkanComputeAPITest,
+    test_tensor_over_external_buffer_rejects_overflow) {
+  vTensor owner = make_buffer_owner();
+
+  std::vector<int64_t> too_large = kExternalSizes;
+  too_large.back() *= 2;
+
+  EXPECT_THROW(
+      vTensor(
+          context(),
+          too_large,
+          vkapi::kFloat,
+          utils::kBuffer,
+          utils::kWidthPacked,
+          /*allocate_memory = */ false,
+          utils::kDefaultAxisMap,
+          /*external_image = */ nullptr,
+          &owner.buffer()),
+      vkapi::Error);
+}
+
+TEST_F(VulkanComputeAPITest, test_tensor_over_external_buffer_does_not_own_it) {
+  vTensor owner = make_buffer_owner();
+
+  {
+    vTensor view(
+        context(),
+        kExternalSizes,
+        vkapi::kFloat,
+        utils::kBuffer,
+        utils::kWidthPacked,
+        /*allocate_memory = */ false,
+        utils::kDefaultAxisMap,
+        /*external_image = */ nullptr,
+        &owner.buffer());
+
+    EXPECT_TRUE(view.buffer().is_copy_of(owner.buffer()));
+  }
+
+  context()->flush();
+  EXPECT_TRUE(owner.buffer());
+  EXPECT_FALSE(owner.buffer().is_copy());
+}
+
 TEST(VulkanComputeGraphTest, test_values_scalars) {
   GraphConfig config;
   ComputeGraph graph(config);