xref: /aosp_15_r20/external/executorch/backends/vulkan/runtime/graph/ops/BlitNode.cpp (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1 /*
2  * Copyright (c) Meta Platforms, Inc. and affiliates.
3  * All rights reserved.
4  *
5  * This source code is licensed under the BSD-style license found in the
6  * LICENSE file in the root directory of this source tree.
7  */
8 
9 #include <executorch/backends/vulkan/runtime/graph/ops/BlitNode.h>
10 
11 #include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
12 
13 namespace vkcompute {
14 
BlitNode(ComputeGraph & graph,ValueRef src,ValueRef dst,const ResizeFunction & resize_fn,const std::vector<ValueRef> & resize_args)15 BlitNode::BlitNode(
16     ComputeGraph& graph,
17     ValueRef src,
18     ValueRef dst,
19     // const vkapi::ScalarType& dtype,
20     const ResizeFunction& resize_fn,
21     const std::vector<ValueRef>& resize_args)
22     : ExecuteNode(resize_fn, resize_args, {}, "Blit Node"),
23       src_(src),
24       dst_(dst) {
25   (void)graph;
26 }
27 
encode(ComputeGraph * graph)28 void BlitNode::encode(ComputeGraph* graph) {
29   auto src_tensor = graph->get_tensor(src_);
30   auto dst_tensor = graph->get_tensor(dst_);
31   VK_CHECK_COND(
32       src_tensor->storage_type() != utils::kBuffer &&
33           dst_tensor->storage_type() != utils::kBuffer,
34       "BlitNode: Only texture backed tensors are supported.");
35 
36   api::Context* const context = graph->context();
37   vkapi::PipelineBarrier pipeline_barrier{};
38 
39   std::unique_lock<std::mutex> cmd_lock = context->dispatch_lock();
40 
41   // Hack to get timing data for non shader op
42   std::string kernel_name("Blit_");
43   kernel_name.reserve(32);
44   kernel_name += vkapi::to_string(src_tensor->dtype());
45   kernel_name += "_to_";
46   kernel_name += vkapi::to_string(dst_tensor->dtype());
47 
48   context->report_shader_dispatch_start(
49       kernel_name, utils::uvec3(), utils::uvec3(), node_id_);
50 
51   context->register_blit(
52       pipeline_barrier,
53       src_tensor->image(
54           pipeline_barrier, vkapi::PipelineStage::TRANSFER, vkapi::kRead),
55       dst_tensor->image(
56           pipeline_barrier, vkapi::PipelineStage::TRANSFER, vkapi::kWrite));
57 
58   context->report_shader_dispatch_end();
59 }
60 
61 } // namespace vkcompute
62