Add CPU/GPU transfer shaders for width and height packing (#2515)

Summary:
Pull Request resolved: #2515

## Context

Begin to lay the foundation for `TENSOR_HEIGHT_PACKED` and `TENSOR_WIDTH_PACKED` GPU memory layouts by generating `nchw_to_image` and `image_to_nchw` shader variants for those packing schemes.
ghstack-source-id: 219281500
exported-using-ghexport

Reviewed By: jorgep31415

Differential Revision: D55018341

fbshipit-source-id: 98e18754be35466adcb3fe2a7fefc549f444c075

Author: SS-JIA

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

@@ -44,18 +44,21 @@ void main() {
 
   const int base_index = COORD_TO_BUFFER_IDX(coord, cpu_sizes.data);
   const ivec4 buf_indices =
-      base_index + ivec4(0, 1, 2, 3) * (gpu_sizes.data.x * gpu_sizes.data.y);
+      base_index + ivec4(0, 1, 2, 3) * STRIDE_${PACKING}(cpu_sizes.data);
 
-  if (coord.z < cpu_sizes.data.z) {
+  const int packed_dim_size = PACKED_DIM_${PACKING}(cpu_sizes.data);
+  int packed_coord = PACKED_DIM_${PACKING}(coord);
+
+  if (packed_coord < packed_dim_size) {
     buffer_out.data[buf_indices.x] = intex.x;
   }
-  if (coord.z + 1 < cpu_sizes.data.z) {
+  if (packed_coord + 1 < packed_dim_size) {
     buffer_out.data[buf_indices.y] = intex.y;
   }
-  if (coord.z + 2 < cpu_sizes.data.z) {
+  if (packed_coord + 2 < packed_dim_size) {
     buffer_out.data[buf_indices.z] = intex.z;
   }
-  if (coord.z + 3 < cpu_sizes.data.z) {
+  if (packed_coord + 3 < packed_dim_size) {
     buffer_out.data[buf_indices.w] = intex.w;
   }
 }

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

@@ -10,14 +10,21 @@ image_to_nchw:
     DTYPE: float
     PACKING: CHANNELS_PACKED
   generate_variant_forall:
+    PACKING:
+      - VALUE: CHANNELS_PACKED
+        SUFFIX: C_packed
+      - VALUE: WIDTH_PACKED
+        SUFFIX: W_packed
+      - VALUE: HEIGHT_PACKED
+        SUFFIX: H_packed
     DTYPE:
-      - VALUE: "half"
-        SUFFIX: "half"
-      - VALUE: "float"
-        SUFFIX: "float"
-      - VALUE: "int"
-        SUFFIX: "int"
+      - VALUE: half
+        SUFFIX: half
+      - VALUE: float
+        SUFFIX: float
+      - VALUE: int
+        SUFFIX: int
   shader_variants:
-    - NAME: image3d_to_nchw_C_packed
-    - NAME: image2d_to_nchw_C_packed
+    - NAME: image3d_to_nchw
+    - NAME: image2d_to_nchw
       NDIM: 2

backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h

@@ -6,12 +6,30 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#define PACKED_DIM_CHANNELS_PACKED(vec) vec.z
+
+#define PACKED_DIM_WIDTH_PACKED(vec) vec.x
+
+#define PACKED_DIM_HEIGHT_PACKED(vec) vec.y
+
 #define POS_TO_COORD_CHANNELS_PACKED(pos, sizes) \
   ivec4(pos.x, pos.y, (pos.z * 4) % sizes.z, (pos.z * 4) / sizes.z)
 
+#define POS_TO_COORD_WIDTH_PACKED(pos, sizes) \
+  ivec4((pos.x * 4), pos.y, pos.z % sizes.z, pos.z / sizes.z)
+
+#define POS_TO_COORD_HEIGHT_PACKED(pos, sizes) \
+  ivec4(pos.x, (pos.y * 4), pos.z % sizes.z, pos.z / sizes.z)
+
 #define COORD_TO_POS_CHANNELS_PACKED(coord, sizes) \
   ivec3(coord.x, coord.y, (coord.z + coord.w * sizes.z) / 4)
 
 #define COORD_TO_BUFFER_IDX(coord, sizes)                  \
   coord.x + coord.y* sizes.x + coord.z* sizes.y* sizes.x + \
       coord.w* sizes.z* sizes.y* sizes.x;
+
+#define STRIDE_CHANNELS_PACKED(vec) (vec.x * vec.y)
+
+#define STRIDE_WIDTH_PACKED(vec) (1)
+
+#define STRIDE_HEIGHT_PACKED(vec) (vec.x)

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

@@ -42,7 +42,7 @@ void main() {
 
   const int base_index = COORD_TO_BUFFER_IDX(coord, cpu_sizes.data);
   const ivec4 buf_indices =
-      base_index + ivec4(0, 1, 2, 3) * (gpu_sizes.data.x * gpu_sizes.data.y);
+      base_index + ivec4(0, 1, 2, 3) * STRIDE_${PACKING}(cpu_sizes.data);
 
   ${T[DTYPE]} val_x = buffer_in.data[buf_indices.x];
   ${T[DTYPE]} val_y = buffer_in.data[buf_indices.y];
@@ -51,10 +51,13 @@ void main() {
 
   ${VEC4_T[DTYPE]} texel = ${VEC4_T[DTYPE]}(val_x, val_y, val_z, val_w);
 
-  if (coord.z + 3 >= cpu_sizes.data.z) {
-    ivec4 c_ind = ivec4(coord.z) + ivec4(0, 1, 2, 3);
-    ${VEC4_T[DTYPE]} valid_c = ${VEC4_T[DTYPE]}(lessThan(c_ind, ivec4(cpu_sizes.data.z)));
-    texel = texel * valid_c;
+  const int packed_dim_size = PACKED_DIM_${PACKING}(cpu_sizes.data);
+  int packed_coord = PACKED_DIM_${PACKING}(coord);
+
+  if (packed_coord + 3 >= packed_dim_size) {
+    ivec4 packed_ind = ivec4(packed_coord) + ivec4(0, 1, 2, 3);
+    ${VEC4_T[DTYPE]} valid_idx = ${VEC4_T[DTYPE]}(lessThan(packed_ind, ivec4(packed_dim_size)));
+    texel = texel * valid_idx;
   }
 
   imageStore(image_out, ${GET_POS[NDIM]("pos")}, texel);

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

@@ -10,14 +10,21 @@ nchw_to_image:
     DTYPE: float
     PACKING: CHANNELS_PACKED
   generate_variant_forall:
+    PACKING:
+      - VALUE: CHANNELS_PACKED
+        SUFFIX: C_packed
+      - VALUE: WIDTH_PACKED
+        SUFFIX: W_packed
+      - VALUE: HEIGHT_PACKED
+        SUFFIX: H_packed
     DTYPE:
-      - VALUE: "half"
-        SUFFIX: "half"
-      - VALUE: "float"
-        SUFFIX: "float"
-      - VALUE: "int"
-        SUFFIX: "int"
+      - VALUE: half
+        SUFFIX: half
+      - VALUE: float
+        SUFFIX: float
+      - VALUE: int
+        SUFFIX: int
   shader_variants:
-    - NAME: nchw_to_image3d_C_packed
-    - NAME: nchw_to_image2d_C_packed
+    - NAME: nchw_to_image3d
+    - NAME: nchw_to_image2d
       NDIM: 2

backends/vulkan/test/glsl/all_shaders.yaml

@@ -33,30 +33,23 @@ fill_texture__test:
   shader_variants:
     - NAME: fill_texture__test
 
-image_to_nchw__test:
+idx_fill_texture:
   parameter_names_with_default_values:
-    NDIM: 3
     DTYPE: float
-    PACKING: CHANNELS_PACKED
-  generate_variant_forall:
-    DTYPE:
-      - VALUE: "half"
-        SUFFIX: "half"
-      - VALUE: "float"
-        SUFFIX: "float"
-  shader_variants:
-    - NAME: image3d_to_nchw__test_C_packed
-
-nchw_to_image__test:
-  parameter_names_with_default_values:
     NDIM: 3
-    DTYPE: float
     PACKING: CHANNELS_PACKED
   generate_variant_forall:
+    PACKING:
+      - VALUE: "CHANNELS_PACKED"
+        SUFFIX: "C_packed"
+      - VALUE: "WIDTH_PACKED"
+        SUFFIX: "W_packed"
+      - VALUE: "HEIGHT_PACKED"
+        SUFFIX: "H_packed"
     DTYPE:
       - VALUE: "half"
         SUFFIX: "half"
       - VALUE: "float"
         SUFFIX: "float"
   shader_variants:
-    - NAME: nchw_to_image3d__test_C_packed
+    - NAME: idx_fill_texture

backends/vulkan/test/glsl/nchw_to_image__test.glsl renamed to backends/vulkan/test/glsl/idx_fill_texture.glsl

@@ -7,28 +7,21 @@
  */
 
 #version 450 core
-// clang-format off
+
 #define PRECISION ${PRECISION}
-// clang-format on
 
 #include "indexing_utils.h"
 
 layout(std430) buffer;
 
-// clang-format off
 layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} image_out;
-// clang-format on
-layout(set = 0, binding = 1) buffer  PRECISION restrict readonly Buffer {
-  ${T[DTYPE]} data[];
-}
-buffer_in;
 
-layout(set = 0, binding = 2) uniform PRECISION restrict GpuSizes {
+layout(set = 0, binding = 1) uniform PRECISION restrict GpuSizes {
   ivec4 data;
 }
 gpu_sizes;
 
-layout(set = 0, binding = 3) uniform PRECISION restrict CpuSizes {
+layout(set = 0, binding = 2) uniform PRECISION restrict CpuSizes {
   ivec4 data;
 }
 cpu_sizes;
@@ -45,20 +38,9 @@ void main() {
 
   const int base_index = COORD_TO_BUFFER_IDX(coord, cpu_sizes.data);
   const ivec4 buf_indices =
-      base_index + ivec4(0, 1, 2, 3) * (gpu_sizes.data.x * gpu_sizes.data.y);
-
-  ${T[DTYPE]} val_x = buffer_in.data[buf_indices.x];
-  ${T[DTYPE]} val_y = buffer_in.data[buf_indices.y];
-  ${T[DTYPE]} val_z = buffer_in.data[buf_indices.z];
-  ${T[DTYPE]} val_w = buffer_in.data[buf_indices.w];
+      base_index + ivec4(0, 1, 2, 3) * PLANE_SIZE_${PACKING}(gpu_sizes.data);
 
-  ${VEC4_T[DTYPE]} texel = ${VEC4_T[DTYPE]}(val_x, val_y, val_z, val_w);
-
-  if (coord.z + 3 >= cpu_sizes.data.z) {
-    ivec4 c_ind = ivec4(coord.z) + ivec4(0, 1, 2, 3);
-    vec4 valid_c = vec4(lessThan(c_ind, ivec4(cpu_sizes.data.z)));
-    texel = texel * valid_c;
-  }
+  ${VEC4_T[DTYPE]} texel = ${VEC4_T[DTYPE]}(buf_indices);
 
   imageStore(image_out, ${GET_POS[NDIM]("pos")}, texel);
 }

backends/vulkan/test/glsl/image_to_nchw__test.glsl

@@ -6,9 +6,27 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#define PACKED_DIM_CHANNELS_PACKED(vec) vec.z
+
+#define PACKED_DIM_WIDTH_PACKED(vec) vec.x
+
+#define PACKED_DIM_HEIGHT_PACKED(vec) vec.y
+
 #define POS_TO_COORD_CHANNELS_PACKED(pos, sizes) \
   ivec4(pos.x, pos.y, (pos.z * 4) % sizes.z, (pos.z * 4) / sizes.z)
 
+#define POS_TO_COORD_WIDTH_PACKED(pos, sizes) \
+  ivec4((pos.x * 4), pos.y, pos.z % sizes.z, pos.z / sizes.z)
+
+#define POS_TO_COORD_HEIGHT_PACKED(pos, sizes) \
+  ivec4(pos.x, (pos.y * 4), pos.z % sizes.z, pos.z / sizes.z)
+
 #define COORD_TO_BUFFER_IDX(coord, sizes)                  \
   coord.x + coord.y* sizes.x + coord.z* sizes.y* sizes.x + \
       coord.w* sizes.z* sizes.y* sizes.x;
+
+#define PLANE_SIZE_CHANNELS_PACKED(vec) (vec.x * vec.y)
+
+#define PLANE_SIZE_WIDTH_PACKED(vec) (1)
+
+#define PLANE_SIZE_HEIGHT_PACKED(vec) (vec.x)

backends/vulkan/test/glsl/indexing_utils.h

@@ -73,13 +73,9 @@ void record_nchw_to_image_op(
     api::VulkanBuffer& src_buffer,
     vTensor& v_dst) {
   api::PipelineBarrier pipeline_barrier{};
-  api::ShaderInfo compute_shader =
-      VK_KERNEL(nchw_to_image3d__test_C_packed_half);
-  if (v_dst.image().format() == VK_FORMAT_R32G32B32A32_SFLOAT) {
-    compute_shader = VK_KERNEL(nchw_to_image3d__test_C_packed_float);
-  }
+
   context->submit_compute_job(
-      compute_shader,
+      get_nchw_to_image_shader(v_dst),
       pipeline_barrier,
       v_dst.virtual_extents(),
       adaptive_work_group_size(v_dst.virtual_extents()),
@@ -97,14 +93,9 @@ void record_image_to_nchw_op(
     api::Context* const context,
     vTensor& v_src,
     api::VulkanBuffer& dst_buffer) {
-  api::ShaderInfo compute_shader =
-      VK_KERNEL(image3d_to_nchw__test_C_packed_half);
-  if (v_src.image().format() == VK_FORMAT_R32G32B32A32_SFLOAT) {
-    compute_shader = VK_KERNEL(image3d_to_nchw__test_C_packed_float);
-  }
   api::PipelineBarrier pipeline_barrier{};
   context->submit_compute_job(
-      compute_shader,
+      get_image_to_nchw_shader(v_src),
       pipeline_barrier,
       v_src.virtual_extents(),
       adaptive_work_group_size(v_src.virtual_extents()),
@@ -147,12 +138,6 @@ void execute_and_check_add(
     vTensor& c,
     float a_val,
     float b_val) {
-  // Add shader kernel
-  api::ShaderInfo kernel = VK_KERNEL(binary_add_nobroadcast__test_half);
-  if (c.image().format() == VK_FORMAT_R32G32B32A32_SFLOAT) {
-    kernel = VK_KERNEL(nchw_to_image3d__test_C_packed_float);
-  }
-
   // Fill input tensors
   fill_vtensor(a, a_val);
   fill_vtensor(b, b_val);
@@ -164,8 +149,8 @@ void execute_and_check_add(
   std::vector<float> data_out = extract_vtensor(c);
 
   // Check output
-  for (const auto& d : data_out) {
-    EXPECT_TRUE(d == (a_val + b_val));
+  for (size_t i = 0; i < data_out.size(); ++i) {
+    CHECK_VALUE(data_out, i, (a_val + b_val));
   }
 }