From f95bd50e402241c167727e55a991b44e08f9dbbb Mon Sep 17 00:00:00 2001 From: Vivek Trivedi <5340687+trivedivivek@users.noreply.github.com> Date: Mon, 24 Feb 2025 13:03:42 -0800 Subject: [PATCH] [ET-VK] Moving PushConstantData class implementation a separate h and cpp file. This diff moves the implementation of the `PushConstantData` class from the `DispatchNode` file to a separate `PushConstantData.h` and `PushConstantData.cpp` files. This is done to enable push constants for PrepackNodes in follow up changes. Differential Revision: [D70102031](https://our.internmc.facebook.com/intern/diff/D70102031/) [ghstack-poisoned] --- .../graph/containers/PushConstantData.cpp | 29 +++++++++ .../graph/containers/PushConstantData.h | 65 +++++++++++++++++++ .../vulkan/runtime/graph/ops/DispatchNode.cpp | 16 ----- .../vulkan/runtime/graph/ops/DispatchNode.h | 49 +------------- 4 files changed, 95 insertions(+), 64 deletions(-) create mode 100644 backends/vulkan/runtime/graph/containers/PushConstantData.cpp create mode 100644 backends/vulkan/runtime/graph/containers/PushConstantData.h diff --git a/backends/vulkan/runtime/graph/containers/PushConstantData.cpp b/backends/vulkan/runtime/graph/containers/PushConstantData.cpp new file mode 100644 index 00000000000..7999118443b --- /dev/null +++ b/backends/vulkan/runtime/graph/containers/PushConstantData.cpp @@ -0,0 +1,29 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +namespace vkcompute { + +uint32_t PushConstantDataInfo::write( + void* dst, + const uint32_t dst_offset, + const uint32_t max_dst_size) const { + if (tensorUniformData != nullptr) { + return tensorUniformData->write_attribute( + dst, dst_offset, max_dst_size, payload_.attr); + } + + VK_CHECK_COND( + (dst_offset + payload_.dataSize) <= max_dst_size, + "Attempting to write push constant data outside data boundary."); + memcpy((uint8_t*)dst + dst_offset, payload_.data, payload_.dataSize); + return payload_.dataSize; +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/containers/PushConstantData.h b/backends/vulkan/runtime/graph/containers/PushConstantData.h new file mode 100644 index 00000000000..39cde4722a7 --- /dev/null +++ b/backends/vulkan/runtime/graph/containers/PushConstantData.h @@ -0,0 +1,65 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace vkcompute { + +class ComputeGraph; + +constexpr uint32_t kMaxPushConstantSize = 128; +/* + * Represents a push constant data entry + * Which is either shared pointer to a tensor's uniform data with an attribute + * Or data with a maximum size of 16 bytes + */ +class PushConstantDataInfo { + std::shared_ptr tensorUniformData; + union Payload { + struct { + api::vTensor::Attribute attr; + }; + struct { + uint8_t data[16]; + uint32_t dataSize; + }; + }; + + Payload payload_; + + public: + explicit PushConstantDataInfo( + const std::shared_ptr& tensorUniformData, + api::vTensor::Attribute attr) + : tensorUniformData(tensorUniformData) { + payload_.attr = attr; + } + + explicit PushConstantDataInfo( + const void* data, + uint32_t dataLen, + uint32_t pushConstantLen = 0) + : tensorUniformData(nullptr) { + VK_CHECK_COND( + dataLen <= 16, "Single push constant data size must be <= 16 bytes"); + payload_.dataSize = pushConstantLen ? pushConstantLen : dataLen; + memcpy(payload_.data, data, dataLen); + } + + /* + * Function writes push constant data to the destination buffer + */ + uint32_t write( + void* dst, + const uint32_t dst_offset, + const uint32_t max_dst_size) const; +}; + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/DispatchNode.cpp b/backends/vulkan/runtime/graph/ops/DispatchNode.cpp index 63b8798f2c1..6730d851483 100644 --- a/backends/vulkan/runtime/graph/ops/DispatchNode.cpp +++ b/backends/vulkan/runtime/graph/ops/DispatchNode.cpp @@ -14,22 +14,6 @@ namespace vkcompute { -uint32_t PushConstantDataInfo::write( - void* dst, - const uint32_t dst_offset, - const uint32_t max_dst_size) const { - if (tensorUniformData != nullptr) { - return tensorUniformData->write_attribute( - dst, dst_offset, max_dst_size, payload_.attr); - } - - VK_CHECK_COND( - (dst_offset + payload_.dataSize) <= max_dst_size, - "Attempting to write push constant data outside data boundary."); - memcpy((uint8_t*)dst + dst_offset, payload_.data, payload_.dataSize); - return payload_.dataSize; -} - DispatchNode::DispatchNode( ComputeGraph& graph, const vkapi::ShaderInfo& shader, diff --git a/backends/vulkan/runtime/graph/ops/DispatchNode.h b/backends/vulkan/runtime/graph/ops/DispatchNode.h index 4661b5bf9cf..e3794e9a9e4 100644 --- a/backends/vulkan/runtime/graph/ops/DispatchNode.h +++ b/backends/vulkan/runtime/graph/ops/DispatchNode.h @@ -10,6 +10,7 @@ #include +#include #include #include @@ -18,54 +19,6 @@ namespace vkcompute { class ComputeGraph; -constexpr uint32_t kMaxPushConstantSize = 128; -/* - * Represents a push constant data entry - * Which is either shared pointer to a tensor's uniform data with an attribute - * Or data with a maximum size of 16 bytes - */ -class PushConstantDataInfo { - std::shared_ptr tensorUniformData; - union Payload { - struct { - api::vTensor::Attribute attr; - }; - struct { - uint8_t data[16]; - uint32_t dataSize; - }; - }; - - Payload payload_; - - public: - explicit PushConstantDataInfo( - const std::shared_ptr& tensorUniformData, - api::vTensor::Attribute attr) - : tensorUniformData(tensorUniformData) { - payload_.attr = attr; - } - - explicit PushConstantDataInfo( - const void* data, - uint32_t dataLen, - uint32_t pushConstantLen = 0) - : tensorUniformData(nullptr) { - VK_CHECK_COND( - dataLen <= 16, "Single push constant data size must be <= 16 bytes"); - payload_.dataSize = pushConstantLen ? pushConstantLen : dataLen; - memcpy(payload_.data, data, dataLen); - } - - /* - * Function writes push constant data to the destination buffer - */ - uint32_t write( - void* dst, - const uint32_t dst_offset, - const uint32_t max_dst_size) const; -}; - /* * Represents a single shader execution op in a ML model. */