Skip to content

Commit ec04231

Browse files
committed
Update on "[ET-VK] Clean up shader library and introduce some new conventions"
## Context This changeset introduces some fairly mechnical improvements to the Vulkan compute graph shader library in order to introduce some new conventions. **Note that backwards compatibility with existing shader authoring methods is preserved**. ### Only List `VALUE` in the `.yaml` files Previously, to generate variants for a combination of vales, the YAML file will contain ``` PACKING: - VALUE: CHANNELS_PACKED SUFFIX: C_packed - VALUE: WIDTH_PACKED SUFFIX: W_packed - VALUE: HEIGHT_PACKED SUFFIX: H_packed ``` however, the shader code generation script will use the `VALUE` as the `SUFFIX` if no `SUFFIX` is provided. Therefore, only the below is needed: ``` PACKING: - VALUE: C_packed - VALUE: W_packed - VALUE: H_packed ``` ### Change indexing utility macros to lowercase Indexing utility macros have been changed to lowercase, and the packing identifiers have been changed due to the change in YAML files. The change to lowercase is to make calls to the macro read more like functions (and indeed they are typically used as functions) in order to help make the code more readable. ``` POS_TO_COORD_${PACKING} -> pos_to_coord_${PACKING} ``` ### Use convention of defining macros in order to reduce Python code blocks usage Previously python code blocks were used in the GLSL code itself in order to vary the shader between different settings. However, usage of Python code blocks negatively impact code readability. Therefore, this diff seeks to introduce a convention of defining macros near the top of the shader to reduce the usage of Python code blocks, i.e. ``` #define pos_to_coord pos_to_coord_${PACKING} #define get_packed_dim get_packed_dim_${PACKING} #define get_packed_stride get_packed_stride_${PACKING} ``` ### Improve GLSL type definitions Previously, the following Python code blocks were used to determine appropriate vectorized and scalar types: ``` ${VEC4_T[DTYPE}} texel = ... ${T[DTYPE]} scalar = ... ``` This changeset replaces that with: ``` #define BUF_T ${buffer_scalar_type(DTYPE)} #define VEC4_T ${texel_type(DTYPE)} #define SCALAR_T ${texel_component_type(DTYPE)} layout(set = 0, binding = 1) buffer PRECISION restrict readonly Buffer { BUF_T data[]; } buffer_in; VEC4_T texel = ... SCALAR_T scalar = ... ``` The main differences are as such: * `buffer_scalar_type()` produces the same result as `T[DTYPE]` * `texel_type()` is not determined from a mapping with `DTYPE`, but is determined indirectly based on the image format that is associated with the `DTYPE`. * `texel_component_type()` is based on the result of `texel_type(DTYPE)` Essentially, the mapping is more in-line with what happens in code. The reason for this change is to enable FP16 support and is a bit complicated. Basically, we need a way to distinguish the scalar type used for buffer storage, vs the scalar type used to store a component of a vec4 type (hence `BUF_T` vs `SCALAR_T`). The reason this is required is that to support half-precision tensors, the buffer representation will use a 16-bit float type but textures will still extract to `vec4` (i.e. 4x34bit floats). Differential Revision: [D56082461](https://our.internmc.facebook.com/intern/diff/D56082461/) [ghstack-poisoned]
2 parents 8a6ae21 + 9243f74 commit ec04231

13 files changed

+16
-75
lines changed

backends/vulkan/runtime/api/gen_vulkan_spv.py

Lines changed: 0 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -77,24 +77,6 @@
7777
"int8": "rgba8i",
7878
"uint8": "rgba8ui",
7979
},
80-
# Kept for backwards compatibility
81-
# TODO(ssjia): remove when no more shaders use these
82-
"VEC4_T": {
83-
"float": "vec4",
84-
"half": "vec4",
85-
"int": "ivec4",
86-
"uint": "uvec4",
87-
"int8": "vec4",
88-
"uint8": "uvec4",
89-
},
90-
"T": {
91-
"float": "float",
92-
"half": "float",
93-
"int": "int",
94-
"uint": "uint",
95-
"int8": "int",
96-
"uint8": "uint8",
97-
},
9880
}
9981

10082

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,13 @@
1010

1111
#define PRECISION ${PRECISION}
1212

13-
#define op(X, Y, A) ${OPERATOR}
14-
1513
#define VEC4_T ${texel_type(DTYPE)}
1614

1715
#define to_tensor_idx to_tensor_idx_${PACKING}
1816
#define to_texture_pos to_texture_pos_${PACKING}
1917

18+
#define op(X, Y, A) ${OPERATOR}
19+
2020
#include "broadcasting_utils.h"
2121
#include "indexing_utils.h"
2222

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#version 450 core
1010

1111
#define PRECISION ${PRECISION}
12+
1213
#define VEC4_T ${texel_type(DTYPE)}
1314

1415
#include "indexing_utils.h"

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#version 450 core
1010

1111
#define PRECISION ${PRECISION}
12+
1213
#define VEC4_T ${texel_type(DTYPE)}
1314

1415
#include "indexing_utils.h"

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#version 450 core
1010

1111
#define PRECISION ${PRECISION}
12+
1213
#define VEC4_T ${texel_type(DTYPE)}
1314

1415
#include "indexing_utils.h"

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#version 450 core
1010

1111
#define PRECISION ${PRECISION}
12+
1213
#define VEC4_T ${texel_type(DTYPE)}
1314

1415
#include "indexing_utils.h"

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#version 450 core
1010

