Optimized axis map

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

@@ -7,6 +7,7 @@
  */
 
 #include <executorch/backends/vulkan/runtime/api/containers/Tensor.h>
+#include <cassert>
 #include <cstring>
 
 namespace vkcompute {
@@ -99,12 +100,31 @@ std::vector<int64_t> calculate_strides(
  *
  * The axis mapping allows for permuted views of texture-backed tensors.
  */
-std::vector<int64_t> default_axis_map() {
-  // Currently, all compute shaders have an assumption that the channels dim is
-  // used to combine with the batch dim of a tensor. However, once dim mapping
-  // is integrated into the tensor indexing logic for each compute shader, we
-  // can be more flexible with mapping the batch dim to different texture axes
-  // in order to improve performance or memory footprint.
+std::vector<int64_t> calculate_axis_map(
+    const std::vector<int64_t>& sizes,
+    utils::AxisMapLayout axis_map_layout) {
+  if (axis_map_layout == utils::AxisMapLayout::OPTIMIZED) {
+    std::vector<int64_t> axis_map(sizes.size() + 1);
+    std::iota(axis_map.begin(), axis_map.end() - 1, 0);
+
+    std::stable_sort(
+        axis_map.begin(), axis_map.end() - 1, [&sizes](size_t i1, size_t i2) {
+          return sizes[i1] < sizes[i2];
+        });
+
+    assert(axis_map.size() > 0);
+    // Find the index of the channel dimension
+    for (size_t i = 0; i < axis_map.size() - 1; ++i) {
+      assert(sizes.size() > axis_map[i]);
+      if (sizes[axis_map[i]] == 2) {
+        axis_map.back() = i;
+        break;
+      }
+    }
+
+    return axis_map;
+  }
+  // default
   return {0, 1, 2, 2};
 }
 
@@ -439,13 +459,14 @@ vTensor::vTensor(
     const vkapi::ScalarType dtype,
     const utils::StorageType storage_type,
     const utils::GPUMemoryLayout memory_layout,
-    const bool allocate_memory)
+    const bool allocate_memory,
+    const utils::AxisMapLayout axis_map_layout)
     : dtype_(dtype),
       // Calculate tensor metadata
       sizes_(sizes.begin(), sizes.end()),
       packed_dim_(utils::to_packed_dim<int32_t>(memory_layout)),
       dim_order_(calculate_dim_order(sizes_.size(), packed_dim_)),
-      axis_map_(default_axis_map()),
+      axis_map_(calculate_axis_map(sizes_, axis_map_layout)),
       strides_(calculate_strides(sizes, dim_order_)),
       padded_sizes_{calculate_padded_sizes(sizes, packed_dim_)},
       unsqueezed_strides_{
@@ -484,13 +505,14 @@ vTensor::vTensor(
 vTensor::vTensor(
     Context* context,
     const vkapi::VulkanImage& image,
-    const utils::GPUMemoryLayout memory_layout)
+    const utils::GPUMemoryLayout memory_layout,
+    const utils::AxisMapLayout axis_map_layout)
     : dtype_(vkapi::element_scalartype(image.format())),
       // Calculate tensor metadata
       sizes_(calculate_sizes(image, memory_layout)),
       packed_dim_(utils::to_packed_dim<int32_t>(memory_layout)),
       dim_order_(),
-      axis_map_(default_axis_map()),
+      axis_map_(calculate_axis_map(sizes_, axis_map_layout)),
       strides_(),
       padded_sizes_(calculate_padded_sizes(sizes_, packed_dim_)),
       unsqueezed_strides_(),
@@ -547,7 +569,7 @@ vTensor::vTensor(
       sizes_(sizes.begin(), sizes.end()),
       packed_dim_(other.packed_dim_),
       dim_order_(dim_order.begin(), dim_order.end()),
-      axis_map_(default_axis_map()),
+      axis_map_(calculate_axis_map(sizes_, utils::kDefaultAxisMap)),
       strides_(calculate_strides(sizes_, dim_order_)),
       padded_sizes_{calculate_padded_sizes(sizes, packed_dim_)},
       unsqueezed_strides_{

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

@@ -183,14 +183,16 @@ class vTensor final {
       const vkapi::ScalarType dtype,
       const utils::StorageType storage_type = utils::kTexture3D,
       const utils::GPUMemoryLayout memory_layout = utils::kChannelsPacked,
-      const bool allocate_memory = true);
+      const bool allocate_memory = true,
+      const utils::AxisMapLayout axis_map_layout = utils::kDefaultAxisMap);
 
   vTensor(const vTensor& other) = delete;
 
   explicit vTensor(
       Context* context,
       const vkapi::VulkanImage& image,
-      const utils::GPUMemoryLayout memory_layout = utils::kChannelsPacked);
+      const utils::GPUMemoryLayout memory_layout = utils::kChannelsPacked,
+      const utils::AxisMapLayout axis_map_layout = utils::kDefaultAxisMap);
 
   /*
    * This constructor allows for the creation of a vTensor that references the

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -239,13 +239,20 @@ ValueRef ComputeGraph::add_tensor(
     const vkapi::ScalarType dtype,
     const utils::StorageType storage_type,
     const utils::GPUMemoryLayout memory_layout,
-    const int64_t shared_object_idx) {
+    const int64_t shared_object_idx,
+    const utils::AxisMapLayout axis_map_layout) {
   bool allocate_memory = shared_object_idx < 0;
 
   ValueRef idx(static_cast<int>(values_.size()));
   check_no_active_value_ptrs();
   values_.emplace_back(api::vTensor(
-      context(), sizes, dtype, storage_type, memory_layout, allocate_memory));
+      context(),
+      sizes,
+      dtype,
+      storage_type,
+      memory_layout,
+      allocate_memory,
+      axis_map_layout));
 
   if (!allocate_memory) {
     get_shared_object(shared_object_idx).add_user(this, idx);
@@ -257,44 +264,70 @@ ValueRef ComputeGraph::add_tensor(
     const std::vector<int64_t>& sizes,
     const vkapi::ScalarType dtype,
     const utils::StorageType storage_type,
-    const int64_t shared_object_idx) {
+    const int64_t shared_object_idx,
+    const utils::AxisMapLayout axis_map_layout) {
   return add_tensor(
       sizes,
       dtype,
       storage_type,
       suggested_memory_layout(sizes),
-      shared_object_idx);
+      shared_object_idx,
+      axis_map_layout);
 }
 
 ValueRef ComputeGraph::add_tensor(
     const std::vector<int64_t>& sizes,
     const vkapi::ScalarType dtype,
     const utils::GPUMemoryLayout memory_layout,
-    const int64_t shared_object_idx) {
+    const int64_t shared_object_idx,
+    const utils::AxisMapLayout axis_map_layout) {
   return add_tensor(
-      sizes, dtype, suggested_storage_type(), memory_layout, shared_object_idx);
+      sizes,
+      dtype,
+      suggested_storage_type(),
+      memory_layout,
+      shared_object_idx,
+      axis_map_layout);
 }
 
 ValueRef ComputeGraph::add_tensor_like(
     const ValueRef idx,
     const utils::StorageType storage_type,
-    const utils::GPUMemoryLayout memory_layout) {
-  return add_tensor(sizes_of(idx), dtype_of(idx), storage_type, memory_layout);
+    const utils::GPUMemoryLayout memory_layout,
+    const utils::AxisMapLayout axis_map_layout) {
+  return add_tensor(
+      sizes_of(idx),
+      dtype_of(idx),
+      storage_type,
+      memory_layout,
+      -1,
+      axis_map_layout);
 }
 
 ValueRef ComputeGraph::add_tensor_like(
     const ValueRef idx,
-    const utils::GPUMemoryLayout memory_layout) {
+    const utils::GPUMemoryLayout memory_layout,
+    const utils::AxisMapLayout axis_map_layout) {
   return add_tensor(
-      sizes_of(idx), dtype_of(idx), storage_type_of(idx), memory_layout);
+      sizes_of(idx),
+      dtype_of(idx),
+      storage_type_of(idx),
+      memory_layout,
+      -1,
+      axis_map_layout);
 }
 
 ValueRef ComputeGraph::add_tensor(
     const std::vector<int64_t>& sizes,
     const vkapi::ScalarType dtype,
-    const int64_t shared_object_idx) {
+    const int64_t shared_object_idx,
+    const utils::AxisMapLayout axis_map_layout) {
   return add_tensor(
-      sizes, dtype, suggested_memory_layout(sizes), shared_object_idx);
+      sizes,
+      dtype,
+      suggested_memory_layout(sizes),
+      shared_object_idx,
+      axis_map_layout);
 }
 
 ValueRef ComputeGraph::add_tensor(const vkapi::VulkanImage& image) {

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -461,7 +461,8 @@ class ComputeGraph final {
       const vkapi::ScalarType dtype,
       const utils::StorageType storage_type,
       const utils::GPUMemoryLayout memory_layout,
-      const int64_t shared_object_idx = -1);
+      const int64_t shared_object_idx = -1,
+      const utils::AxisMapLayout axis_map_layout = utils::kDefaultAxisMap);
 
   /*
    * Add a `api::vTensor` value to the graph with the specified properties. The
@@ -471,7 +472,8 @@ class ComputeGraph final {
       const std::vector<int64_t>& sizes,
       const vkapi::ScalarType dtype,
       const utils::StorageType storage_type,
-      const int64_t shared_object_idx = -1);
+      const int64_t shared_object_idx = -1,
+      const utils::AxisMapLayout axis_map_layout = utils::kDefaultAxisMap);
 
   /*
    * Add a `api::vTensor` value to the graph with the specified properties. The
@@ -481,7 +483,8 @@ class ComputeGraph final {
       const std::vector<int64_t>& sizes,
       const vkapi::ScalarType dtype,
       const utils::GPUMemoryLayout memory_layout,
-      const int64_t shared_object_idx = -1);
+      const int64_t shared_object_idx = -1,
+      const utils::AxisMapLayout axis_map_layout = utils::kDefaultAxisMap);
 
   /*
    * Add a `api::vTensor` value to the graph with the specified properties. The
@@ -491,7 +494,8 @@ class ComputeGraph final {
   ValueRef add_tensor(
       const std::vector<int64_t>& sizes,
       const vkapi::ScalarType dtype,
-      const int64_t shared_object_idx = -1);
+      const int64_t shared_object_idx = -1,
+      const utils::AxisMapLayout axis_map_layout = utils::kDefaultAxisMap);
 
   /*
    * Add a `api::vTensor` value to the graph with the specified image.
@@ -504,15 +508,17 @@ class ComputeGraph final {
   ValueRef add_tensor_like(
       const ValueRef vref,
       const utils::StorageType storage_type,
-      const utils::GPUMemoryLayout memory_layout);
+      const utils::GPUMemoryLayout memory_layout,
+      const utils::AxisMapLayout axis_map_layout = utils::kDefaultAxisMap);
 
   /*
    * Add a `api::vTensor` value to the graph with the properties of `vref`. The
    * suggested storage type will be used to construct the `api::vTensor`.
    */
   ValueRef add_tensor_like(
       const ValueRef vref,
-      const utils::GPUMemoryLayout memory_layout);
+      const utils::GPUMemoryLayout memory_layout,
+      const utils::AxisMapLayout axis_map_layout = utils::kDefaultAxisMap);
 
   /*
    * Use the copy constructor of `api::vTensor` to create a "view" of the

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

@@ -146,12 +146,14 @@ ValueRef prepack_standard(
     const ValueRef tensor_data,
     const utils::StorageType storage_type,
     const utils::GPUMemoryLayout layout,
-    const bool passthrough) {
+    const bool passthrough,
+    const utils::AxisMapLayout axis_map_layout) {
   if (passthrough && graph.val_is_tensor(tensor_data)) {
     return tensor_data;
   }
   VK_CHECK_COND(graph.val_is_tref(tensor_data));
-  ValueRef tensor = graph.add_tensor_like(tensor_data, storage_type, layout);
+  ValueRef tensor =
+      graph.add_tensor_like(tensor_data, storage_type, layout, axis_map_layout);
   add_prepack_standard_node(graph, tensor_data, tensor);
   return tensor;
 }

backends/vulkan/runtime/graph/ops/impl/Staging.h

@@ -48,7 +48,8 @@ ValueRef prepack_standard(
     const ValueRef tensor_data,
     const utils::StorageType storage_type,
     const utils::GPUMemoryLayout layout,
-    const bool passthrough = false);
+    const bool passthrough = false,
+    const utils::AxisMapLayout axis_map_layout = utils::kDefaultAxisMap);
 
 /*
  * Equivalent to `prepack_standard()` function, except the `storage_type` and

backends/vulkan/runtime/utils/StorageUtils.h

@@ -146,5 +146,14 @@ inline std::ostream& operator<<(
   return os;
 }
 
+enum class AxisMapLayout : uint8_t {
+  DEFAULT = 0u,
+  OPTIMIZED = 1u,
+};
+
+constexpr AxisMapLayout kDefaultAxisMap = AxisMapLayout::DEFAULT;
+
+constexpr AxisMapLayout kOptimizedAxisMap = AxisMapLayout::OPTIMIZED;
+
 } // namespace utils
 } // namespace vkcompute