mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-24 22:17:32 +00:00
a1db87b382
* Add SafeInt bounds checking to memory allocation size calculations. * Fix TensorRT library includes
103 lines
2.6 KiB
C++
103 lines
2.6 KiB
C++
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
#include "core/codegen/common/utils.h"
|
|
#include "core/common/cpuid_info.h"
|
|
#include "core/common/make_unique.h"
|
|
#include "core/common/safeint.h"
|
|
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
|
|
namespace onnxruntime {
|
|
|
|
std::unique_ptr<char[]> GetEnv(const char* var) {
|
|
char* val = nullptr;
|
|
#if _MSC_VER
|
|
size_t len;
|
|
|
|
if (_dupenv_s(&val, &len, var)) {
|
|
// Something went wrong, just return nullptr.
|
|
return nullptr;
|
|
}
|
|
#else
|
|
val = getenv(var);
|
|
#endif // _MSC_VER
|
|
|
|
if (val == nullptr) {
|
|
return nullptr;
|
|
}
|
|
|
|
// On windows, we will have to explicitly free val. Instead of returning val
|
|
// to its caller and make distinguish between windows and linux, we return
|
|
// a unique_ptr, and it will be destroyed automatically after the caller
|
|
// completes.
|
|
size_t len_val = strlen(val) + 1;
|
|
auto p = onnxruntime::make_unique<char[]>(len_val);
|
|
// use explicit loop to get ride of VC's warning on unsafe copy
|
|
for (size_t i = 0; i < len_val; ++i) {
|
|
p[i] = val[i];
|
|
}
|
|
return p;
|
|
}
|
|
|
|
bool IsEnvVarDefined(const char* var) {
|
|
auto val = GetEnv(var);
|
|
return val != nullptr;
|
|
}
|
|
|
|
int64_t TotalSize(const std::vector<int64_t>& shape) {
|
|
SafeInt<int64_t> total = 1;
|
|
for (auto s : shape) {
|
|
total *= s;
|
|
}
|
|
return total;
|
|
}
|
|
|
|
// Return the strides for the input shape, i.e. the number of
|
|
// elements contained by a single element of current dimension.
|
|
// For example, for shape[3][4][5][6], strides will be
|
|
// [4*5*6, 5*6, 6, 1], i.e. [120, 30, 6, 1]
|
|
void GetStrides(const int64_t* shape, int ndim, std::vector<int64_t>& strides) {
|
|
strides.resize(ndim);
|
|
strides[ndim - 1] = 1;
|
|
for (int64_t i = ndim - 2; i >= 0; i--) {
|
|
strides[i] = strides[i + 1] * shape[i + 1];
|
|
}
|
|
}
|
|
|
|
// Common utils to get target option
|
|
TargetFeature GetTargetInfo(const codegen::CodeGenSettings& settings) {
|
|
TargetFeature feature;
|
|
|
|
std::string target_str = "";
|
|
if (settings.HasOption(nuphar::kNupharCodeGenTarget)) {
|
|
target_str = settings.GetOptionValue(nuphar::kNupharCodeGenTarget);
|
|
}
|
|
|
|
bool isAVX = false;
|
|
bool isAVX2 = false;
|
|
bool isAVX512 = false;
|
|
if (target_str == "avx") {
|
|
isAVX = true;
|
|
} else if (target_str == "avx2") {
|
|
isAVX = true;
|
|
isAVX2 = true;
|
|
} else if (target_str == "avx512") {
|
|
isAVX = true;
|
|
isAVX2 = true;
|
|
isAVX512 = true;
|
|
} else {
|
|
isAVX = CPUIDInfo::GetCPUIDInfo().HasAVX();
|
|
isAVX2 = CPUIDInfo::GetCPUIDInfo().HasAVX2();
|
|
isAVX512 = CPUIDInfo::GetCPUIDInfo().HasAVX512Skylake();
|
|
}
|
|
|
|
feature.hasAVX = isAVX;
|
|
feature.hasAVX2 = isAVX2;
|
|
feature.hasAVX512 = isAVX512;
|
|
|
|
return feature;
|
|
}
|
|
|
|
} // namespace onnxruntime
|