Skip to content

Commit c3baa77

Browse files
trivedivivektrivedivivek
committed
[ET-VK] Adding push constant and ubo verison of select and slice ops to improve memory and performance.
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. Differential Revision: [D78095262](https://our.internmc.facebook.com/intern/diff/D78095262/) [ghstack-poisoned]
1 parent 4cd537b commit c3baa77

File tree

5 files changed

+84
-31
lines changed

5 files changed

+84
-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: 42 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -40,34 +40,52 @@ 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(
85+
PushConstantDataInfo(&transfer_params, sizeof(transfer_params)));
86+
} else {
87+
push_constants.emplace_back(PushConstantDataInfo(
88+
&transfer_params.dim, sizeof(transfer_params.dim)));
7189
}
7290

7391
vkapi::SpecVarList spec_vars = {
@@ -82,6 +100,9 @@ void add_transfer_copy_node(
82100
} else { // TransferType::SLICE
83101
kernel_name = "slice";
84102
}
103+
if (!param_is_scalar) {
104+
kernel_name += "_ubo";
105+
}
85106
add_storage_type_suffix(kernel_name, graph.storage_type_of(out));
86107
add_dtype_suffix(kernel_name, graph.dtype_of(out));
87108

0 commit comments

Comments
 (0)