[ET-VK][ez] Fix Vulkan Validation layer errors due to consecutive command buffer encoding

Pull Request resolved: #11401

## Changes

* In `VulkanBackend.cpp` do not call `encode_execute()` during model load if the model compile spec specifies `requires_dynamic_shapes` as true
* `propagate_resize()` will not call `encode_execute()` if dynamic shapes are not expected.

## Motivation

> In `VulkanBackend.cpp` do not call `encode_execute()` during model load if the model compile spec specifies `requires_dynamic_shapes` as true

Recently, it was discovered that a command buffer re-encode was required to update push constant values. This means that for dynamic shapes to work correctly, `encode_execute()` must be called after updating tensor sizes.

As a result, `propagate_resize()` now calls `encode_execute()` internally. This results in scenarios where `encode_execute()` is called once during model load, then again right before the first inference during `propagate_resize()`, without actually executing the command buffer in-between.

This causes Validation layer errors like

```

UNASSIGNED-CoreValidation-DrawState-InvalidImageLayout(ERROR / SPEC): msgNum: 1303270965 - Validation Error: [ UNASSIGNED-CoreValidation-DrawState-InvalidImageLayout ] Object 0: handle = 0x24086224ec0, type = VK_OBJECT_TYPE_COMMAND_BUFFER; Object 1: handle = 0x88d2b500000000e2, type = VK_OBJECT_TYPE_IMAGE; | MessageID = 0x4dae5635 | vkQueueSubmit(): pSubmits[0].pCommandBuffers[0] command buffer VkCommandBuffer 0x24086224ec0[] expects VkImage 0x88d2b500000000e2[] (subresource: aspectMask VK_IMAGE_ASPECT_COLOR_BIT array layer 0, mip level 0) to be in layout VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL--instead, current layout is VK_IMAGE_LAYOUT_UNDEFINED.
    Objects: 2
        [0] 0x24086224ec0, type: 6, name: NULL
        [1] 0x88d2b500000000e2, type: 10, name: NULL
UNASSIGNED-CoreValidation-DrawState-InvalidImageLayout(ERROR / SPEC): msgNum: 1303270965 - Validation Error: [ UNASSIGNED-CoreValidation-DrawState-InvalidImageLayout ] Object 0: handle = 0x24086224ec0, type = VK_OBJECT_TYPE_COMMAND_BUFFER; Object 1: handle = 0x6caffc00000000e3, type = VK_OBJECT_TYPE_IMAGE; | MessageID = 0x4dae5635 | vkQueueSubmit(): pSubmits[0].pCommandBuffers[0] command buffer VkCommandBuffer 0x24086224ec0[] expects VkImage 0x6caffc00000000e3[] (subresource: aspectMask VK_IMAGE_ASPECT_COLOR_BIT array layer 0, mip level 0) to be in layout VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL--instead, current layout is VK_IMAGE_LAYOUT_UNDEFINED.
    Objects: 2
        [0] 0x24086224ec0, type: 6, name: NULL
        [1] 0x6caffc00000000e3, type: 10, name: NULL

```

because the last access information of image/buffer resources are inaccurate during the second command buffer encoding, since the first command buffer never executed.

>  `propagate_resize()` will not call `encode_execute()` if dynamic shapes are not expected.

Llama models will need to call `propagate_resize()`, to account for changes in the second input tensor (a scalar tensor which represents input token position). However, the cost of `encode_execute()` is costly, so we should avoid calling it if possible.

## Perf Impact

* Performance improvement for first inference of dynamic shape models if actual tensor sizes are much smaller than maximum possible sizes
* No impact for non-dynamic shape models
ghstack-source-id: 289110008

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

Author: SS-JIA

backends/vulkan/runtime/VulkanBackend.cpp

@@ -140,6 +140,14 @@ GraphConfig get_graph_config(ArrayRef<CompileSpec>& compile_specs) {
 
       config.set_memory_layout_override(memory_layout);
     }
+    if (strcmp(spec.key, "require_dynamic_shapes") == 0) {
+      ET_CHECK_MSG(value_size == sizeof(uint8_t), "Unexpected value size!");
+      bool value = getBool(value_data);
+
+      if (value) {
+        config.expect_dynamic_shapes = true;
+      }
+    }
   }
 #ifdef ET_EVENT_TRACER_ENABLED
   config.enable_querypool = true;
@@ -500,9 +508,12 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface {
     compute_graph->encode_prepack();
     compute_graph->prepack();
 
-    // TODO(ssjia): remove this once we can batch compile compute pipelines
-    // during prepare().
-    compute_graph->encode_execute();
+    // If dynamic shapes are not expected, then the command buffer only needs to
+    // be encoded once. Otherwise, wait until the first inference to encode the
+    // the command buffer, when actual input shapes are known.
+    if (!compute_graph->graphconfig().expect_dynamic_shapes) {
+      compute_graph->encode_execute();
+    }
 
     return Error::Ok;
   }
