From 321e0a65b6e9844cab6a609701fb58f89a9b55d3 Mon Sep 17 00:00:00 2001 From: Vivek Trivedi <5340687+trivedivivek@users.noreply.github.com> Date: Fri, 6 Dec 2024 08:53:18 -0800 Subject: [PATCH] [ET-VK] Adding convenience functions in Compute graph to get PushConstantDataInfo for various attributes of a tensor. This diff adds convenience functions in the Compute graph to get PushConstantDataInfo for various attributes of a tensor. Differential Revision: [D66853502](https://our.internmc.facebook.com/intern/diff/D66853502/) [ghstack-poisoned] --- backends/vulkan/runtime/graph/ComputeGraph.h | 23 ++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/backends/vulkan/runtime/graph/ComputeGraph.h b/backends/vulkan/runtime/graph/ComputeGraph.h index cb958cefea..201278ac61 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.h +++ b/backends/vulkan/runtime/graph/ComputeGraph.h @@ -20,6 +20,7 @@ #include #include +#include #include #include @@ -350,6 +351,28 @@ class ComputeGraph final { return values_.at(idx).toTensor().logical_limits_ubo(); } + inline PushConstantDataInfo sizes_pc_of(const ValueRef idx) const { + return PushConstantDataInfo( + values_.at(idx).toConstTensor().get_uniform_data(), api::kTensorSizes); + } + + inline PushConstantDataInfo strides_pc_of(const ValueRef idx) const { + return PushConstantDataInfo( + values_.at(idx).toConstTensor().get_uniform_data(), + api::kTensorStrides); + } + + inline PushConstantDataInfo logical_limits_pc_of(const ValueRef idx) const { + return PushConstantDataInfo( + values_.at(idx).toConstTensor().get_uniform_data(), + api::kTensorLogicalLimits); + } + + inline PushConstantDataInfo numel_pc_of(const ValueRef idx) const { + return PushConstantDataInfo( + values_.at(idx).toConstTensor().get_uniform_data(), api::kTensorNumel); + } + // // Scalar Value Extraction //