Update on "[executorch] Custom op for fast hadamard transform kernel"

Custom op support for Fast Hadamard Transform.

Differential Revision: [D60530438](https://our.internmc.facebook.com/intern/diff/D60530438/)

[ghstack-poisoned]

Author: swolchok

.github/workflows/android-perf.yml

@@ -220,6 +220,8 @@ jobs:
         device: ${{ fromJson(needs.set-parameters.outputs.devices) }}
       fail-fast: false
     with:
+      # Due to scheduling a job may be pushed beyond the default 60m threshold
+      timeout: 120
       device-type: android
       runner: linux.2xlarge
       test-infra-ref: ''

.github/workflows/apple-perf.yml

@@ -292,6 +292,8 @@ jobs:
         device: ${{ fromJson(needs.set-parameters.outputs.devices) }}
       fail-fast: false
     with:
+      # Due to scheduling a job may be pushed beyond the default 60m threshold
+      timeout: 120
       device-type: ios
       # For iOS testing, the runner just needs to call AWS Device Farm, so there is no need to run this on macOS
       runner: linux.2xlarge

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

@@ -27,6 +27,8 @@ class StagingBuffer final {
   size_t nbytes_;
   vkapi::VulkanBuffer vulkan_buffer_;
 
+  void* mapped_data_;
+
  public:
   StagingBuffer(
       Context* context_p,
@@ -37,7 +39,8 @@ class StagingBuffer final {
         numel_(numel),
         nbytes_(element_size(dtype_) * numel_),
         vulkan_buffer_(
-            context_p_->adapter_ptr()->vma().create_staging_buffer(nbytes_)) {}
+            context_p_->adapter_ptr()->vma().create_staging_buffer(nbytes_)),
+        mapped_data_(nullptr) {}
 
   StagingBuffer(const StagingBuffer&) = delete;
   StagingBuffer& operator=(const StagingBuffer&) = delete;
@@ -58,7 +61,10 @@ class StagingBuffer final {
   }
 
   inline void* data() {
-    return vulkan_buffer_.allocation_info().pMappedData;
+    if (!mapped_data_) {
+      mapped_data_ = vulkan_buffer_.allocation_info().pMappedData;
+    }
+    return mapped_data_;
   }
 
   inline size_t numel() {

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

@@ -436,17 +436,6 @@ size_t vTensor::staging_buffer_numel() const {
   return padded_numel_;
 }
 
-VmaAllocationCreateInfo vTensor::get_allocation_create_info() const {
-  switch (storage_type()) {
-    case utils::kBuffer:
-      return storage_.buffer_.allocation_create_info();
-    case utils::kTexture2D:
-    case utils::kTexture3D:
-      return storage_.image_.allocation_create_info();
-  }
-  return {};
-}
-
 VkMemoryRequirements vTensor::get_memory_requirements() const {
   switch (storage_type()) {
     case utils::kBuffer:
@@ -567,6 +556,48 @@ void vTensor::virtual_resize(const std::vector<int64_t>& new_sizes) {
   update_metadata();
 }
 
+/*
+ * Transposing the dim order is a bit unintuitive. dim0 and dim1 have swapped
+ * their "identities", so we need to swap the values of dim0 and dim1 wherever
+ * they appear in the dim order vector. Compare this to just swapping the
+ * elements at dim0 and dim1 in the `sizes` vectors.
+ */
+void transpose_dim_order_inplace(
+    std::vector<int64_t>& dim_order,
+    const int64_t dim0,
+    const int64_t dim1) {
+  for (int i = 0; i < dim_order.size(); ++i) {
+    if (dim_order[i] == dim0) {
+      dim_order[i] = dim1;
+    } else if (dim_order[i] == dim1) {
+      dim_order[i] = dim0;
+    }
+  }
+}
+
+void vTensor::virtual_transpose(const int64_t dim0, const int64_t dim1) {
+  std::iter_swap(sizes_.begin() + dim0, sizes_.begin() + dim1);
+  if (storage_type() == utils::kBuffer) {
+    transpose_dim_order_inplace(dim_order_, dim0, dim1);
+  } else {
+    const int dim0_whcn = sizes_.size() - 1 - dim0;
+    const int dim1_whcn = sizes_.size() - 1 - dim1;
+    // Cannot transpose batch dimension for texture storage
+    VK_CHECK_COND(dim0_whcn < 3 && dim1_whcn < 3);
+
+    std::iter_swap(
+        axis_map_.begin() + dim0_whcn, axis_map_.begin() + dim1_whcn);
+
+    if (packed_dim_whcn_idx() == dim0_whcn) {
+      memory_layout_ = utils::GPUMemoryLayout(dim1_whcn);
+    }
+    if (packed_dim_whcn_idx() == dim1_whcn) {
+      memory_layout_ = utils::GPUMemoryLayout(dim0_whcn);
+    }
+  }
+  update_metadata();
+}
+
 void vTensor::reallocate(const std::vector<int64_t>& new_sizes) {
   sizes_ = new_sizes;
   update_metadata();

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

@@ -530,6 +530,11 @@ class vTensor final {
    */
   void virtual_resize(const std::vector<int64_t>& new_sizes);
 
+  /*
+   * Transpose the tensor in-place by updating its metadata.
+   */
+  void virtual_transpose(const int64_t dim0, const int64_t dim1);
+
   /*
    * Discard the underlying VkImage or VkBuffer and re-allocate based on new
    * tensor sizes

backends/vulkan/runtime/graph/containers/SharedObject.cpp

@@ -15,47 +15,27 @@ namespace vkcompute {
 void SharedObject::add_user(ComputeGraph* const graph, const ValueRef idx) {
   vTensorPtr t = graph->get_tensor(idx);
 
-  //
   // Aggregate Memory Requirements
-  //
-
   const VkMemoryRequirements mem_reqs = t->get_memory_requirements();
   aggregate_memory_requirements.size =
       std::max(mem_reqs.size, aggregate_memory_requirements.size);
   aggregate_memory_requirements.alignment =
       std::max(mem_reqs.alignment, aggregate_memory_requirements.alignment);
   aggregate_memory_requirements.memoryTypeBits |= mem_reqs.memoryTypeBits;
 
-  //
-  // Aggregate Allocation Create Info
-  //
-
-  const VmaAllocationCreateInfo create_info = t->get_allocation_create_info();
-  // Clear out CREATE_STRATEGY bit flags in case of conflict
-  VmaAllocationCreateFlags clear_mask = ~VMA_ALLOCATION_CREATE_STRATEGY_MASK;
-  VmaAllocationCreateFlags create_flags = create_info.flags & clear_mask;
-  // Use the default allocation strategy
-  aggregate_create_info.flags =
-      create_flags | vkapi::DEFAULT_ALLOCATION_STRATEGY;
-
-  // Set the usage flag if it is currently not set
-  if (aggregate_create_info.usage == VMA_MEMORY_USAGE_UNKNOWN) {
-    aggregate_create_info.usage = create_info.usage;
-  }
-  // Otherwise check that there is no conflict regarding usage
-  VK_CHECK_COND(aggregate_create_info.usage == create_info.usage);
-  aggregate_create_info.requiredFlags |= create_info.requiredFlags;
-  aggregate_create_info.preferredFlags |= create_info.preferredFlags;
-
   users.emplace_back(idx);
 }
 
 void SharedObject::allocate(ComputeGraph* const graph) {
   if (aggregate_memory_requirements.size == 0) {
     return;
   }
+
+  VmaAllocationCreateInfo alloc_create_info =
+      graph->context()->adapter_ptr()->vma().gpuonly_resource_create_info();
+
   allocation = graph->context()->adapter_ptr()->vma().create_allocation(
-      aggregate_memory_requirements, aggregate_create_info);
+      aggregate_memory_requirements, alloc_create_info);
 }
 
 void SharedObject::bind_users(ComputeGraph* const graph) {

backends/vulkan/runtime/graph/containers/SharedObject.h

@@ -28,7 +28,6 @@ struct SharedObject {
   explicit SharedObject() = default;
 
   VkMemoryRequirements aggregate_memory_requirements;
-  VmaAllocationCreateInfo aggregate_create_info;
   std::vector<ValueRef> users;
   vkapi::Allocation allocation;
 

backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h

@@ -270,7 +270,7 @@ ivec3 to_texture_pos(
 
   // Adjust batch dim if needed
   if (sizes.w > 1) {
-    pos[axis_map[axis_map[3]]] += idx.w * sizes.w;
+    pos[axis_map[axis_map.w]] += idx.w * sizes[axis_map.w];
   }
 
   // Adjust packed dim. Moving 1 texel unit along the packed dim traverses 4
@@ -317,7 +317,7 @@ ivec4 to_texture_elem_pos(
 
   // Adjust batch dim if needed
   if (sizes.w > 1) {
-    pos[axis_map[axis_map[3]]] += idx.w * sizes.w;
+    pos[axis_map[axis_map.w]] += idx.w * sizes[axis_map.w];
   }
 
   // Adjust packed dim. Moving 1 texel unit along the packed dim traverses 4

backends/vulkan/runtime/vk_api/memory/Allocation.cpp

@@ -26,58 +26,37 @@ namespace vkcompute {
 namespace vkapi {
 
 Allocation::Allocation()
-    : memory_requirements{},
-      create_info{},
-      allocator(VK_NULL_HANDLE),
-      allocation(VK_NULL_HANDLE),
-      allocation_info({}),
-      is_copy_(false) {}
+    : allocator(VK_NULL_HANDLE), allocation(VK_NULL_HANDLE), is_copy_(false) {}
 
 Allocation::Allocation(
     VmaAllocator vma_allocator,
     const VkMemoryRequirements& mem_props,
     const VmaAllocationCreateInfo& create_info)
-    : memory_requirements(mem_props),
-      create_info(create_info),
-      allocator(vma_allocator),
-      allocation(VK_NULL_HANDLE),
-      allocation_info({}),
-      is_copy_(false) {
+    : allocator(vma_allocator), allocation(VK_NULL_HANDLE), is_copy_(false) {
   VK_CHECK(vmaAllocateMemory(
-      allocator, &memory_requirements, &create_info, &allocation, nullptr));
+      allocator, &mem_props, &create_info, &allocation, nullptr));
 }
 
 Allocation::Allocation(const Allocation& other) noexcept
-    : memory_requirements(other.memory_requirements),
-      create_info(other.create_info),
-      allocator(other.allocator),
+    : allocator(other.allocator),
       allocation(other.allocation),
-      allocation_info(other.allocation_info),
       is_copy_(true) {}
 
 Allocation::Allocation(Allocation&& other) noexcept
-    : memory_requirements(other.memory_requirements),
-      create_info(other.create_info),
-      allocator(other.allocator),
+    : allocator(other.allocator),
       allocation(other.allocation),
-      allocation_info(other.allocation_info),
       is_copy_(other.is_copy_) {
   other.allocation = VK_NULL_HANDLE;
-  other.allocation_info = {};
 }
 
 Allocation& Allocation::operator=(Allocation&& other) noexcept {
   VmaAllocation tmp_allocation = allocation;
 
-  memory_requirements = other.memory_requirements;
-  create_info = other.create_info;
   allocator = other.allocator;
   allocation = other.allocation;
-  allocation_info = other.allocation_info;
   is_copy_ = other.is_copy_;
 
   other.allocation = tmp_allocation;
-  other.allocation_info = {};
 
   return *this;
 }

backends/vulkan/runtime/vk_api/memory/Allocation.h

@@ -55,15 +55,10 @@ struct Allocation final {
 
   ~Allocation();
 
-  VkMemoryRequirements memory_requirements;
-  // The properties this allocation was created with
-  VmaAllocationCreateInfo create_info;
   // The allocator object this was allocated from
   VmaAllocator allocator;
   // Handles to the allocated memory
   VmaAllocation allocation;
-  // Information about the allocated memory
-  VmaAllocationInfo allocation_info;
 
  private:
   // Indicates whether this class instance is a copy of another class instance,

backends/vulkan/runtime/vk_api/memory/Allocator.cpp

@@ -58,6 +58,13 @@ Allocator::~Allocator() {
   vmaDestroyAllocator(allocator_);
 }
 
+VmaAllocationCreateInfo Allocator::gpuonly_resource_create_info() {
+  VmaAllocationCreateInfo alloc_create_info = {};
+  alloc_create_info.flags = DEFAULT_ALLOCATION_STRATEGY;
+  alloc_create_info.usage = VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE;
+  return alloc_create_info;
+}
+
 Allocation Allocator::create_allocation(
     const VkMemoryRequirements& memory_requirements,
     const VmaAllocationCreateInfo& create_info) {
@@ -103,9 +110,7 @@ VulkanImage Allocator::create_image(
         (VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT);
   }
 
-  VmaAllocationCreateInfo alloc_create_info = {};
-  alloc_create_info.flags = DEFAULT_ALLOCATION_STRATEGY;
-  alloc_create_info.usage = VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE;
+  VmaAllocationCreateInfo alloc_create_info = gpuonly_resource_create_info();
 
   const VulkanImage::ImageProperties image_props{
       image_type,
@@ -157,10 +162,7 @@ VulkanBuffer Allocator::create_storage_buffer(
     const bool allocate_memory) {
   const VkBufferUsageFlags buffer_usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
 
-  VmaAllocationCreateInfo alloc_create_info = {};
-  alloc_create_info.flags = DEFAULT_ALLOCATION_STRATEGY;
-  alloc_create_info.usage = VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE;
-
+  VmaAllocationCreateInfo alloc_create_info = gpuonly_resource_create_info();
   return VulkanBuffer(
       allocator_, size, alloc_create_info, buffer_usage, allocate_memory);
 }

backends/vulkan/runtime/vk_api/memory/Allocator.h

@@ -48,6 +48,8 @@ class Allocator final {
   VmaAllocator allocator_;
 
  public:
+  VmaAllocationCreateInfo gpuonly_resource_create_info();
+
   Allocation create_allocation(
       const VkMemoryRequirements& memory_requirements,
       const VmaAllocationCreateInfo& create_info);

backends/vulkan/runtime/vk_api/memory/Buffer.cpp

@@ -58,16 +58,14 @@ VulkanBuffer::VulkanBuffer(
       nullptr, // pQueueFamilyIndices
   };
 
-  memory_.create_info = allocation_create_info;
-
   if (allocate_memory) {
     VK_CHECK(vmaCreateBuffer(
         allocator_,
         &buffer_create_info,
         &allocation_create_info,
         &handle_,
         &(memory_.allocation),
-        &(memory_.allocation_info)));
+        nullptr));
   } else {
     VmaAllocatorInfo allocator_info{};
     vmaGetAllocatorInfo(allocator_, &allocator_info);
@@ -137,6 +135,12 @@ VulkanBuffer::~VulkanBuffer() {
   }
 }
 
+VmaAllocationInfo VulkanBuffer::allocation_info() const {
+  VmaAllocationInfo info;
+  vmaGetAllocationInfo(allocator_, memory_.allocation, &info);
+  return info;
+}
+
 VkMemoryRequirements VulkanBuffer::get_memory_requirements() const {
   VkMemoryRequirements memory_requirements;
   vkGetBufferMemoryRequirements(this->device(), handle_, &memory_requirements);

backends/vulkan/runtime/vk_api/memory/Buffer.h

@@ -114,13 +114,7 @@ class VulkanBuffer final {
     return memory_.allocation;
   }
 
-  inline VmaAllocationInfo allocation_info() const {
-    return memory_.allocation_info;
-  }
-
-  inline VmaAllocationCreateInfo allocation_create_info() const {
-    return VmaAllocationCreateInfo(memory_.create_info);
-  }
+  VmaAllocationInfo allocation_info() const;
 
   inline VkBuffer handle() const {
     return handle_;

backends/vulkan/runtime/vk_api/memory/Image.cpp

@@ -159,8 +159,6 @@ VulkanImage::VulkanImage(
       layout_, // initialLayout
   };
 
-  memory_.create_info = allocation_create_info;
-
   if (allocate_memory) {
     VK_CHECK(vmaCreateImage(
         allocator_,

backends/vulkan/runtime/vk_api/memory/Image.h

@@ -169,10 +169,6 @@ class VulkanImage final {
     return memory_.allocation;
   }
 
-  inline VmaAllocationCreateInfo allocation_create_info() const {
-    return VmaAllocationCreateInfo(memory_.create_info);
-  }
-
   inline VkFormat format() const {
     return image_properties_.image_format;
   }