From 1c28f78a9455d2cd9d093ef1abead3499df4f105 Mon Sep 17 00:00:00 2001
From: Nathanael See <nathanaelsee@fb.com>
Date: Wed, 26 Mar 2025 10:49:35 -0700
Subject: [PATCH] [BE][ET-VK] update clone to use layout gen

Pull Request resolved: https://github.com/pytorch/executorch/pull/9589

TSIA

@pytorchbot label "topic: not user facing"
ghstack-source-id: 274222181
@exported-using-ghexport

Differential Revision: [D71825480](https://our.internmc.facebook.com/intern/diff/D71825480/)
---
 backends/vulkan/runtime/graph/ops/glsl/clone.glsl | 13 ++++++-------
 backends/vulkan/runtime/graph/ops/glsl/clone.yaml |  1 +
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/backends/vulkan/runtime/graph/ops/glsl/clone.glsl b/backends/vulkan/runtime/graph/ops/glsl/clone.glsl
index 64def8d7000..3bd1af8bb0c 100644
--- a/backends/vulkan/runtime/graph/ops/glsl/clone.glsl
+++ b/backends/vulkan/runtime/graph/ops/glsl/clone.glsl
@@ -8,16 +8,15 @@
 
 #version 450 core
 
+#include "indexing_utils.h"
+
 #define PRECISION ${PRECISION}
 
 layout(std430) buffer;
 
-layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} image_out;
-layout(set = 0, binding = 1) uniform PRECISION sampler3D image_in;
-
-layout(set = 0, binding = 2) uniform PRECISION restrict OutLimits {
-  ivec3 out_limits;
-};
+${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)}
+${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)}
+${layout_declare_ubo(B, "ivec3", "out_limits")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
@@ -26,5 +25,5 @@ void main() {
   if (any(greaterThanEqual(pos, out_limits))) {
     return;
   }
-  imageStore(image_out, pos, texelFetch(image_in, pos, 0));
+  imageStore(t_out, pos, load_texel(t_in, pos));
 }
diff --git a/backends/vulkan/runtime/graph/ops/glsl/clone.yaml b/backends/vulkan/runtime/graph/ops/glsl/clone.yaml
index 5dbce0e9d8d..1fdbf506bfd 100644
--- a/backends/vulkan/runtime/graph/ops/glsl/clone.yaml
+++ b/backends/vulkan/runtime/graph/ops/glsl/clone.yaml
@@ -2,6 +2,7 @@ clone:
   parameter_names_with_default_values:
     DTYPE: float
     NDIM: 3
+    STORAGE: texture3d
   generate_variant_forall:
     DTYPE:
       - VALUE: half