xref: /aosp_15_r20/external/executorch/backends/qualcomm/aot/python/PyQnnWrapperAdaptor.h (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
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 
10 #include <executorch/backends/qualcomm/aot/wrappers/OpWrapper.h>
11 #include <executorch/backends/qualcomm/aot/wrappers/TensorWrapper.h>
12 #include <executorch/backends/qualcomm/runtime/Logging.h>
13 #include <pybind11/numpy.h>
14 #include <pybind11/pybind11.h>
15 #include <pybind11/stl.h>
16 namespace py = pybind11;
17 namespace executorch {
18 namespace backends {
19 namespace qnn {
20 class PyQnnOpWrapper {
21  public:
PyQnnOpWrapper(const std::string & name,const std::string & package_name,const std::string & op_type)22   explicit PyQnnOpWrapper(
23       const std::string& name,
24       const std::string& package_name,
25       const std::string& op_type) {
26     op_wrapper_ = std::make_shared<OpWrapper>(name, package_name, op_type);
27   }
AddInputTensors(const std::vector<std::shared_ptr<TensorWrapper>> & tensors)28   void AddInputTensors(
29       const std::vector<std::shared_ptr<TensorWrapper>>& tensors) {
30     op_wrapper_->AddInputTensors(tensors);
31   }
32 
AddOutputTensors(const std::vector<std::shared_ptr<TensorWrapper>> & tensors)33   void AddOutputTensors(
34       const std::vector<std::shared_ptr<TensorWrapper>>& tensors) {
35     op_wrapper_->AddOutputTensors(tensors);
36   }
37 
AddTensorParam(const std::string & name,Qnn_DataType_t data_type,std::uint32_t rank,const std::vector<uint32_t> & dims,py::array & data,bool copy_data)38   void AddTensorParam(
39       const std::string& name,
40       Qnn_DataType_t data_type,
41       std::uint32_t rank,
42       const std::vector<uint32_t>& dims,
43       py::array& data,
44       bool copy_data) {
45     op_wrapper_->AddTensorParam(
46         name, data_type, rank, dims.data(), data.data(), copy_data);
47   }
48 
AddScalarParam(const std::string & name,Qnn_DataType_t data_type,py::dict & attrData)49   void AddScalarParam(
50       const std::string& name,
51       Qnn_DataType_t data_type,
52       py::dict& attrData) {
53     switch (data_type) {
54       case Qnn_DataType_t::QNN_DATATYPE_INT_32:
55         op_wrapper_->AddScalarParam(
56             name, data_type, attrData["data"].cast<int32_t>());
57         break;
58       case Qnn_DataType_t::QNN_DATATYPE_INT_16:
59         op_wrapper_->AddScalarParam(
60             name, data_type, attrData["data"].cast<int16_t>());
61         break;
62       case Qnn_DataType_t::QNN_DATATYPE_INT_8:
63         op_wrapper_->AddScalarParam(
64             name, data_type, attrData["data"].cast<int8_t>());
65         break;
66       case Qnn_DataType_t::QNN_DATATYPE_UINT_32:
67         op_wrapper_->AddScalarParam(
68             name, data_type, attrData["data"].cast<uint32_t>());
69         break;
70       case Qnn_DataType_t::QNN_DATATYPE_UINT_16:
71         op_wrapper_->AddScalarParam(
72             name, data_type, attrData["data"].cast<uint16_t>());
73         break;
74       case Qnn_DataType_t::QNN_DATATYPE_UINT_8:
75         op_wrapper_->AddScalarParam(
76             name, data_type, attrData["data"].cast<uint8_t>());
77         break;
78       case Qnn_DataType_t::QNN_DATATYPE_FLOAT_32:
79       case Qnn_DataType_t::QNN_DATATYPE_FLOAT_16:
80         op_wrapper_->AddScalarParam(
81             name, data_type, attrData["data"].cast<float>());
82         break;
83       case Qnn_DataType_t::QNN_DATATYPE_BOOL_8:
84         op_wrapper_->AddScalarParam(
85             name, data_type, attrData["data"].cast<bool>());
86         break;
87       default:
88         QNN_EXECUTORCH_LOG_ERROR(
89             "%s has invalid data type: %d", name.c_str(), data_type);
90         break;
91     }
92   }
GetOpWrapper()93   std::shared_ptr<OpWrapper>& GetOpWrapper() {
94     return op_wrapper_;
95   }
96 
97  private:
98   std::shared_ptr<OpWrapper> op_wrapper_;
99 };
100 
101 class PyQnnTensorWrapper {
102  public:
PyQnnTensorWrapper(const std::shared_ptr<TensorWrapper> & wrapper)103   explicit PyQnnTensorWrapper(const std::shared_ptr<TensorWrapper>& wrapper) {
104     tensor_wrapper_ = wrapper;
105   }
106   struct EncodingData {
107     float scale;
108     int32_t offset;
109   };
110   struct Encoding {
111     py::array_t<EncodingData> data;
112     int32_t axis;
113   };
114 
GetDims()115   py::array_t<std::uint32_t> GetDims() {
116     std::uint32_t* dim = tensor_wrapper_->GetDims();
117     size_t shape[1]{tensor_wrapper_->GetRank()};
118     size_t stride[1]{sizeof(std::uint32_t)};
119     auto ret = py::array_t<std::uint32_t>(shape, stride);
120     auto view = ret.mutable_unchecked<1>();
121     for (int i = 0; i < ret.shape(0); ++i) {
122       view(i) = dim[i];
123     }
124     return ret;
125   }
GetName()126   std::string GetName() {
127     return tensor_wrapper_->GetName();
128   }
GetDataType()129   Qnn_DataType_t GetDataType() {
130     return tensor_wrapper_->GetDataType();
131   }
GetEncodings()132   Encoding GetEncodings() {
133     auto q_param = tensor_wrapper_->GetQuantizeParams();
134     size_t stride[1]{sizeof(EncodingData)};
135 
136     switch (q_param.quantizationEncoding) {
137       case QNN_QUANTIZATION_ENCODING_SCALE_OFFSET: {
138         Qnn_ScaleOffset_t data = q_param.scaleOffsetEncoding;
139         size_t shape[1]{1};
140         auto enc_data = py::array_t<EncodingData>(shape, stride);
141         auto view = enc_data.mutable_unchecked<1>();
142         view(0) = {data.scale, data.offset};
143         return {enc_data, -1};
144       }
145       case QNN_QUANTIZATION_ENCODING_AXIS_SCALE_OFFSET: {
146         Qnn_AxisScaleOffset_t data = q_param.axisScaleOffsetEncoding;
147         size_t shape[1]{data.numScaleOffsets};
148         auto enc_data = py::array_t<EncodingData>(shape, stride);
149         auto view = enc_data.mutable_unchecked<1>();
150         for (int i = 0; i < enc_data.shape(0); ++i) {
151           view(i) = {data.scaleOffset[i].scale, data.scaleOffset[i].offset};
152         }
153         return {enc_data, data.axis};
154       }
155       case QNN_QUANTIZATION_ENCODING_BW_SCALE_OFFSET: {
156         Qnn_BwScaleOffset_t data = q_param.bwScaleOffsetEncoding;
157         size_t shape[1]{1};
158         auto enc_data = py::array_t<EncodingData>(shape, stride);
159         auto view = enc_data.mutable_unchecked<1>();
160         view(0) = {data.scale, data.offset};
161         return {enc_data, -1};
162       }
163       case QNN_QUANTIZATION_ENCODING_BW_AXIS_SCALE_OFFSET: {
164         Qnn_BwAxisScaleOffset_t data = q_param.bwAxisScaleOffsetEncoding;
165         size_t shape[1]{data.numElements};
166         auto enc_data = py::array_t<EncodingData>(shape, stride);
167         auto view = enc_data.mutable_unchecked<1>();
168         for (int i = 0; i < enc_data.shape(0); ++i) {
169           view(i) = {data.scales[i], data.offsets[i]};
170         }
171         return {enc_data, data.axis};
172       }
173       default:
174         QNN_EXECUTORCH_LOG_WARN(
175             "%s QNN_QUANTIZATION_ENCODING_UNDEFINED detected",
176             GetName().c_str());
177         break;
178     }
179     return {};
180   }
181 
182  private:
183   std::shared_ptr<TensorWrapper> tensor_wrapper_;
184 };
185 } // namespace qnn
186 } // namespace backends
187 } // namespace executorch
188