Skip to content

Commit 6f68455

Browse files
trivedivivektrivedivivek
committed
[ET-VK] Adding push constant and ubo verison of select and slice ops to improve memory and performance.
Pull Request resolved: #12358 Adding push constant and ubo verison of select and slice ops to improve memory and performance. * Updated `transfer_buffer.yaml` and `transfer_texture.yaml` to include `UBO_PARAMS` parameter and generate variants for `select` and `slice` ops with UBO parameters. * Updated `transfer.glsl` to generate ubo and push constant versions of `select` and `slice` ops with UBO parameters. ghstack-source-id: 295655470 Differential Revision: [D78095262](https://our.internmc.facebook.com/intern/diff/D78095262/)
1 parent b1cd9eb commit 6f68455

File tree

5 files changed

+83
-31
lines changed

5 files changed

+83
-31
lines changed

‎backends/vulkan/runtime/graph/ops/glsl/transfer_buffer.glsl

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#version 450 core
1010

1111
#define PRECISION ${PRECISION}
12+
#define UBO_PARAMS ${UBO_PARAMS}
1213

1314
#define VEC4_T ${texel_type(DTYPE)}
1415
#define T ${buffer_scalar_type(DTYPE)}
@@ -22,19 +23,27 @@ layout(std430) buffer;
2223
${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")}
2324
${layout_declare_tensor(B, "r", "t_in", DTYPE, "buffer")}
2425

25-
$if OP_NAME == "slice":
26-
${layout_declare_ubo(B, "int", "start")}
27-
${layout_declare_ubo(B, "int", "step")}
26+
$if UBO_PARAMS:
27+
$if OP_NAME == "slice":
28+
${layout_declare_ubo(B, "int", "start")}
29+
${layout_declare_ubo(B, "int", "step")}
2830

29-
$if OP_NAME == "select":
30-
${layout_declare_ubo(B, "int", "index")}
31+
$if OP_NAME == "select":
32+
${layout_declare_ubo(B, "int", "index")}
3133

3234
layout(push_constant) uniform restrict Block {
3335
ivec4 in_sizes;
3436
ivec4 out_strides;
3537
ivec4 in_strides;
3638
int out_numel;
3739
int selected_dim;
40+
$if not UBO_PARAMS:
41+
$if OP_NAME == "slice":
42+
int start;
43+
int step;
44+
45+
$if OP_NAME == "select":
46+
int index;
3847
};
3948

4049
${layout_declare_spec_const(C, "int", "out_layout", "DEFAULT_LAYOUT")}

‎backends/vulkan/runtime/graph/ops/glsl/transfer_buffer.yaml

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ transfer_buffer:
22
parameter_names_with_default_values:
33
DTYPE: float
44
OP_NAME: select
5+
UBO_PARAMS: False
56
generate_variant_forall:
67
DTYPE:
78
- VALUE: half
@@ -11,3 +12,9 @@ transfer_buffer:
1112
OP_NAME: select
1213
- NAME: slice_buffer
1314
OP_NAME: slice
15+
- NAME: select_ubo_buffer
16+
OP_NAME: select
17+
UBO_PARAMS: True
18+
- NAME: slice_ubo_buffer
19+
OP_NAME: slice
20+
UBO_PARAMS: True

‎backends/vulkan/runtime/graph/ops/glsl/transfer_texture.glsl

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#version 450 core
1010

1111
#define PRECISION ${PRECISION}
12+
#define UBO_PARAMS ${UBO_PARAMS}
1213

1314
#define VEC4_T ${texel_type(DTYPE)}
1415
#define T ${buffer_scalar_type(DTYPE)}
@@ -23,17 +24,25 @@ layout(std430) buffer;
2324
${layout_declare_tensor(B, "w", "t_out", DTYPE, "texture3d")}
2425
${layout_declare_tensor(B, "r", "t_in", DTYPE, "texture3d")}
2526

26-
$if OP_NAME == "slice":
27-
${layout_declare_ubo(B, "int", "start")}
28-
${layout_declare_ubo(B, "int", "step")}
27+
$if UBO_PARAMS:
28+
$if OP_NAME == "slice":
29+
${layout_declare_ubo(B, "int", "start")}
30+
${layout_declare_ubo(B, "int", "step")}
2931

30-
$if OP_NAME == "select":
31-
${layout_declare_ubo(B, "int", "index")}
32+
$if OP_NAME == "select":
33+
${layout_declare_ubo(B, "int", "index")}
3234

3335
layout(push_constant) uniform restrict Block {
3436
ivec4 out_sizes;
3537
ivec4 in_sizes;
3638
int selected_dim;
39+
$if not UBO_PARAMS:
40+
$if OP_NAME == "slice":
41+
int start;
42+
int step;
43+
44+
$if OP_NAME == "select":
45+
int index;
3746
};
3847

3948
${layout_declare_spec_const(C, "int", "out_layout", "DEFAULT_LAYOUT")}

‎backends/vulkan/runtime/graph/ops/glsl/transfer_texture.yaml

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ transfer_texture:
22
parameter_names_with_default_values:
33
DTYPE: float
44
OP_NAME: select
5+
UBO_PARAMS: False
56
generate_variant_forall:
67
DTYPE:
78
- VALUE: half
@@ -11,3 +12,9 @@ transfer_texture:
1112
OP_NAME: select
1213
- NAME: slice_texture3d
1314
OP_NAME: slice
15+
- NAME: select_ubo_texture3d
16+
OP_NAME: select
17+
UBO_PARAMS: True
18+
- NAME: slice_ubo_texture3d
19+
OP_NAME: slice
20+
UBO_PARAMS: True

‎backends/vulkan/runtime/graph/ops/impl/Transfer.cpp

Lines changed: 41 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -40,34 +40,51 @@ void add_transfer_copy_node(
4040

4141
int64_t dim_whcn = nchw_dim_to_whcn_dim(dim, ndim);
4242

43+
struct TransferParams {
44+
int32_t dim;
45+
int32_t index_or_start_ref;
46+
int32_t step_ref;
47+
} transfer_params{static_cast<int32_t>(dim_whcn), 0, 0};
48+
49+
const bool param_is_scalar = graph.is_scalar(index_or_start_ref) &&
50+
(transfer_type == TransferType::SELECT || graph.is_scalar(step_ref));
51+
4352
vkapi::ParamsBindList param_buffers;
44-
if (transfer_type == TransferType::SELECT) {
45-
param_buffers = {
46-
graph.get_or_create_int_param_buffer(index_or_start_ref, 0)};
47-
} else { // TransferType::SLICE
48-
param_buffers = {
49-
graph.get_or_create_int_param_buffer(index_or_start_ref, 0),
50-
graph.get_or_create_int_param_buffer(step_ref, 1)};
53+
if (!param_is_scalar) {
54+
if (transfer_type == TransferType::SELECT) {
55+
param_buffers = {
56+
graph.get_or_create_int_param_buffer(index_or_start_ref, 0)};
57+
} else { // TransferType::SLICE
58+
param_buffers = {
59+
graph.get_or_create_int_param_buffer(index_or_start_ref, 0),
60+
graph.get_or_create_int_param_buffer(step_ref, 1)};
61+
}
62+
} else {
63+
transfer_params.index_or_start_ref =
64+
graph.get_or_create_int(index_or_start_ref, 0);
65+
if (transfer_type != TransferType::SELECT) {
66+
transfer_params.step_ref = graph.get_or_create_int(step_ref, 1);
67+
}
5168
}
5269

53-
const struct TransferParams {
54-
const int32_t dim;
55-
} transfer_params{static_cast<int32_t>(dim_whcn)};
56-
5770
std::vector<PushConstantDataInfo> push_constants;
71+
push_constants.reserve(graph.is_buffer_storage(out) ? 5 : 3);
5872

5973
if (graph.is_buffer_storage(out)) {
60-
push_constants = {
61-
graph.sizes_pc_of(in),
62-
graph.strides_pc_of(out),
63-
graph.strides_pc_of(in),
64-
graph.numel_pc_of(out),
65-
PushConstantDataInfo(&transfer_params, sizeof(transfer_params))};
74+
push_constants.emplace_back(graph.sizes_pc_of(in));
75+
push_constants.emplace_back(graph.strides_pc_of(out));
76+
push_constants.emplace_back(graph.strides_pc_of(in));
77+
push_constants.emplace_back(graph.numel_pc_of(out));
6678
} else {
67-
push_constants = {
68-
graph.sizes_pc_of(out),
69-
graph.sizes_pc_of(in),
70-
PushConstantDataInfo(&transfer_params, sizeof(transfer_params))};
79+
push_constants.emplace_back(graph.sizes_pc_of(out));
80+
push_constants.emplace_back(graph.sizes_pc_of(in));
81+
}
82+
83+
if (param_is_scalar) {
84+
push_constants.emplace_back(&transfer_params, sizeof(transfer_params));
85+
} else {
86+
push_constants.emplace_back(
87+
&transfer_params.dim, sizeof(transfer_params.dim));
7188
}
7289

7390
vkapi::SpecVarList spec_vars = {
@@ -82,6 +99,9 @@ void add_transfer_copy_node(
8299
} else { // TransferType::SLICE
83100
kernel_name = "slice";
84101
}
102+
if (!param_is_scalar) {
103+
kernel_name += "_ubo";
104+
}
85105
add_storage_type_suffix(kernel_name, graph.storage_type_of(out));
86106
add_dtype_suffix(kernel_name, graph.dtype_of(out));
87107

0 commit comments

Comments
 (0)