1 /* 2 * Copyright (c) Qualcomm Innovation Center, Inc. 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 #pragma once 9 #include <executorch/backends/qualcomm/aot/wrappers/TensorWrapper.h> 10 #include <executorch/backends/qualcomm/runtime/SharedBuffer.h> 11 #include <executorch/backends/qualcomm/runtime/backends/QnnContextCommon.h> 12 #include <executorch/backends/qualcomm/runtime/backends/QnnImplementation.h> 13 #include <unordered_map> 14 #include "HTP/QnnHtpMem.h" 15 16 namespace executorch { 17 namespace backends { 18 namespace qnn { 19 20 class QnnMemManager { 21 public: QnnMemManager(const QnnImplementation & implementation,QnnContext * context)22 explicit QnnMemManager( 23 const QnnImplementation& implementation, 24 QnnContext* context) 25 : implementation_(implementation), context_(context) {} ~QnnMemManager()26 ~QnnMemManager() { 27 DeRegisterMem(); 28 } 29 30 executorch::runtime::Error RegisterIonMem( 31 const std::shared_ptr<TensorWrapper>& tensor_wrapper, 32 int32_t mem_fd, 33 void* mem_ptr); 34 35 executorch::runtime::Error RegisterCustomMem( 36 const std::shared_ptr<TensorWrapper>& tensor_wrapper, 37 int32_t mem_fd, 38 void* mem_ptr, 39 void* unaligned_custom_mem_base, 40 size_t total_custom_mem_size, 41 size_t tensor_offset); 42 43 // Pre-register custom mem handle from SharedBuffer. Bring forward the 44 // memHandle creating time from execution to initialization. 45 executorch::runtime::Error PreRegisterCustomMemHandle( 46 int32_t mem_fd, 47 void* unaligned_custom_mem_base, 48 size_t total_custom_mem_size, 49 size_t tensor_offset, 50 const CustomMemTensorInfo& info); 51 52 bool IsRegistered(Qnn_MemHandle_t handle, void* mem_ptr); 53 54 void* GetPreRegisteredHandle(const CustomMemTensorInfo& info); 55 56 executorch::runtime::Error SetMemHandle( 57 const std::shared_ptr<TensorWrapper>& tensor_wrapper, 58 void* mem_ptr, 59 Qnn_MemHandle_t handle); 60 61 private: 62 void DeRegisterMem(); 63 64 const QnnImplementation& implementation_; 65 QnnContext* context_; 66 std::unordered_map<Qnn_MemHandle_t, void*> registered_map_; 67 std::unordered_map<CustomMemTensorInfo, void*> pre_registered_handles_; 68 std::unordered_map<executorch::aten::ScalarType, Qnn_DataType_t> 69 scalar_type_to_qnn_dtype_ = { 70 {executorch::aten::ScalarType::Int, 71 Qnn_DataType_t::QNN_DATATYPE_INT_32}, 72 {executorch::aten::ScalarType::Float, 73 Qnn_DataType_t::QNN_DATATYPE_FLOAT_32}, 74 {executorch::aten::ScalarType::Char, 75 Qnn_DataType_t::QNN_DATATYPE_SFIXED_POINT_8}, 76 {executorch::aten::ScalarType::Short, 77 Qnn_DataType_t::QNN_DATATYPE_SFIXED_POINT_16}, 78 {executorch::aten::ScalarType::Byte, 79 Qnn_DataType_t::QNN_DATATYPE_UFIXED_POINT_8}, 80 {executorch::aten::ScalarType::Bits16, 81 Qnn_DataType_t::QNN_DATATYPE_UFIXED_POINT_16}, 82 }; 83 }; 84 } // namespace qnn 85 } // namespace backends 86 } // namespace executorch 87