1111
#define PRECISION ${PRECISION}
12+
1213
#define VEC4_T ${texel_type(DTYPE)}
1314

1415
layout(std430) buffer;

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

Lines changed: 0 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -48,45 +48,3 @@
4848
plane * \
4949
((1 - y) * ((cur % (x * y * plane)) / (y * plane)) + \
5050
(x - 1) * ((cur % (y * plane)) / plane))
51-
52-
// Kept for backwards compatibility
53-
// TODO(ssjia): remove once there are no shaders that use these macros
54-
55-
#define DIVUP4(x) ((x + 3) / 4)
56-
57-
#define PACKED_DIM_CHANNELS_PACKED(vec) vec.z
58-
59-
#define PACKED_DIM_WIDTH_PACKED(vec) vec.x
60-
61-
#define PACKED_DIM_HEIGHT_PACKED(vec) vec.y
62-
63-
#define POS_TO_COORD_CHANNELS_PACKED(pos, sizes) \
64-
ivec4(pos.x, pos.y, (pos.z * 4) % sizes.z, (pos.z * 4) / sizes.z)
65-
66-
#define POS_TO_COORD_WIDTH_PACKED(pos, sizes) \
67-
ivec4((pos.x * 4), pos.y, pos.z % sizes.z, pos.z / sizes.z)
68-
69-
#define POS_TO_COORD_HEIGHT_PACKED(pos, sizes) \
70-
ivec4(pos.x, (pos.y * 4), pos.z % sizes.z, pos.z / sizes.z)
71-
72-
#define COORD_TO_POS_CHANNELS_PACKED(idx, sizes) \
73-
ivec3(idx.x, idx.y, (idx.z + idx.w * sizes.z) / 4)
74-
75-
#define COORD_TO_POS_WIDTH_PACKED(idx, sizes) \
76-
ivec3(idx.x / 4, idx.y, (idx.z + idx.w * sizes.z))
77-
78-
#define COORD_TO_POS_HEIGHT_PACKED(idx, sizes) \
79-
ivec3(idx.x, idx.y / 4, (idx.z + idx.w * sizes.z))
80-
81-
#define COORD_TO_POS_CHANNELS_PACKED(idx, sizes) \
82-
ivec3(idx.x, idx.y, (idx.z + idx.w * sizes.z) / 4)
83-
84-
#define COORD_TO_BUFFER_IDX(idx, sizes) \
85-
idx.x + idx.y* sizes.x + idx.z* sizes.y* sizes.x + \
86-
idx.w* sizes.z* sizes.y* sizes.x;
87-
88-
#define STRIDE_CHANNELS_PACKED(vec) (vec.x * vec.y)
89-
90-
#define STRIDE_WIDTH_PACKED(vec) (1)
91-
92-
#define STRIDE_HEIGHT_PACKED(vec) (vec.x)

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,11 @@
99
#version 450 core
1010

1111
#define PRECISION ${PRECISION}
12-
#define op(X, A, B) ${OPERATOR}
1312

1413
#define VEC4_T ${texel_type(DTYPE)}
1514

15+
#define op(X, A, B) ${OPERATOR}
16+
1617
layout(std430) buffer;
1718

1819
layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} image_out;

backends/vulkan/test/glsl/binary_op_nobroadcast__test.glsl

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,10 @@
77
*/
88

99
#version 450 core
10-
// clang-format off
10+
1111
#define PRECISION ${PRECISION}
1212

13-
#define OP(X, Y) ${OPERATOR}
14-
// clang-format on
13+
#define op(X, Y) ${OPERATOR}
1514

1615
layout(std430) buffer;
1716

@@ -38,5 +37,5 @@ void main() {
3837
vec4 in_texel = texelFetch(image_in, pos, 0);
3938
vec4 other_texel = texelFetch(image_other, pos, 0);
4039

41-
imageStore(image_out, pos, OP(in_texel, other_texel));
40+
imageStore(image_out, pos, op(in_texel, other_texel));
4241
}

backends/vulkan/test/glsl/fill_texture__test.glsl

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,15 +7,12 @@
77
*/
88

99
#version 450 core
10+
1011
#define PRECISION ${PRECISION}
1112

1213
layout(std430) buffer;
1314

14-
/* Qualifiers: layout - storage - precision - memory */
15-
16-
// clang-format off
1715
layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} uOutput;
18-
// clang-format on
1916
layout(set = 0, binding = 1) uniform PRECISION restrict Block {
2017
ivec3 size;
2118
int fill;

backends/vulkan/test/glsl/idx_fill_texture.glsl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,12 @@
1010

1111
#define PRECISION ${PRECISION}
1212

13+
#define VEC4_T ${texel_type(DTYPE)}
14+
1315
#include "indexing_utils.h"
1416

1517
layout(std430) buffer;
1618

17-
#define VEC4_T ${texel_type(DTYPE)}
18-
1919
layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} image_out;
2020

2121
layout(set = 0, binding = 1) uniform PRECISION restrict GpuSizes {

backends/vulkan/test/glsl/test_shader.glsl

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,11 @@
77
*/
88

99
#version 450 core
10+
1011
#define PRECISION ${PRECISION}
1112

1213
layout(std430) buffer;
1314

14-
/* Qualifiers: layout - storage - precision - memory */
15-
1615
layout(set = 0, binding = 0, rgba32f) uniform PRECISION restrict writeonly image3D uOutput;
1716
layout(set = 0, binding = 1) uniform PRECISION sampler3D uInput;
1817
layout(set = 0, binding = 2) uniform PRECISION restrict Block {

0 commit comments

Comments
 (0)