vTensor cleanup 4/N - consolidate texture positions and extents to be logical positions and extents (#5442)

Summary:
Pull Request resolved: #5442

## Context

Tensors are N-dimensional constructs, whereas image textures are limited to 3 dimensions (with an X, Y, and Z axis). Before axis mapping was introduced, converting between tensor index and texture coordinate, as well as converting between tensor sizes and texture limits, was easy because of rigid way texture axes were used to represent tensor dimensions:

1. the X axis of the tensor corresponds to the width dim
2. the Y axis of the tensor corresponds to the height dim
3. the Z axis of the tensor corresponds to the combined batch + channels dim

However, with the introduction of axis mapping, the texture representation of a tensor is far more flexible.

NOTE: summary is still a WIP. Short on time today so will fill in more context later

This diff consolidates all image extent related member variables of `vTensor` to use logical texture limits instead of physical texture limits. The reason for this is to reduce developer confusion about different representations of texture metadata, and also to enforce better practices in shader programming. Using logical texture positions and logical texture limits is usually preferable because of the following:

1. Each element has a clear relationship to a tensor dimension, thus it is easier to understand semantically
2. It's more efficient to translate between logical texture position and tensor indices; in shaders, we often have to convert texture positions to tensor indices to coordinate data loads/write among the different input tensors and output tensor

## Notes for Reviewers

There are a large number of changes in this file, but majority of meaningful changes are in `Tensor.h` and `Tensor.cpp`. The rest is just replacing instance of `image_extents()` and `texture_limits()` with `logical_limits()`.
ghstack-source-id: 243309911
exported-using-ghexport

Reviewed By: jorgep31415

Differential Revision: D62901893

fbshipit-source-id: bfb63151c85255967ed9e26afc82eaa70b1b365b

Author: SS-JIA

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

@@ -418,14 +418,12 @@ vTensor::vTensor(
       padded_sizes_{calculate_padded_sizes(sizes, memory_layout_)},
       unsqueezed_strides_{unsqueeze_strides(strides_, numel_)},
       padded_numel_(utils::multiply_integers(padded_sizes_)),
-      texture_limits_{{0, 0, 0}},
       logical_limits_{{0, 0, 0}},
       // Utility Uniform Buffers that can be passed to shaders as arguments
       sizes_uniform_(),
       strides_uniform_(),
       numel_uniform_(),
       axis_map_uniform_(),
-      texture_limits_uniform_(),
       logical_limits_uniform_(),
       // Construct Tensor storage
       storage_(
@@ -440,12 +438,7 @@ vTensor::vTensor(
       dim_order_is_valid(dim_order_), "computed dim order is invalid");
 
   if (storage_type != utils::kBuffer) {
-    texture_limits_.limits = utils::ivec3{
-        utils::safe_downcast<int32_t>(storage_.image_extents_[0]),
-        utils::safe_downcast<int32_t>(storage_.image_extents_[1]),
-        utils::safe_downcast<int32_t>(storage_.image_extents_[2])};
-
-    update_logical_limits();
+    set_logical_limits(storage_.image_extents_);
   }
 
   if (dtype == vkapi::kHalf) {
@@ -470,14 +463,12 @@ vTensor::vTensor(const vTensor& other)
           other.unsqueezed_strides_.begin(),
           other.unsqueezed_strides_.end()},
       padded_numel_(other.padded_numel_),
-      texture_limits_{other.texture_limits_},
       logical_limits_{other.logical_limits_},
       // Empty initialize Utility Uniform Buffers
       sizes_uniform_(),
       strides_uniform_(),
       numel_uniform_(),
       axis_map_uniform_(),
-      texture_limits_uniform_(),
       logical_limits_uniform_(),
       // Copy Tensor storage
       storage_(other.storage_) {}
@@ -498,14 +489,12 @@ vTensor::vTensor(
       padded_sizes_{calculate_padded_sizes(sizes, memory_layout_)},
       unsqueezed_strides_{unsqueeze_strides(strides_, numel_)},
       padded_numel_(utils::multiply_integers(padded_sizes_)),
-      texture_limits_{other.texture_limits_},
       logical_limits_(other.logical_limits_),
       // Empty initialize Utility Uniform Buffers
       sizes_uniform_(),
       strides_uniform_(),
       numel_uniform_(),
       axis_map_uniform_(),
-      texture_limits_uniform_(),
       logical_limits_uniform_(),
       // Copy Tensor storage
       storage_(other.storage_, vkapi::element_size(dtype_) * offset_numel) {
@@ -547,18 +536,10 @@ vkapi::VulkanBuffer& vTensor::buffer(
   return storage_.buffer_;
 }
 
-void vTensor::update_logical_limits() {
-  logical_limits_.limits[0] = texture_limits_.limits[axis_map_.at(0)];
-  logical_limits_.limits[1] = texture_limits_.limits[axis_map_.at(1)];
-  logical_limits_.limits[2] = texture_limits_.limits[axis_map_.at(2)];
-}
-
-utils::uvec3 vTensor::logical_extents() const {
-  utils::uvec3 logical_extents(
-      {utils::safe_downcast<uint32_t>(logical_limits_.limits[0]),
-       utils::safe_downcast<uint32_t>(logical_limits_.limits[1]),
-       utils::safe_downcast<uint32_t>(logical_limits_.limits[2])});
-  return logical_extents;
+void vTensor::set_logical_limits(const utils::uvec3& image_extents) {
+  logical_limits_.limits[0] = image_extents[axis_map_.at(0)];
+  logical_limits_.limits[1] = image_extents[axis_map_.at(1)];
+  logical_limits_.limits[2] = image_extents[axis_map_.at(2)];
 }
 
 const vkapi::BufferBindInfo vTensor::sizes_ubo() {
@@ -585,13 +566,6 @@ const vkapi::BufferBindInfo vTensor::axis_map_ubo() {
   return vkapi::BufferBindInfo(axis_map_uniform_.buffer());
 }
 
-const vkapi::BufferBindInfo vTensor::texture_limits_ubo() {
-  if (!texture_limits_uniform_.buffer()) {
-    texture_limits_uniform_ = ParamsBuffer(storage_.context_, texture_limits_);
-  }
-  return vkapi::BufferBindInfo(texture_limits_uniform_.buffer());
-}
-
 const vkapi::BufferBindInfo vTensor::logical_limits_ubo() {
   if (!logical_limits_uniform_.buffer()) {
     logical_limits_uniform_ = ParamsBuffer(storage_.context_, logical_limits_);
@@ -655,18 +629,10 @@ void vTensor::update_metadata() {
   unsqueezed_strides_ = unsqueeze_strides(strides_, numel_);
   padded_numel_ = utils::multiply_integers(padded_sizes_);
 
-  // Calculate the extents of the image texture that would have been required
-  // for a tensor of the new sizes.
-  utils::uvec3 virtual_extents =
-      calculate_image_extents(padded_sizes_, axis_map_, memory_layout_);
-
-  // Update the texture limits to reflect the new virtual extents.
-  texture_limits_.limits = utils::ivec3{
-      utils::safe_downcast<int32_t>(virtual_extents[0]),
-      utils::safe_downcast<int32_t>(virtual_extents[1]),
-      utils::safe_downcast<int32_t>(virtual_extents[2])};
-
-  update_logical_limits();
+  // Calculate the image extents that would have been used to allocate a texture
+  // withthe current sizes, and use that to set the logical limits.
+  set_logical_limits(
+      calculate_image_extents(padded_sizes_, axis_map_, memory_layout_));
 
   if (sizes_uniform_.buffer()) {
     sizes_uniform_.update(utils::make_whcn_ivec4(sizes_));
@@ -680,9 +646,6 @@ void vTensor::update_metadata() {
   if (axis_map_uniform_.buffer()) {
     axis_map_uniform_.update(utils::make_ivec4(axis_map_));
   }
-  if (texture_limits_uniform_.buffer()) {
-    texture_limits_uniform_.update(texture_limits_);
-  }
   if (logical_limits_uniform_.buffer()) {
     logical_limits_uniform_.update(logical_limits_);
   }
@@ -695,9 +658,11 @@ void vTensor::check_sizes(const std::vector<int64_t>& sizes) const {
     utils::uvec3 virtual_extents =
         calculate_image_extents(padded_sizes_, axis_map_, memory_layout_);
 
-    bool valid_resize = virtual_extents[0] <= image_extents()[0];
-    valid_resize = valid_resize && virtual_extents[1] <= image_extents()[1];
-    valid_resize = valid_resize && virtual_extents[2] <= image_extents()[2];
+    bool valid_resize = virtual_extents[0] <= storage_.image_extents_[0];
+    valid_resize =
+        valid_resize && virtual_extents[1] <= storage_.image_extents_[1];
+    valid_resize =
+        valid_resize && virtual_extents[2] <= storage_.image_extents_[2];
 
     VK_CHECK_COND(
         valid_resize,

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

@@ -276,9 +276,7 @@ class vTensor final {
   // Contains the number of elements in the tensor according to the padded
   // sizes.
   size_t padded_numel_;
-  // See the comments documenting image_extents() for more context.
-  TextureLimits texture_limits_;
-  // See the comments documenting logical_extents() for more context.
+  // See the comments documenting logical_limits() for more context.
   TextureLimits logical_limits_;
 
   /*
@@ -294,7 +292,6 @@ class vTensor final {
   ParamsBuffer strides_uniform_;
   ParamsBuffer numel_uniform_;
   ParamsBuffer axis_map_uniform_;
-  ParamsBuffer texture_limits_uniform_;
   ParamsBuffer logical_limits_uniform_;
 
   vTensorStorage storage_;
@@ -342,28 +339,30 @@ class vTensor final {
     return storage_.storage_type_ == utils::kBuffer;
   }
 
-  /*
-   * Returns the raw image extents of the underlying image texture used to store
-   * the tensor's data. Note that due to axis mapping, the X, Y, and Z extents
-   * may not correspond to the width, height, or channels dimension of the
-   * tensor.
-   */
-  inline const utils::uvec3& image_extents() const {
-    return storage_.image_extents_;
-  }
-
  private:
-  void update_logical_limits();
+  void set_logical_limits(const utils::uvec3& image_extents);
 
  public:
   /*
-   * Returns the image extents of the underlying image texture, but re-ordered
-   * such that the first element is the extent of the axis used to represent the
-   * tensor's width dimension, the second element is the extent of the axis used
-   * to represent the tensor's height dimension, and the third element is the
-   * extent of the axis used to represent the tensor's channels dimension.
+   * The logical limits of the tensor are derived from the image extents of the
+   * image texture used to store the tensor, but with two key differences.
+   *
+   * First, the image extents are permuted according to the axis map. This
+   * makes it so that the first element of the logical limit is the limit of the
+   * texture axis corresponding to the width dimension of the tensor, the next
+   * element is the limit of the texture axis corresponding to the height
+   * dimension and the last element is the limit of the texture axis that
+   * corresponds to the channels dimension of the tensor.
+   *
+   * Second, the logical limits may use smaller extents than the actual image
+   * extents of the image texture. This is due to dynamic shape; if the tensor's
+   * `virtual_resize()` function is called, then the logical limits will reflect
+   * the extents that would be needed to support a tensor with the updated sizes
+   * instead of the original sizes.
    */
-  utils::uvec3 logical_extents() const;
+  inline const utils::ivec3& logical_limits() const {
+    return logical_limits_.limits;
+  }
 
   /*
    * Extract an `vkapi::ScalarType` from the TensorOptions member
@@ -430,18 +429,8 @@ class vTensor final {
   const vkapi::BufferBindInfo axis_map_ubo();
 
   /*
-   * Returns a GPU buffer containing the virtual image extents of the tensor.
-   * Since a tensor can be resized with the virtual_resize() function, this
-   * GPU buffer contains the image extents of the tensor calculated using the
-   * virtual_resize() function. This allows shaders to exit early if they are
-   * working outside the limits of the texture.
-   */
-  const vkapi::BufferBindInfo texture_limits_ubo();
-
-  /*
-   * Returns a GPU buffer containing the logical image extents of the tensor.
-   * It contains the same data as texture_limits_ubo(), but with the data
-   * re-ordered. See the comments for logical_extents() for more context.
+   * Returns a GPU buffer containing the logical limits of the tensor. See the
+   * comments for logical_limits() for more context.
    */
   const vkapi::BufferBindInfo logical_limits_ubo();
 
@@ -450,10 +439,6 @@ class vTensor final {
    */
   const vkapi::BufferBindInfo numel_ubo();
 
-  inline const utils::ivec3 texture_limits() const {
-    return texture_limits_.limits;
-  }
-
   inline size_t numel() const {
     return numel_;
   }

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -428,7 +428,7 @@ utils::uvec3 ComputeGraph::create_global_wg_size(const ValueRef idx) {
   if (is_buffer_storage(idx)) {
     return {uint32_t(numel_of(idx)), 1u, 1u};
   }
-  return image_extents_of(idx);
+  return logical_limits_of(idx);
 }
 
 utils::uvec3 ComputeGraph::create_local_wg_size(

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -284,12 +284,8 @@ class ComputeGraph final {
 
   vkapi::ScalarType dtype_of(const ValueRef idx) const;
 
-  inline utils::uvec3 image_extents_of(const ValueRef idx) const {
-    return values_.at(idx).toConstTensor().image_extents();
-  }
-
-  inline utils::uvec3 logical_extents_of(const ValueRef idx) const {
-    return values_.at(idx).toConstTensor().logical_extents();
+  inline const utils::ivec3& logical_limits_of(const ValueRef idx) const {
+    return values_.at(idx).toConstTensor().logical_limits();
   }
 
   inline int32_t numel_of(const ValueRef idx) const {
@@ -335,10 +331,6 @@ class ComputeGraph final {
     return values_.at(idx).toTensor().axis_map_ubo();
   }
 
-  inline vkapi::BufferBindInfo texture_limits_ubo(const ValueRef idx) {
-    return values_.at(idx).toTensor().texture_limits_ubo();
-  }
-
   inline vkapi::BufferBindInfo logical_limits_ubo(const ValueRef idx) {
     return values_.at(idx).toTensor().logical_limits_ubo();
   }

backends/vulkan/runtime/graph/ops/impl/BatchNorm.cpp

@@ -88,7 +88,7 @@ void add_native_batch_norm_node(
       {{out_ref, vkapi::MemoryAccessType::WRITE},
        {{in_ref, arg_weight, arg_bias, arg_mean, arg_var},
         vkapi::MemoryAccessType::READ}},
-      {t_out->texture_limits_ubo(),
+      {t_out->logical_limits_ubo(),
        graph.create_params_buffer(epsilon),
        graph.create_params_buffer(num_texel_per_batch)}));
 }

backends/vulkan/runtime/graph/ops/impl/Cat.cpp

@@ -40,7 +40,7 @@ void add_cat_default_node(
 
     for (ValueRef input_ref : *input_list) {
       vTensorPtr t_in = graph.get_tensor(input_ref);
-      utils::ivec3 range = t_in->texture_limits();
+      utils::ivec3 range = t_in->logical_limits();
       add_copy_offset_node(
           graph, input_ref, range, src_offset, dst_offset, out);
       dst_offset[0] += range[0];
@@ -52,7 +52,7 @@ void add_cat_default_node(
 
     for (ValueRef input_ref : *input_list) {
       vTensorPtr t_in = graph.get_tensor(input_ref);
-      utils::ivec3 range = t_in->texture_limits();
+      utils::ivec3 range = t_in->logical_limits();
       add_copy_offset_node(
           graph, input_ref, range, src_offset, dst_offset, out);
       dst_offset[1] += range[1];
@@ -63,7 +63,7 @@ void add_cat_default_node(
 
     for (ValueRef input_ref : *input_list) {
       vTensorPtr t_in = graph.get_tensor(input_ref);
-      utils::ivec3 range = t_in->texture_limits();
+      utils::ivec3 range = t_in->logical_limits();
       add_copy_offset_node(
           graph, input_ref, range, src_offset, dst_offset, out);
       dst_offset[2] += range[2];

backends/vulkan/runtime/graph/ops/impl/Clone.cpp

@@ -32,7 +32,7 @@ void add_clone_node(
       graph.create_local_wg_size(out),
       {{out, vkapi::MemoryAccessType::WRITE},
        {in, vkapi::MemoryAccessType::READ}},
-      {t_out->texture_limits_ubo()}));
+      {t_out->logical_limits_ubo()}));
 }
 
 void clone(ComputeGraph& graph, const std::vector<ValueRef>& args) {

backends/vulkan/runtime/graph/ops/impl/Convolution.cpp

@@ -291,7 +291,7 @@ utils::uvec3 create_conv2d_global_wg_size(
     const Conv2dMethod method,
     const ValueRef out) {
   if (method == Conv2dMethod::Pointwise) {
-    const utils::uvec3 image_extents = graph.image_extents_of(out);
+    const utils::uvec3 image_extents = graph.logical_limits_of(out);
     return {
         utils::div_up(image_extents[0u], 2u),
         utils::div_up(image_extents[1u], 2u),
@@ -376,7 +376,7 @@ void add_conv2d_node(
        {{arg_in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}},
       // Shader params buffers
       {
-          t_out->texture_limits_ubo(),
+          t_out->logical_limits_ubo(),
           t_in->sizes_ubo(),
           graph.create_params_buffer(kernel_params),
           graph.create_params_buffer(extra_params),
@@ -474,7 +474,7 @@ void add_conv1d_node(
        {{arg_in, arg_weight, arg_bias}, vkapi::MemoryAccessType::READ}},
       // Shader params buffers
       {
-          t_out->texture_limits_ubo(),
+          t_out->logical_limits_ubo(),
           t_in->sizes_ubo(),
           graph.create_params_buffer(kernel_params),
           graph.create_params_buffer(out_params),

backends/vulkan/runtime/graph/ops/impl/Linear.cpp

@@ -103,7 +103,7 @@ void add_addmm_naive_node(
   add_storage_type_suffix(kernel_name, graph.storage_type_of(out));
   add_dtype_suffix(kernel_name, graph.dtype_of(out));
 
-  utils::uvec3 global_wg_size = graph.logical_extents_of(out);
+  utils::uvec3 global_wg_size = graph.logical_limits_of(out);
   graph.execute_nodes().emplace_back(new ExecuteNode(
       graph,
       VK_KERNEL_FROM_STR(kernel_name),
@@ -182,7 +182,7 @@ void add_addmm_optimized_node(
 
   add_dtype_suffix(kernel_name, graph.dtype_of(out));
 
-  utils::uvec3 global_size;
+  utils::uvec3 global_size = graph.logical_limits_of(out);
 
   // Each thread computes a W=(2/4) x H=4 x C=(1/4) output tile. Therefore, the
   // total number of threads is W/(2 or 4) x H/4 x C/1. Since the out tensor is
@@ -193,9 +193,9 @@ void add_addmm_optimized_node(
   if (mat1_sizes.at(mat1_dims - 2) < 8) {
     // Use `logical_extents` instead of `image_extents` because the workgroup
     // axes need to correspond to tensor dimensions.
-    global_size = utils::divup_vec(graph.logical_extents_of(out), {4, 2, 1});
+    global_size = utils::divup_vec(global_size, {4, 2, 1});
   } else {
-    global_size = utils::divup_vec(graph.logical_extents_of(out), {4, 4, 1});
+    global_size = utils::divup_vec(global_size, {4, 4, 1});
   }
   utils::uvec3 local_size = adaptive_work_group_size(global_size);