Update base for Update on "[ET-VK] Use shared pointer for vTensorStorage"

## Changes

* In `vTensor`, store the `storage_` member using a shared pointer instead of a direct `vTensorStorage` object
* Remove the ability to construct a tensor view with a buffer offset

## Motivation

> * In `vTensor`, store the `storage_` member using a shared pointer instead of a direct `vTensorStorage` object

Previously, to support the ability to create tensor views I implemented copy constructors for `VulkanImage`, `VulkanBuffer`, and `vTensorStorage`. The idea was that the copied object instance would have the same handles as the original, but would not own the handle (and therefore not be responsible for destroying it).

However, the problem with this approach is that in order to construct pipeline barriers and  image layout transitions properly, we must know the details of how a resource was last accessed. Since tensor views make copies, accessing the original tensor will not update the access information of any copies it may have. Therefore, if the copy is then accessed, the last access information it stores would then be out of date.

I originally tried to solve this problem with crude assumptions in the `transition()` function, but unfortunately the solution was not robust enough. I discovered validation errors such as

```
[ RUN      ] VulkanComputeGraphOpsTest.test_transpose_with_mm
VUID-VkImageMemoryBarrier-oldLayout-01197(ERROR / SPEC): msgNum: 307231540 - Validation Error: [ VUID-VkImageMemoryBarrier-oldLayout-01197 ] Object 0: handle = 0x21eba07e160, type = VK_OBJECT_TYPE_COMMAND_BUFFER; Object 1: handle = 0x35643f0000000111, type = VK_OBJECT_TYPE_IMAGE; | MessageID = 0x124ffb34 | vkCmdPipelineBarrier(): pImageMemoryBarriers[2].image (VkImage 0x35643f0000000111[]) cannot transition the layout of aspect=1, level=0, layer=0 from VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL when the previous known layout is VK_IMAGE_LAYOUT_GENERAL.
The Vulkan spec states: If srcQueueFamilyIndex and dstQueueFamilyIndex define a queue family ownership transfer or oldLayout and newLayout define an image layout transition, oldLayout must be VK_IMAGE_LAYOUT_UNDEFINED or the current layout of the image subresources affected by the barrier (https://vulkan.lunarg.com/doc/view/1.3.296.0/windows/1.3-extensions/vkspec.html#VUID-VkImageMemoryBarrier-oldLayout-01197)
    Objects: 2
        [0] 0x21eba07e160, type: 6, name: NULL
        [1] 0x35643f0000000111, type: 10, name: NULL
UNASSIGNED-CoreValidation-DrawState-InvalidImageLayout(ERROR / SPEC): msgNum: 1303270965 - Validation Error: [ UNASSIGNED-CoreValidation-DrawState-InvalidImageLayout ] Object 0: handle = 0x21eba07e160, type = VK_OBJECT_TYPE_COMMAND_BUFFER; Object 1: handle = 0x35643f0000000111, type = VK_OBJECT_TYPE_IMAGE; | MessageID = 0x4dae5635 | vkQueueSubmit(): pSubmits[0].pCommandBuffers[0] command buffer VkCommandBuffer 0x21eba07e160[] expects VkImage 0x35643f0000000111[] (subresource: aspectMask VK_IMAGE_ASPECT_COLOR_BIT array layer 0, mip level 0) to be in layout VK_IMAGE_LAYOUT_GENERAL--instead, current layout is VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL.
    Objects: 2
        [0] 0x21eba07e160, type: 6, name: NULL
        [1] 0x35643f0000000111, type: 10, name: NULL
VUID-VkImageMemoryBarrier-oldLayout-01197(ERROR / SPEC): msgNum: 307231540 - Validation Error: [ VUID-VkImageMemoryBarrier-oldLayout-01197 ] Object 0: handle = 0x21eba07e160, type = VK_OBJECT_TYPE_COMMAND_BUFFER; Object 1: handle = 0x35643f0000000111, type = VK_OBJECT_TYPE_IMAGE; | MessageID = 0x124ffb34 | vkCmdPipelineBarrier(): pImageMemoryBarriers[2].image (VkImage 0x35643f0000000111[]) cannot transition the layout of aspect=1, level=0, layer=0 from VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL when the previous known layout is VK_IMAGE_LAYOUT_GENERAL.
The Vulkan spec states: If srcQueueFamilyIndex and dstQueueFamilyIndex define a queue family ownership transfer or oldLayout and newLayout define an image layout transition, oldLayout must be VK_IMAGE_LAYOUT_UNDEFINED or the current layout of the image subresources affected by the barrier (https://vulkan.lunarg.com/doc/view/1.3.296.0/windows/1.3-extensions/vkspec.html#VUID-VkImageMemoryBarrier-oldLayout-01197)
    Objects: 2
        [0] 0x21eba07e160, type: 6, name: NULL
        [1] 0x35643f0000000111, type: 10, name: NULL
UNASSIGNED-CoreValidation-DrawState-InvalidImageLayout(ERROR / SPEC): msgNum: 1303270965 - Validation Error: [ UNASSIGNED-CoreValidation-DrawState-InvalidImageLayout ] Object 0: handle = 0x21eba07e160, type = VK_OBJECT_TYPE_COMMAND_BUFFER; Object 1: handle = 0x35643f0000000111, type = VK_OBJECT_TYPE_IMAGE; | MessageID = 0x4dae5635 | vkQueueSubmit(): pSubmits[0].pCommandBuffers[0] command buffer VkCommandBuffer 0x21eba07e160[] expects VkImage 0x35643f0000000111[] (subresource: aspectMask VK_IMAGE_ASPECT_COLOR_BIT array layer 0, mip level 0) to be in layout VK_IMAGE_LAYOUT_GENERAL--instead, current layout is VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL.
    Objects: 2
        [0] 0x21eba07e160, type: 6, name: NULL
        [1] 0x35643f0000000111, type: 10, name: NULL
VUID-VkImageMemoryBarrier-oldLayout-01197(ERROR / SPEC): msgNum: 307231540 - Validation Error: [ VUID-VkImageMemoryBarrier-oldLayout-01197 ] Object 0: handle = 0x21eba07e160, type = VK_OBJECT_TYPE_COMMAND_BUFFER; Object 1: handle = 0x35643f0000000111, type = VK_OBJECT_TYPE_IMAGE; | MessageID = 0x124ffb34 | vkCmdPipelineBarrier(): pImageMemoryBarriers[2].image (VkImage 0x35643f0000000111[]) cannot transition the layout of aspect=1, level=0, layer=0 from VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL when the previous known layout is VK_IMAGE_LAYOUT_GENERAL.
The Vulkan spec states: If srcQueueFamilyIndex and dstQueueFamilyIndex define a queue family ownership transfer or oldLayout and newLayout define an image layout transition, oldLayout must be VK_IMAGE_LAYOUT_UNDEFINED or the current layout of the image subresources affected by the barrier (https://vulkan.lunarg.com/doc/view/1.3.296.0/windows/1.3-extensions/vkspec.html#VUID-VkImageMemoryBarrier-oldLayout-01197)
    Objects: 2
        [0] 0x21eba07e160, type: 6, name: NULL
        [1] 0x35643f0000000111, type: 10, name: NULL
UNASSIGNED-CoreValidation-DrawState-InvalidImageLayout(ERROR / SPEC): msgNum: 1303270965 - Validation Error: [ UNASSIGNED-CoreValidation-DrawState-InvalidImageLayout ] Object 0: handle = 0x21eba07e160, type = VK_OBJECT_TYPE_COMMAND_BUFFER; Object 1: handle = 0x35643f0000000111, type = VK_OBJECT_TYPE_IMAGE; | MessageID = 0x4dae5635 | vkQueueSubmit(): pSubmits[0].pCommandBuffers[0] command buffer VkCommandBuffer 0x21eba07e160[] expects VkImage 0x35643f0000000111[] (subresource: aspectMask VK_IMAGE_ASPECT_COLOR_BIT array layer 0, mip level 0) to be in layout VK_IMAGE_LAYOUT_GENERAL--instead, current layout is VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL.
    Objects: 2
        [0] 0x21eba07e160, type: 6, name: NULL
        [1] 0x35643f0000000111, type: 10, name: NULL
VUID-VkImageMemoryBarrier-oldLayout-01197(ERROR / SPEC): msgNum: 307231540 - Validation Error: [ VUID-VkImageMemoryBarrier-oldLayout-01197 ] Object 0: handle = 0x21eba07e160, type = VK_OBJECT_TYPE_COMMAND_BUFFER; Object 1: handle = 0x35643f0000000111, type = VK_OBJECT_TYPE_IMAGE; | MessageID = 0x124ffb34 | vkCmdPipelineBarrier(): pImageMemoryBarriers[2].image (VkImage 0x35643f0000000111[]) cannot transition the layout of aspect=1, level=0, layer=0 from VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL when the previous known layout is VK_IMAGE_LAYOUT_GENERAL.
The Vulkan spec states: If srcQueueFamilyIndex and dstQueueFamilyIndex define a queue family ownership transfer or oldLayout and newLayout define an image layout transition, oldLayout must be VK_IMAGE_LAYOUT_UNDEFINED or the current layout of the image subresources affected by the barrier (https://vulkan.lunarg.com/doc/view/1.3.296.0/windows/1.3-extensions/vkspec.html#VUID-VkImageMemoryBarrier-oldLayout-01197)
    Objects: 2
        [0] 0x21eba07e160, type: 6, name: NULL
        [1] 0x35643f0000000111, type: 10, name: NULL
UNASSIGNED-CoreValidation-DrawState-InvalidImageLayout(ERROR / SPEC): msgNum: 1303270965 - Validation Error: [ UNASSIGNED-CoreValidation-DrawState-InvalidImageLayout ] Object 0: handle = 0x21eba07e160, type = VK_OBJECT_TYPE_COMMAND_BUFFER; Object 1: handle = 0x35643f0000000111, type = VK_OBJECT_TYPE_IMAGE; | MessageID = 0x4dae5635 | vkQueueSubmit(): pSubmits[0].pCommandBuffers[0] command buffer VkCommandBuffer 0x21eba07e160[] expects VkImage 0x35643f0000000111[] (subresource: aspectMask VK_IMAGE_ASPECT_COLOR_BIT array layer 0, mip level 0) to be in layout VK_IMAGE_LAYOUT_GENERAL--instead, current layout is VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL.
    Objects: 2
        [0] 0x21eba07e160, type: 6, name: NULL
        [1] 0x35643f0000000111, type: 10, name: NULL
VUID-VkImageMemoryBarrier-oldLayout-01197(ERROR / SPEC): msgNum: 307231540 - Validation Error: [ VUID-VkImageMemoryBarrier-oldLayout-01197 ] Object 0: handle = 0x21eba07e160, type = VK_OBJECT_TYPE_COMMAND_BUFFER; Object 1: handle = 0x35643f0000000111, type = VK_OBJECT_TYPE_IMAGE; | MessageID = 0x124ffb34 | vkCmdPipelineBarrier(): pImageMemoryBarriers[2].image (VkImage 0x35643f0000000111[]) cannot transition the layout of aspect=1, level=0, layer=0 from VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL when the previous known layout is VK_IMAGE_LAYOUT_GENERAL.
The Vulkan spec states: If srcQueueFamilyIndex and dstQueueFamilyIndex define a queue family ownership transfer or oldLayout and newLayout define an image layout transition, oldLayout must be VK_IMAGE_LAYOUT_UNDEFINED or the current layout of the image subresources affected by the barrier (https://vulkan.lunarg.com/doc/view/1.3.296.0/windows/1.3-extensions/vkspec.html#VUID-VkImageMemoryBarrier-oldLayout-01197)
    Objects: 2
        [0] 0x21eba07e160, type: 6, name: NULL
        [1] 0x35643f0000000111, type: 10, name: NULL
UNASSIGNED-CoreValidation-DrawState-InvalidImageLayout(ERROR / SPEC): msgNum: 1303270965 - Validation Error: [ UNASSIGNED-CoreValidation-DrawState-InvalidImageLayout ] Object 0: handle = 0x21eba07e160, type = VK_OBJECT_TYPE_COMMAND_BUFFER; Object 1: handle = 0x35643f0000000111, type = VK_OBJECT_TYPE_IMAGE; | MessageID = 0x4dae5635 | vkQueueSubmit(): pSubmits[0].pCommandBuffers[0] command buffer VkCommandBuffer 0x21eba07e160[] expects VkImage 0x35643f0000000111[] (subresource: aspectMask VK_IMAGE_ASPECT_COLOR_BIT array layer 0, mip level 0) to be in layout VK_IMAGE_LAYOUT_GENERAL--instead, current layout is VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL.
    Objects: 2
        [0] 0x21eba07e160, type: 6, name: NULL
        [1] 0x35643f0000000111, type: 10, name: NULL
[       OK ] VulkanComputeGraphOpsTest.test_transpose_with_mm (26 ms)
```

