Skip to content

Commit 38341fa

Browse files
ultajainAtul Jainjpuneet
authored
[DelitePy] Add fn set_xnnpack_num_threads() (#167)
* Exposed get_hardware_info and set_xnnpack_num_threads functions in script Signed-off-by: Atul Jain <atul.jain@nimbleedgehq.ai> * Fixes Signed-off-by: Puneet Jindal <puneet.jindal@nimbleedgehq.ai> * Enable XNNPACK selectively for Android and iOS Signed-off-by: Puneet Jindal <puneet.jindal@nimbleedgehq.ai> --------- Signed-off-by: Puneet Jindal <puneet.jindal@nimbleedgehq.ai> Co-authored-by: Atul Jain <atul.jain@nimbleedgehq.ai> Co-authored-by: Puneet Jindal <puneet.jindal@nimbleedgehq.ai>
1 parent 0eb1feb commit 38341fa

8 files changed

Lines changed: 88 additions & 26 deletions

File tree

‎coreruntime/nimblenet/data_variable/include/data_variable_enums.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,7 @@ enum MemberFuncType {
134134
ADD_CONTEXT,
135135
LIST_COMPATIBLE_LLMS,
136136
GET_HARDWARE_INFO,
137+
SET_XNNPACK_NUM_THREADS,
137138
#if DELITEAI_TARGET_OS_ANDROID || DELITEAI_TARGET_OS_IOS
138139
CONVERT_TEXT_TO_PHONEMES,
139140
#endif // DELITEAI_TARGET_OS_ANDROID || DELITEAI_TARGET_OS_IOS

‎coreruntime/nimblenet/data_variable/include/model_nimble_net_variable.hpp‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,15 @@ class ModelNimbleNetVariable final : public DataVariable {
112112
nlohmann::json to_json() const override { return "[Model]"; }
113113

114114
public:
115+
/**
116+
* @brief Sets the number of intra-op threads for XNNPACK
117+
*
118+
* Only supported for ONNX models. Throws if not supported.
119+
*
120+
* @param num_threads Number of threads to use
121+
*/
122+
static void set_xnnpack_intra_op_num_threads(int num_threads);
123+
115124
/*
116125
DELITEPY_DOC_BLOCK_BEGIN
117126

‎coreruntime/nimblenet/data_variable/include/nimble_net_data_variable.hpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,8 @@ def zeros(shape: list[int], dtype: str) -> Tensor:
113113

114114
OpReturnType get_hardware_info(const std::vector<OpReturnType>& arguments, CallStack& stack);
115115

116+
OpReturnType set_xnnpack_num_threads(const std::vector<OpReturnType>& arguments, CallStack& stack);
117+
116118
OpReturnType create_concurrent_executor(const std::vector<OpReturnType>& arguments);
117119

118120
OpReturnType set_threads(const std::vector<OpReturnType>& arguments);

‎coreruntime/nimblenet/data_variable/src/data_variable.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@ std::map<std::string, int> DataVariable::_memberFuncMap = {
106106
{"add_context", MemberFuncType::ADD_CONTEXT},
107107
{"list_compatible_llms", MemberFuncType::LIST_COMPATIBLE_LLMS},
108108
{"get_hardware_info", MemberFuncType::GET_HARDWARE_INFO},
109+
{"set_xnnpack_num_threads", MemberFuncType::SET_XNNPACK_NUM_THREADS},
109110
#if DELITEAI_TARGET_OS_ANDROID || DELITEAI_TARGET_OS_IOS
110111
{"convert_text_to_phonemes", MemberFuncType::CONVERT_TEXT_TO_PHONEMES},
111112
#endif // DELITEAI_TARGET_OS_ANDROID || DELITEAI_TARGET_OS_IOS
@@ -199,6 +200,7 @@ std::map<int, std::string> DataVariable::_inverseMemberFuncMap = {
199200
{MemberFuncType::ADD_CONTEXT, "add_context"},
200201
{MemberFuncType::LIST_COMPATIBLE_LLMS, "list_compatible_llms"},
201202
{MemberFuncType::GET_HARDWARE_INFO, "get_hardware_info"},
203+
{MemberFuncType::SET_XNNPACK_NUM_THREADS, "set_xnnpack_num_threads"},
202204
#if DELITEAI_TARGET_OS_ANDROID || DELITEAI_TARGET_OS_IOS
203205
{MemberFuncType::CONVERT_TEXT_TO_PHONEMES, "convert_text_to_phonemes"},
204206
#endif // DELITEAI_TARGET_OS_ANDROID || DELITEAI_TARGET_OS_IOS

‎coreruntime/nimblenet/data_variable/src/model_nimble_net_variable.cpp‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,15 @@
77
#include "model_nimble_net_variable.hpp"
88

99
#include "asset_load_job.hpp"
10+
#include "task_onnx_model.hpp"
11+
12+
void ModelNimbleNetVariable::set_xnnpack_intra_op_num_threads(int num_threads) {
13+
#ifdef ONNX_EXECUTOR
14+
TaskONNXModel::set_xnnpack_intra_op_num_threads(num_threads);
15+
return;
16+
#endif // ONNX_EXECUTOR
17+
THROW("set_xnnpack_intra_op_num_threads not supported");
18+
}
1019

1120
std::shared_ptr<FutureDataVariable> ModelNimbleNetVariable::load_async(
1221
const std::string& modelName, CommandCenter* commandCenter) {

‎coreruntime/nimblenet/data_variable/src/nimble_net_data_variable.cpp‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -468,6 +468,14 @@ OpReturnType NimbleNetDataVariable::get_hardware_info(const std::vector<OpReturn
468468
return DataVariable::get_map_from_json_object(std::move(j));
469469
}
470470

471+
OpReturnType NimbleNetDataVariable::set_xnnpack_num_threads(
472+
const std::vector<OpReturnType>& arguments, CallStack& stack) {
473+
THROW_ARGUMENTS_NOT_MATCH(arguments.size(), 1, MemberFuncType::SET_XNNPACK_NUM_THREADS);
474+
475+
ModelNimbleNetVariable::set_xnnpack_intra_op_num_threads(arguments[0]->get_int32());
476+
return std::make_shared<NoneVariable>();
477+
}
478+
471479
/*
472480
* 1. Get device tier
473481
* 2. Get the LLMs in deployment from asset manager in cloud
@@ -595,6 +603,8 @@ OpReturnType NimbleNetDataVariable::call_function(int memberFuncIndex,
595603
return list_compatible_llms(arguments);
596604
case MemberFuncType::GET_HARDWARE_INFO:
597605
return get_hardware_info(arguments, stack);
606+
case MemberFuncType::SET_XNNPACK_NUM_THREADS:
607+
return set_xnnpack_num_threads(arguments, stack);
598608
#if DELITEAI_TARGET_OS_ANDROID || DELITEAI_TARGET_OS_IOS
599609
case MemberFuncType::CONVERT_TEXT_TO_PHONEMES:
600610
// TODO (jpuneet): Should this be a part of another DelitePy module?

‎coreruntime/nimblenet/executors/onnx/include/task_onnx_model.hpp‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,13 @@ class TaskONNXModel : public TaskBaseModel {
9393
static OpReturnType get_tensor_variable_from_onnx_tensor(Ort::Value onnx_tensor);
9494

9595
public:
96+
/**
97+
* @brief Sets the number of intra-op threads for the XNNPACK backend.
98+
*
99+
* @param num_threads Number of threads to use for intra-op parallelism
100+
*/
101+
static void set_xnnpack_intra_op_num_threads(int num_threads);
102+
96103
/**
97104
* @brief Constructs a TaskONNXModel instance.
98105
*

‎coreruntime/nimblenet/executors/onnx/src/task_onnx_model.cpp‎

Lines changed: 48 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,50 @@
88

99
#include "data_variable.hpp"
1010
#include "nimble_net_util.hpp"
11+
#include "nimble_net/config.h"
1112
#include "onnx_operators.hpp"
1213
#include "tensor_data_variable.hpp"
1314

15+
#ifdef ORT_EXTENSIONS
16+
EXTERN_C_BEGIN
17+
18+
OrtStatus* ORT_API_CALL RegisterCustomOps(OrtSessionOptions* options, const OrtApiBase* api);
19+
20+
int ORT_API_CALL GetActiveOrtAPIVersion();
21+
22+
EXTERN_C_END
23+
#endif // ORT_EXTENSIONS
24+
25+
// =================================================================================================
26+
27+
namespace {
28+
29+
// Number of threads for the XNNPACK backend
30+
int xnnpack_intra_op_num_threads = 6;
31+
32+
/**
33+
* @brief Configures common ONNX session options.
34+
*
35+
* @param sessionOptions Session options to configure.
36+
*/
37+
void add_common_session_options(Ort::SessionOptions& sessionOptions) {
38+
sessionOptions.AddConfigEntry("session.use_ort_model_bytes_directly", "1");
39+
#if DELITEAI_TARGET_OS_ANDROID || DELITEAI_TARGET_OS_IOS
40+
sessionOptions.AppendExecutionProvider(
41+
"XNNPACK", {std::pair<std::string, std::string>(
42+
"intra_op_num_threads", ne::fmt("%d", xnnpack_intra_op_num_threads).str)});
43+
#endif // DELITEAI_TARGET_OS_ANDROID || DELITEAI_TARGET_OS_IOS
44+
#ifdef ORT_EXTENSIONS
45+
Ort::ThrowOnError(RegisterCustomOps((OrtSessionOptions*)sessionOptions, OrtGetApiBase()));
46+
#endif // ORT_EXTENSIONS
47+
}
48+
49+
} // namespace
50+
51+
// =================================================================================================
52+
1453
Ort::Env TaskONNXModel::_myEnv =
15-
Ort::Env(OrtLoggingLevel::ORT_LOGGING_LEVEL_FATAL, "ONNX Inference Environment");
54+
Ort::Env(OrtLoggingLevel::ORT_LOGGING_LEVEL_FATAL, "ONNX Inference Environment");
1655

1756
int TaskONNXModel::create_input_tensor_and_set_data_ptr(const OpReturnType req, int modelInputIndex,
1857
Ort::Value&& returnedInputTensor) {
@@ -117,7 +156,7 @@ Ort::SessionOptions TaskONNXModel::get_session_options_from_json(const nlohmann:
117156
std::string configProviderName = epConfig.find("providerName") != epConfig.end()
118157
? epConfig["providerName"].get<std::string>()
119158
: "";
120-
#if defined(__ANDROID__) && defined(NNAPI)
159+
#if DELITEAI_TARGET_OS_ANDROID && defined(NNAPI)
121160
if (configProviderName == "XNNPACK") {
122161
std::unordered_map<std::string, std::string> provider_options = {};
123162
if (epConfig.find("providerSettings") != epConfig.end()) {
@@ -132,7 +171,7 @@ Ort::SessionOptions TaskONNXModel::get_session_options_from_json(const nlohmann:
132171
}
133172
Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_Nnapi(sessionOptions, nnapi_flags));
134173
}
135-
#elif IOS
174+
#elif DELITEAI_TARGET_OS_IOS
136175
#endif
137176
if (epConfig.find("intraOpNumThreads") != epConfig.end()) {
138177
sessionOptions.SetIntraOpNumThreads(epConfig["intraOpNumThreads"].get<int>());
@@ -148,36 +187,15 @@ Ort::SessionOptions TaskONNXModel::get_session_options_from_json(const nlohmann:
148187
return sessionOptions;
149188
}
150189

151-
#ifdef ORT_EXTENSIONS
152-
#ifdef __cplusplus
153-
extern "C" {
154-
#endif // __cplusplus
155-
156-
OrtStatus* ORT_API_CALL RegisterCustomOps(OrtSessionOptions* options, const OrtApiBase* api);
157-
158-
int ORT_API_CALL GetActiveOrtAPIVersion();
159-
160-
#ifdef __cplusplus
161-
}
162-
#endif // __cplusplus
163-
#endif // ORT_EXTENSIONS
164-
165-
void add_common_session_options(Ort::SessionOptions& sessionOptions) {
166-
sessionOptions.AddConfigEntry("session.use_ort_model_bytes_directly", "1");
167-
#ifdef ORT_EXTENSIONS
168-
Ort::ThrowOnError(RegisterCustomOps((OrtSessionOptions*)sessionOptions, OrtGetApiBase()));
169-
#endif // ORT_EXTENSIONS
170-
}
171-
172190
void TaskONNXModel::load_model_from_buffer() {
173191
Ort::CustomOpDomain deliteai_operator_domain{"dev.deliteai"};
174192
register_custom_onnx_operators(deliteai_operator_domain);
175193
nlohmann::json epConfigListToCheck = nlohmann::json::array();
176-
#if defined(__ANDROID__)
194+
#if DELITEAI_TARGET_OS_ANDROID
177195
if (_epConfig.contains("android")) {
178196
epConfigListToCheck = _epConfig.at("android");
179197
}
180-
#elif IOS
198+
#elif DELITEAI_TARGET_OS_IOS
181199
if (_epConfig.contains("ios")) {
182200
epConfigListToCheck = _epConfig.at("ios");
183201
}
@@ -249,6 +267,10 @@ void TaskONNXModel::load_model_meta_data() {
249267
}
250268
}
251269

270+
void TaskONNXModel::set_xnnpack_intra_op_num_threads(int num_threads) {
271+
xnnpack_intra_op_num_threads = num_threads;
272+
}
273+
252274
TaskONNXModel::TaskONNXModel(const std::string& plan, const std::string& version,
253275
const std::string& modelId,
254276
const nlohmann::json& executionProviderConfig,

0 commit comments

Comments
 (0)