@@ -574,7 +585,9 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface {
     // constants are updated and DynamicDispatchNode can update the compute
     // shader, global workgroup size, and local workgroup size to perform the
     // model inference.
-    if (should_propagate_resize) {
+    if (should_propagate_resize ||
+        (compute_graph->graphconfig().expect_dynamic_shapes &&
+         compute_graph->execute_count() == 0u)) {
       compute_graph->propagate_resize();
     }
 

backends/vulkan/runtime/VulkanDelegateHeader.cpp

@@ -60,6 +60,10 @@ uint32_t getUInt16LE(const uint8_t* data) {
   return (uint32_t)data[0] | ((uint32_t)data[1] << 8);
 }
 
+bool getBool(const uint8_t* data) {
+  return data[0] != 0;
+}
+
 bool VulkanDelegateHeader::is_valid() const {
   if (header_size < kExpectedSize) {
     return false;

backends/vulkan/runtime/VulkanDelegateHeader.h

@@ -19,6 +19,9 @@ uint64_t getUInt64LE(const uint8_t* data);
 uint32_t getUInt32LE(const uint8_t* data);
 uint32_t getUInt16LE(const uint8_t* data);
 
+// Bool is serialized as a single byte
+bool getBool(const uint8_t* data);
+
 struct VulkanDelegateHeader {
   bool is_valid() const;
 

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -761,7 +761,10 @@ void ComputeGraph::propagate_resize() {
   for (std::unique_ptr<ExecuteNode>& node : execute_nodes_) {
     node->trigger_resize(this);
   }
-  encode_execute();
+  // Only re-encode on resize if dynamic shapes are expected
+  if (config_.expect_dynamic_shapes) {
+    encode_execute();
+  }
 }
 
 } // namespace vkcompute

backends/vulkan/runtime/graph/GraphConfig.cpp

@@ -63,6 +63,8 @@ GraphConfig::GraphConfig() {
 
   enable_local_wg_size_override = false;
   local_wg_size_override = {};
+
+  expect_dynamic_shapes = false;
 }
 
 void GraphConfig::set_storage_type_override(utils::StorageType storage_type) {

backends/vulkan/runtime/graph/GraphConfig.h

@@ -33,6 +33,9 @@ struct GraphConfig final {
   bool enable_local_wg_size_override;
   utils::uvec3 local_wg_size_override;
 
+  // Whether or not the ComputeGraph should expect input shapes to be dynamic
+  bool expect_dynamic_shapes;
+
   // Generate a default graph config with pre-configured settings
   explicit GraphConfig();
 

backends/vulkan/test/test_vulkan_delegate.py

@@ -43,6 +43,9 @@ def lower_module(
     model: torch.nn.Module, sample_inputs: Tuple[torch.Tensor], dynamic_shapes=None
 ) -> EdgeProgramManager:
     compile_options = {}
+    if dynamic_shapes is not None:
+        compile_options["require_dynamic_shapes"] = True
+
     edge_compile_config = EdgeCompileConfig(
         _skip_dim_order=False,  # TODO(T182928844): Delegate dim order op to backend.
     )
@@ -70,6 +73,9 @@ def quantize_and_lower_module(
     dynamic_shapes=None,
 ) -> EdgeProgramManager:
     compile_options = {}
+    if dynamic_shapes is not None:
+        compile_options["require_dynamic_shapes"] = True
+
     edge_compile_config = EdgeCompileConfig(
         _skip_dim_order=False,  # TODO(T182928844): Delegate dim order op to backend.
     )

backends/vulkan/test/utils/test_utils.cpp

@@ -512,6 +512,7 @@ vkcompute::ComputeGraph build_mm_graph(
     const bool prepack_mat2) {
   using namespace vkcompute;
   GraphConfig config;
+  config.expect_dynamic_shapes = true;
   ComputeGraph graph(config);
 
   std::vector<int64_t> mat1_size = {M, K};

backends/vulkan/test/vulkan_compute_api_test.cpp

@@ -2937,6 +2937,7 @@ void test_transpose_view_mm(
     const int N,
     utils::StorageType storage_type) {
   GraphConfig config;
+  config.expect_dynamic_shapes = true;
   config.set_storage_type_override(storage_type);
   ComputeGraph graph(config);
 
@@ -2993,7 +2994,6 @@ void test_transpose_view_mm(
   graph.prepare();
   graph.encode_prepack();
   graph.prepack();
-  graph.encode_execute();
 
   for (int i = 1; i < 4; i++) {
     float val_mat1 = i;