when using tensor views. The simplest solution is to use a shared pointer to store the tensor storage object, that way tensor views can point to the same underlying object instance, and last access metadata will be consistent among all tensor views.

> * Remove the ability to construct a tensor view with a buffer offset

Removed this because it turns out that buffer properties cannot have arbitrary offsets; there are restrictions for what offsets can be used. The validation errors say it all:

```
[ RUN      ] VulkanComputeGraphTest.test_simple_graph_with_view
VUID-VkWriteDescriptorSet-descriptorType-00328(ERROR / SPEC): msgNum: -368569266 - Validation Error: [ VUID-VkWriteDescriptorSet-descriptorType-00328 ] | MessageID = 0xea08144e | vkUpdateDescriptorSets(): pDescriptorWrites[1].pBufferInfo[0].offset (84) must be a multiple of device limit minStorageBufferOffsetAlignment 16 when descriptor type is VK_DESCRIPTOR_TYPE_STORAGE_BUFFER.
The Vulkan spec states: If descriptorType is VK_DESCRIPTOR_TYPE_STORAGE_BUFFER or VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC, the offset member of each element of pBufferInfo must be a multiple of VkPhysicalDeviceLimits::minStorageBufferOffsetAlignment (https://vulkan.lunarg.com/doc/view/1.3.296.0/windows/1.3-extensions/vkspec.html#VUID-VkWriteDescriptorSet-descriptorType-00328)
VUID-VkDescriptorBufferInfo-range-00342(ERROR / SPEC): msgNum: -371195848 - Validation Error: [ VUID-VkDescriptorBufferInfo-range-00342 ] Object 0: handle = 0xee647e0000000009, type = VK_OBJECT_TYPE_BUFFER; | MessageID = 0xe9e00038 | vkUpdateDescriptorSets(): pDescriptorWrites[1].pBufferInfo[0].range (224) is larger than buffer size (224) + offset (84).
The Vulkan spec states: If range is not equal to VK_WHOLE_SIZE, range must be less than or equal to the size of buffer minus offset (https://vulkan.lunarg.com/doc/view/1.3.296.0/windows/1.3-extensions/vkspec.html#VUID-VkDescriptorBufferInfo-range-00342)
    Objects: 1
        [0] 0xee647e0000000009, type: 9, name: NULL
VUID-VkBufferMemoryBarrier-size-01189(ERROR / SPEC): msgNum: -1238074894 - Validation Error: [ VUID-VkBufferMemoryBarrier-size-01189 ] Object 0: handle = 0x25d7d900870, type = VK_OBJECT_TYPE_COMMAND_BUFFER; | MessageID = 0xb63479f2 | vkCmdPipelineBarrier(): pBufferMemoryBarriers[0].size VkBuffer 0xee647e0000000009[] has offset 0x54 and size 0xe0 whose sum is greater than total size 0xe0.
The Vulkan spec states: If size is not equal to VK_WHOLE_SIZE, size must be less than or equal to than the size of buffer minus offset (https://vulkan.lunarg.com/doc/view/1.3.296.0/windows/1.3-extensions/vkspec.html#VUID-VkBufferMemoryBarrier-size-01189)
    Objects: 1
        [0] 0x25d7d900870, type: 6, name: NULL
[       OK ] VulkanComputeGraphTest.test_simple_graph_with_view (8 ms)
```

