Fix Validation Layer warnings about wrong image layout (#2854)

Summary:
Pull Request resolved: #2854

## Context

Currently, when executing a `ComputeGraph` with prepacked tensors with [Vulkan Validation Layers](https://github.com/KhronosGroup/Vulkan-ValidationLayers) turned on, the following Validation Errors can be observed. Note that Validation Layers can be turned on by running Vulkan binaries on Mac with the `vkconfig` app opened.

```
UNASSIGNED-CoreValidation-DrawState-InvalidImageLayout(ERROR / SPEC): msgNum: 1303270965 - Validation Error: [ UNASSIGNED-CoreValidation-DrawState-InvalidImageLayout ] Object 0: handle = 0x7fb76dbbf988, type = VK_OBJECT_TYPE_COMMAND_BUFFER; | MessageID = 0x4dae5635 | vkQueueSubmit(): pSubmits[0].pCommandBuffers[0] command buffer VkCommandBuffer 0x7fb76dbbf988[] expects VkImage 0xd79c8a0000000f09[] (subresource: aspectMask 0x1 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: 1
        [0] 0x7fb76dbbf988, type: 6, name: NULL
```

The reason for this is that prepacked textures are written to with `WRITE` memory access during packing, which means they will be in the `VK_IMAGE_LAYOUT_GENERAL` layout. However, they will subsequently be read from during `graph.execute()`, meaning the texture will have transitioned to `VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL`, but will be bound using the `VK_IMAGE_LAYOUT_GENERAL` layout. Subsequent calls to `execute()` will therefore see that the prepacked texture has been bound with the wrong layout, since after the first graph execution the texture will have the `VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL` layout.

The solution is to submit a no-op shader dispatch during prepacking to trigger a transition to the `READ_ONLY_OPTIMAL` layout.
ghstack-source-id: 221871426

bypass-github-pytorch-ci-checks

Reviewed By: jorgep31415

Differential Revision: D55772003

fbshipit-source-id: f9c69e6e571ca0d0d28a6c25716766af98e82d41

backends/vulkan/runtime/graph/ops/PrepackNode.cpp

@@ -11,10 +11,19 @@
 #include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
 
 #include <executorch/backends/vulkan/runtime/graph/ops/utils/BindingUtils.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
 #include <executorch/backends/vulkan/runtime/graph/ops/utils/StagingUtils.h>
 
 namespace vkcompute {
 
+api::ShaderInfo get_noop_shader(ComputeGraph& graph, const ValueRef packed) {
+  std::stringstream noop_shader_name;
+  noop_shader_name << "no_op";
+  apply_ndim_suffix(noop_shader_name, graph.get_val(packed).toTensor());
+  apply_dtype_suffix(noop_shader_name, graph.get_val(packed).toTensor());
+  return VK_KERNEL_FROM_STR(noop_shader_name.str());
+}
+
 PrepackNode::PrepackNode(
     ComputeGraph& graph,
     const api::ShaderInfo& shader,
@@ -24,17 +33,18 @@ PrepackNode::PrepackNode(
     const ValueRef packed,
     const std::vector<std::shared_ptr<api::UniformParamsBuffer>>& params)
     : shader_(shader),
+      noop_shader_(get_noop_shader(graph, packed)),
       global_workgroup_size_(global_workgroup_size),
       local_workgroup_size_(local_workgroup_size),
       tref_(tref),
       packed_(packed),
       params_(params) {
   graph.update_descriptor_counts(shader, /*execute = */ false);
+  graph.update_descriptor_counts(noop_shader_, /*execute = */ false);
 }
 
 void PrepackNode::encode(ComputeGraph* graph) {
   api::Context* const context = graph->context();
-  api::PipelineBarrier pipeline_barrier{};
 
   TensorRef& tref = graph->get_val(tref_).toTensorRef();
   vTensor& packed = graph->get_val(packed_).toTensor();
@@ -46,21 +56,44 @@ void PrepackNode::encode(ComputeGraph* graph) {
 
   std::unique_lock<std::mutex> cmd_lock = context->dispatch_lock();
 
-  api::DescriptorSet descriptor_set =
-      context->get_descriptor_set(shader_, local_workgroup_size_);
-
-  uint32_t idx = 0;
-  bind_tensor_to_descriptor_set(
-      packed,
-      pipeline_barrier,
-      api::MemoryAccessType::WRITE,
-      descriptor_set,
-      idx++);
-  bind_staging_to_descriptor_set(staging, descriptor_set, idx++);
-  bind_params_to_descriptor_set(params_, descriptor_set, idx);
-
-  context->register_shader_dispatch(
-      descriptor_set, pipeline_barrier, shader_, global_workgroup_size_);
+  {
+    api::PipelineBarrier pipeline_barrier{};
+    api::DescriptorSet descriptor_set =
+        context->get_descriptor_set(shader_, local_workgroup_size_);
+
+    uint32_t idx = 0;
+    bind_tensor_to_descriptor_set(
+        packed,
+        pipeline_barrier,
+        api::MemoryAccessType::WRITE,
+        descriptor_set,
+        idx++);
+    bind_staging_to_descriptor_set(staging, descriptor_set, idx++);
+    bind_params_to_descriptor_set(params_, descriptor_set, idx);
+
+    context->register_shader_dispatch(
+        descriptor_set, pipeline_barrier, shader_, global_workgroup_size_);
+  }
+
+  // Submit a compute shader that performs a no-op with the packed tensor in
+  // order to trigger a image layout transition from GENERAL to
+  // READ_ONLY_OPTIMAL. This ensures that future uses of the tensor will be
+  // bound with the correct image layout.
+  {
+    api::PipelineBarrier pipeline_barrier{};
+    api::DescriptorSet descriptor_set =
+        context->get_descriptor_set(noop_shader_, {1, 1, 1});
+
+    bind_tensor_to_descriptor_set(
+        packed,
+        pipeline_barrier,
+        api::MemoryAccessType::READ,
+        descriptor_set,
+        0);
+
+    context->register_shader_dispatch(
+        descriptor_set, pipeline_barrier, noop_shader_, {1, 1, 1});
+  }
 }
 
 } // namespace vkcompute

backends/vulkan/runtime/graph/ops/PrepackNode.h

@@ -41,6 +41,7 @@ class PrepackNode final {
 
  protected:
   const api::ShaderInfo shader_;
+  api::ShaderInfo noop_shader_;
   const api::utils::uvec3 global_workgroup_size_;
   const api::utils::uvec3 local_workgroup_size_;
   const ValueRef tref_;

backends/vulkan/runtime/graph/ops/glsl/no_op.glsl

@@ -0,0 +1,24 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#version 450 core
+
+#include "broadcasting_utils.h"
+#include "indexing_utils.h"
+
+#define PRECISION ${PRECISION}
+
+#define OP(X, Y, A) ${OPERATOR}
+
+layout(std430) buffer;
+
+layout(set = 0, binding = 0) uniform PRECISION ${SAMPLER_T[NDIM][DTYPE]} image_in;
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+void main() {}

backends/vulkan/runtime/graph/ops/glsl/no_op.yaml

@@ -0,0 +1,26 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+no_op:
+  parameter_names_with_default_values:
+    OPERATOR: X + A * Y
+    NDIM: 3
+    DTYPE: float
+  generate_variant_forall:
+    NDIM:
+      - VALUE: 3
+        SUFFIX: 3d
+      - VALUE: 2
+        SUFFIX: 2d
+    DTYPE:
+      - VALUE: half
+        SUFFIX: half
+      - VALUE: float
+        SUFFIX: float
+      - VALUE: int
+        SUFFIX: int
+  shader_variants:
+    - NAME: no_op

backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.cpp

@@ -26,6 +26,19 @@ void apply_dtype_suffix(std::stringstream& kernel_name, const vTensor& tensor) {
   }
 }
 
+void apply_ndim_suffix(std::stringstream& kernel_name, const vTensor& tensor) {
+  switch (tensor.storage_type()) {
+    case api::StorageType::TEXTURE_3D:
+      kernel_name << "_3d";
+      break;
+    case api::StorageType::TEXTURE_2D:
+      kernel_name << "_2d";
+      break;
+    default:
+      break;
+  }
+}
+
 void apply_memory_layout_suffix(
     std::stringstream& kernel_name,
     const vTensor& tensor) {

backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h

@@ -16,6 +16,8 @@ namespace vkcompute {
 
 void apply_dtype_suffix(std::stringstream& kernel_name, const vTensor& tensor);
 
+void apply_ndim_suffix(std::stringstream& kernel_name, const vTensor& tensor);
+
 void apply_memory_layout_suffix(
     std::stringstream& kernel_name,
     const vTensor& tensor);