[ET-VK] Fix pack_fp_linear_weight for devices without VK_KHR_16bit_storage (pytorch#18653)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: pytorch#18642 by
@SS-JIA
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/SS-JIA/515/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/SS-JIA/515/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/main
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/SS-JIA/515/orig
Differential Revision:
[D99133993](https://our.internmc.facebook.com/intern/diff/D99133993/)
@diff-train-skip-merge

Co-authored-by: ssjia <ssjia@devvm26340.ftw0.facebook.com>

Author: pytorchbot

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

@@ -9,15 +9,16 @@
 #version 450 core
 
 #define PRECISION ${PRECISION}
-#define VEC4_T ${texel_load_type(DTYPE, "buffer")}
-#define T ${texel_load_component_type(DTYPE, "buffer")}
+#define BUF_T ${buffer_scalar_type(BUF_DTYPE)}
+#define VEC4_T ${texel_load_type(DTYPE, PACKED_STORAGE)}
+#define T ${texel_load_component_type(DTYPE, PACKED_STORAGE)}
 
 $if PACKED_STORAGE == "buffer":
   #define OUTPUT_BUFFER
 
 #extension GL_EXT_control_flow_attributes : require
 
-${define_required_extensions("buffer", DTYPE)}
+${define_required_extensions("buffer", BUF_DTYPE)}
 $if PACKED_STORAGE != "buffer":
   ${define_required_extensions(PACKED_STORAGE, DTYPE)}
 
@@ -29,7 +30,7 @@ $if PACKED_STORAGE == "buffer":
   ${layout_declare_tensor(B, "w", "t_weight_packed", DTYPE, "buffer", is_scalar_array=False)}
 $else:
   ${layout_declare_tensor(B, "w", "t_weight_packed", DTYPE, PACKED_STORAGE, is_scalar_array=False)}
-${layout_declare_tensor(B, "r", "t_weight_src", DTYPE, "buffer", is_scalar_array=True)}
+${layout_declare_tensor(B, "r", "t_weight_src", BUF_DTYPE, "buffer", is_scalar_array=True)}
 
 layout(push_constant) uniform restrict Block {
   int N;

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

@@ -7,13 +7,16 @@
 pack_fp_linear_weight:
   parameter_names_with_default_values:
     DTYPE: float
+    BUF_DTYPE: float
     PACKED_STORAGE: texture2d
   generate_variant_forall:
-    PACKED_STORAGE:
-      - VALUE: texture2d
-      - VALUE: buffer
-    DTYPE:
-      - VALUE: float
-      - VALUE: half
+    combination:
+      parameter_names: [PACKED_STORAGE, DTYPE, BUF_DTYPE]
+      combos:
+        - parameter_values: [texture2d, float, float]
+        - parameter_values: [texture2d, half, half]
+        - parameter_values: [texture2d, half, float]
+        - parameter_values: [buffer, float, float]
+        - parameter_values: [buffer, half, half]
   shader_variants:
     - NAME: pack_fp_linear_weight

backends/vulkan/runtime/graph/ops/impl/Conv1dPW.cpp

@@ -71,6 +71,7 @@ static ValueRef prepack_conv1d_pw_weight(
   std::string kernel_name = "pack_fp_linear_weight";
   add_storage_type_suffix(kernel_name, weight_storage);
   add_dtype_suffix(kernel_name, graph.dtype_of(weight_data));
+  add_dtype_suffix(kernel_name, graph.get_staging_dtype_for(weight_data));
 
   graph.prepack_nodes().emplace_back(new PrepackNode(
       graph,

backends/vulkan/runtime/graph/ops/impl/Conv2dPW.cpp

@@ -130,6 +130,7 @@ ValueRef prepack_conv2d_pw_weight(
   std::string pack_kernel_name = "pack_fp_linear_weight";
   add_storage_type_suffix(pack_kernel_name, weight_storage);
   add_dtype_suffix(pack_kernel_name, graph.dtype_of(weight_data));
+  add_dtype_suffix(pack_kernel_name, graph.get_staging_dtype_for(weight_data));
 
   graph.prepack_nodes().emplace_back(new PrepackNode(
       graph,

backends/vulkan/runtime/graph/ops/impl/Linear.cpp

@@ -81,6 +81,7 @@ ValueRef prepack_fp_linear_weight(
   std::string kernel_name = "pack_fp_linear_weight";
   add_storage_type_suffix(kernel_name, weight_storage);
   add_dtype_suffix(kernel_name, graph.dtype_of(weight_data));
+  add_dtype_suffix(kernel_name, graph.get_staging_dtype_for(weight_data));
 
   graph.prepack_nodes().emplace_back(new PrepackNode(
       graph,