## Impact

* Heap allocation upon tensor construction
* Increased pointer chasing when accessing `vTensor` members

I will validate that the impact of this change does not regress load/inference latency.

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

[ghstack-poisoned]

Author: SS-JIA

.ci/scripts/build_llama_android.sh

@@ -42,7 +42,6 @@ build_llama_runner() {
     popd
     ANDROID_ABI=arm64-v8a
     cmake -DBUCK2="${BUCK2}" \
-    -DBUILD_TESTING=OFF \
     -DCMAKE_TOOLCHAIN_FILE="$ANDROID_NDK"/build/cmake/android.toolchain.cmake  \
     -DANDROID_ABI="${ANDROID_ABI}" \
     -DCMAKE_INSTALL_PREFIX=cmake-android-out \

.ci/scripts/test_llama.sh

@@ -169,7 +169,6 @@ cmake_build_llama_runner() {
     popd
     dir="examples/models/llama"
     retry cmake \
-        -DBUILD_TESTING=OFF \
         -DCMAKE_INSTALL_PREFIX=cmake-out \
         -DCMAKE_BUILD_TYPE="$CMAKE_BUILD_TYPE" \
         -Bcmake-out/${dir} \

.ci/scripts/test_llama_torchao_lowbit.sh

@@ -40,7 +40,6 @@ cmake --build cmake-out -j16 --target install --config Release
 
 # Install llama runner with torchao
 cmake -DPYTHON_EXECUTABLE=python \
-    -DBUILD_TESTING=OFF \
     -DCMAKE_BUILD_TYPE=Release \
     -DEXECUTORCH_BUILD_KERNELS_CUSTOM=ON \
     -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \

.ci/scripts/test_llava.sh

@@ -64,10 +64,9 @@ cmake_install_executorch_libraries_for_android() {
 
 
 LLAVA_COMMON_CMAKE_ARGS="                        \
-        -DBUILD_TESTING=OFF                      \
         -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \
         -DCMAKE_INSTALL_PREFIX=${BUILD_DIR}      \
-        -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}   \
+        -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}         \
         -DEXECUTORCH_BUILD_KERNELS_CUSTOM=ON     \
         -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON  \
         -DEXECUTORCH_BUILD_XNNPACK=ON"

backends/vulkan/test/vulkan_compute_api_test.cpp

@@ -3177,8 +3177,10 @@ void test_to_copy() {
 
   EXPECT_EQ(data_in.size(), output_data.size());
 
+#ifdef VULKAN_DEBUG
   float mse_ex = 0.0f;
   float mse_vk = 0.0f;
+#endif
 
   // check results
   for (size_t i = 0; i < output_data.size(); ++i) {
@@ -3200,6 +3202,9 @@ void test_to_copy() {
         std::bitset<16>(*output_bits).to_string() + ")";
 
     std::cout << msg << std::endl;
+
+    mse_ex += std::pow(expected_output - input, 2);
+    mse_vk += std::pow(output - input, 2);
 #endif
 
     // Note: Torch executor half "rounds up" when converting to fp16 whereas
@@ -3221,13 +3226,12 @@ void test_to_copy() {
     EXPECT_TRUE(
         (*output_bits == *expected_bits) ||
         /*rounding error*/ ((*output_bits + 1u) == *expected_bits));
-    mse_ex += std::pow(expected_output - input, 2);
-    mse_vk += std::pow(output - input, 2);
   }
 
+#ifdef VULKAN_DEBUG
   mse_ex /= output_data.size();
   mse_vk /= output_data.size();
-#ifdef VULKAN_DEBUG
+
   std::cout << "========================================================="
             << std::endl;
   std::cout << "mse_ex = " << mse_ex << ", mse_vk = " << mse_vk << std::endl;

examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaMARunner.mm

@@ -14,7 +14,6 @@
 
 using executorch::extension::llm::GenerationConfig;
 using executorch::extension::llm::Image;
-using executorch::extension::llm::TextLLMRunner;
 using executorch::runtime::Error;
 
 NSErrorDomain const LLaMARunnerErrorDomain = @"LLaMARunnerErrorDomain";
@@ -24,15 +23,15 @@ @interface LLaMARunner ()<ExecuTorchLogSink>
 @end
 
 @implementation LLaMARunner {
-  std::unique_ptr<TextLLMRunner> _runner;
+  std::unique_ptr<example::Runner> _runner;
 }
 
 - (instancetype)initWithModelPath:(NSString*)modelPath
                     tokenizerPath:(NSString*)tokenizerPath {
   self = [super init];
   if (self) {
     [ExecuTorchLog.sharedLog addSink:self];
-    _runner = example::create_llama_runner(
+    _runner = example::Runner::create(
         modelPath.UTF8String, tokenizerPath.UTF8String);
   }
   return self;

examples/models/llama/CMakeLists.txt

@@ -220,6 +220,7 @@ endif()
 target_include_directories(
   llama_main
   PUBLIC ${_common_include_directories}
+         ${EXECUTORCH_ROOT}/extension/llm/tokenizers/include
 )
 target_link_libraries(llama_main PUBLIC llama_runner ${link_libraries})
 target_compile_options(llama_main PUBLIC ${_common_compile_options})

examples/models/llama/main.cpp

@@ -81,13 +81,8 @@ int32_t main(int32_t argc, char** argv) {
   }
 #endif
   // create llama runner
-  std::unique_ptr<::executorch::extension::llm::TextLLMRunner> runner =
-      example::create_llama_runner(model_path, tokenizer_path, data_path);
-
-  if (runner == nullptr) {
-    ET_LOG(Error, "Failed to create llama runner");
-    return 1;
-  }
+  std::unique_ptr<example::Runner> runner =
+      example::Runner::create(model_path, tokenizer_path, data_path);
 
   if (warmup) {
     runner->warmup(prompt, /*max_new_tokens=*/seq_len);

examples/models/llama/runner/CMakeLists.txt

@@ -52,20 +52,23 @@ else()
   add_library(llama_runner SHARED ${_llama_runner__srcs})
 endif()
 
-# For extension_llm_runner
-if(NOT TARGET extension_llm_runner)
-  add_subdirectory(
-    ${EXECUTORCH_ROOT}/extension/llm/runner
-    ${CMAKE_CURRENT_BINARY_DIR}/../../../../extension/llm/runner
-  )
-endif()
-
 set(llama_runner_deps executorch_core extension_data_loader extension_module
-                      extension_tensor extension_flat_tensor extension_llm_runner
+                      extension_tensor extension_flat_tensor
 )
 
 target_link_libraries(llama_runner PUBLIC ${llama_runner_deps})
 
+target_include_directories(
+  llama_runner
+  INTERFACE ${_common_include_directories}
+)
+
+# Include tokenizers dependency
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+add_subdirectory(
+  ${EXECUTORCH_ROOT}/extension/llm/tokenizers
+  ${CMAKE_CURRENT_BINARY_DIR}/tokenizers
+)
 target_link_libraries(
   llama_runner PUBLIC tokenizers
 )