Skip to content

Commit ce461ed

Browse files
committed
[ET-VK] Using push constants for convolution prepack dispatches.
Pull Request resolved: #11306 This diff switches the `conv2d_prepack_weights`, `conv2d_dw_prepack_weights`, and `conv_transpose2d_prepack_weights` shaders to use push constants instead of uniform buffers for the `sizes` and `original_sizes` variables. ghstack-source-id: 288155883 @exported-using-ghexport Differential Revision: [D74523774](https://our.internmc.facebook.com/intern/diff/D74523774/)
1 parent 0ef4ac4 commit ce461ed

File tree

4 files changed

+8
-16
lines changed

4 files changed

+8
-16
lines changed

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

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,8 @@ layout(set = 0, binding = 1) buffer PRECISION restrict readonly Buffer {
2626
BUF_T buffer_in[];
2727
};
2828

29-
layout(set = 0, binding = 2) uniform PRECISION restrict Sizes {
29+
layout(push_constant) uniform PRECISION restrict Block {
3030
ivec4 sizes;
31-
};
32-
33-
layout(set = 0, binding = 3) uniform PRECISION restrict OriginalSizes {
3431
ivec4 original_sizes;
3532
};
3633

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

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,8 @@ layout(set = 0, binding = 1) buffer PRECISION restrict readonly Buffer {
2626
BUF_T buffer_in[];
2727
};
2828

29-
layout(set = 0, binding = 2) uniform PRECISION restrict Sizes {
29+
layout(push_constant) uniform PRECISION restrict Block {
3030
ivec4 sizes;
31-
};
32-
33-
layout(set = 0, binding = 3) uniform PRECISION restrict OriginalSizes {
3431
ivec4 original_sizes;
3532
};
3633

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

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,8 @@ layout(set = 0, binding = 1) buffer PRECISION restrict readonly Buffer {
2626
BUF_T buffer_in[];
2727
};
2828

29-
layout(set = 0, binding = 2) uniform PRECISION restrict Sizes {
29+
layout(push_constant) uniform PRECISION restrict Block {
3030
ivec4 sizes;
31-
};
32-
33-
layout(set = 0, binding = 3) uniform PRECISION restrict OriginalSizes {
3431
ivec4 original_sizes;
3532
};
3633

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -211,18 +211,19 @@ ValueRef prepack_weights(
211211
vkapi::ShaderInfo shader =
212212
get_conv2d_shader(graph, *t, /*prepack_weights = */ true, method, vref);
213213

214+
const auto sizes_pc = utils::make_ivec4(original_sizes, /*reverse = */ true);
214215
graph.prepack_nodes().emplace_back(new PrepackNode(
215216
graph,
216217
shader,
217218
graph.create_global_wg_size(v),
218219
graph.create_local_wg_size(v),
219220
vref,
220221
v,
221-
{t->sizes_ubo(),
222-
graph.create_params_buffer(
223-
utils::make_ivec4(original_sizes, /*reverse = */ true))},
222+
{},
224223
// Specialization constants
225-
{SV(t->packed_dim())}));
224+
{SV(t->packed_dim())},
225+
{graph.sizes_pc_of(v),
226+
PushConstantDataInfo(&sizes_pc, sizeof(sizes_pc))}));
226227

227228
return v;
228229
}

0 commit comments

Comments
 (0)