File tree Expand file tree Collapse file tree 1 file changed +9
-1
lines changed
backends/vulkan/runtime/graph/ops/impl Expand file tree Collapse file tree 1 file changed +9
-1
lines changed Original file line number Diff line number Diff line change @@ -51,6 +51,14 @@ void add_staging_to_tensor_node(
51
51
{}));
52
52
}
53
53
54
+ const std::string kBitw8PrefixStr = " bitw8_image_to_nchw_nobitw8buffer" ;
55
+
56
+ bool is_bitw8_shader (const vkapi::ShaderInfo& shader) {
57
+ const auto size = shader.kernel_name .size ();
58
+ const std::string& shader_prefix_str = shader.kernel_name .substr (0 , size);
59
+ return shader_prefix_str == kBitw8PrefixStr ;
60
+ }
61
+
54
62
void add_tensor_to_staging_node (
55
63
ComputeGraph& graph,
56
64
const ValueRef in_tensor,
@@ -80,7 +88,7 @@ void add_tensor_to_staging_node(
80
88
// output buffer. Therefore, the global work group size for this shader will
81
89
// be the number of elements in the output buffer divided by 4, as opposed to
82
90
// the extents of the input texture.
83
- if (shader. kernel_name . starts_with ( " bitw8_image_to_nchw_nobitw8buffer " )) {
91
+ if (is_bitw8_shader (shader )) {
84
92
uint32_t buffer_len = graph.get_staging (out_staging)->numel () / 4 ;
85
93
global_wg_size = {buffer_len, 1 , 1 };
86
94
ubos.append ({graph.numel_ubo (in_tensor)});
You can’t perform that action at this time.
0 commit comments