saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-19 19:00:47 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge remote-tracking branch 'upstream/master' into DmlDev

Browse source
This commit is contained in:
ISS Build Account 2020-08-23 11:05:16 +00:00
commit 7c0ff3742f
76 changed files with 1056 additions and 321 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
cmake/onnxruntime_framework.cmake
View file

@ -11,6 +11,8 @@ if (onnxruntime_MINIMAL_BUILD)
file(GLOB onnxruntime_framework_src_exclude
"${ONNXRUNTIME_ROOT}/core/framework/provider_bridge_ort.cc"
"${ONNXRUNTIME_ROOT}/core/framework/graph_partitioner.*"
"${ONNXRUNTIME_INCLUDE_DIR}/core/framework/customregistry.h"
"${ONNXRUNTIME_ROOT}/core/framework/customregistry.cc"
)
list(REMOVE_ITEM onnxruntime_framework_srcs ${onnxruntime_framework_src_exclude})

5
cmake/onnxruntime_graph.cmake
View file

@ -11,10 +11,13 @@ file(GLOB_RECURSE onnxruntime_graph_src CONFIGURE_DEPENDS
set(onnxruntime_graph_src_exclude_patterns)
if (onnxruntime_MINIMAL_BUILD)
# remove schema registration of contrib ops
# remove schema registration support
list(APPEND onnxruntime_graph_src_exclude_patterns
"${ONNXRUNTIME_INCLUDE_DIR}/core/graph/schema_registry.h"
"${ONNXRUNTIME_ROOT}/core/graph/schema_registry.cc"
"${ONNXRUNTIME_ROOT}/core/graph/contrib_ops/*defs.h"
"${ONNXRUNTIME_ROOT}/core/graph/contrib_ops/*defs.cc"
)
# no Function support initially

7
docs/C_API.md
View file

@ -21,6 +21,13 @@ is as follows
* Create env using ```CreateEnvWithGlobalThreadPools()```
* Create session and call ```DisablePerSessionThreads()``` on the session options object
* Call ```Run()``` as usual
* **Share allocator(s) between sessions:** Allow multiple sessions in the same process to use the same allocator(s). This
allocator is first registered in the env and then reused by all sessions that use the same env instance unless a session
chooses to override this by setting ```session_state.use_env_allocators``` to "0". Usage of this feature is as follows
* Register an allocator created by ORT using the ```CreateAndRegisterAllocator``` API.
* Set ```session.use_env_allocators``` to "1" for each session that wants to use the env registered allocators.
* See test ```TestSharedAllocatorUsingCreateAndRegisterAllocator``` in
onnxruntime/test/shared_lib/test_inference.cc for an example.
## Usage Overview

7
include/onnxruntime/core/common/common.h
View file

@ -242,4 +242,11 @@ inline std::wstring ToWideString(const std::wstring& s) { return s; }
inline std::string ToWideString(const std::string& s) { return s; }
#endif
// from http://www.open-std.org/jtc1/sc22/wg21/docs/papers/2014/n3876.pdf
template <class T>
inline void HashCombine(std::uint64_t& seed, const T& v) {
std::hash<T> hasher;
seed ^= hasher(v) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
}
} // namespace onnxruntime

11
include/onnxruntime/core/framework/execution_provider.h
View file

@ -22,7 +22,8 @@ class KernelRegistryManager;
/**
Logical device representation.
*/
typedef std::map<int, AllocatorPtr> AllocatorMap;
using AllocatorMap = std::map<int, AllocatorPtr>;
using MemoryInfoSet = std::set<OrtMemoryInfo>;
// if we are export the fused function to dll, the function will still in the same binary as onnxruntime
// use std function to give execution provider some chance to capture some state.
@ -34,8 +35,8 @@ using DestroyFunctionStateFunc = std::function<void(FunctionState)>;
using UnorderedMapStringToString = std::unordered_map<std::string, std::string>;
//data types for execution provider options
using ProviderOptionsVector = std::vector<UnorderedMapStringToString>;
using ProviderOptionsMap = std::unordered_map<std::string, UnorderedMapStringToString>;
using ProviderOptionsVector = std::vector<UnorderedMapStringToString>;
using ProviderOptionsMap = std::unordered_map<std::string, UnorderedMapStringToString>;
struct NodeComputeInfo {
CreateFunctionStateFunc create_state_func;
@ -113,7 +114,7 @@ class IExecutionProvider {
/**
Store execution provider's configurations.
*/
void SetProviderOptions(UnorderedMapStringToString& options) {
void SetProviderOptions(UnorderedMapStringToString& options) {
provider_options_ = options;
}
@ -165,6 +166,7 @@ class IExecutionProvider {
virtual common::Status OnSessionInitializationEnd();
void InsertAllocator(AllocatorPtr allocator);
void ReplaceAllocator(AllocatorPtr allocator);
/**
Given a list of fused_node, return create_state/compute/release_state func for each node.
@ -193,6 +195,7 @@ class IExecutionProvider {
private:
const std::string type_;
AllocatorMap allocators_;
MemoryInfoSet mem_info_set_; // to ensure only allocators with unique OrtMemoryInfo are registered in the provider.
//It will be set when this object is registered to a session
const logging::Logger* logger_ = nullptr;
// convenience list of the allocators so GetAllocatorList doesn't have to build a new vector each time

32
include/onnxruntime/core/framework/kernel_def_builder.h
View file

@ -88,9 +88,37 @@ class KernelDef {
bool IsConflict(const KernelDef& other) const;
uint64_t GetHash() const noexcept {
// if we need to support different hash versions we can update CalculateHash to take a version number
// and calculate any non-default versions dynamically. we only use this during kernel lookup so
// it's not performance critical
return hash_;
}
private:
friend class KernelDefBuilder;
// call once the KernelDef has been built
void CalculateHash() {
// use name, start/end, domain, provider and the type constraints.
// we wouldn't have two kernels that only differed by the inplace or alias info or memory types.
// currently nothing sets exec_queue_id either (and would assumably be a runtime thing and not part of the base
// kernel definition)
hash_ = 0; // reset in case this is called multiple times
HashCombine(hash_, op_name_);
HashCombine(hash_, op_since_version_start_);
HashCombine(hash_, op_since_version_end_);
HashCombine(hash_, op_domain_);
HashCombine(hash_, provider_type_);
for (const auto& key_value : type_constraints_) {
HashCombine(hash_, key_value.first);
for (const auto& data_type : key_value.second) {
// need to construct a std::string so it doesn't hash the address of a const char*
HashCombine(hash_, std::string(DataTypeImpl::ToString(data_type)));
}
}
}
// The operator name supported by <*this> kernel..
std::string op_name_;
@ -128,6 +156,9 @@ class KernelDef {
OrtMemType default_inputs_mem_type_{OrtMemTypeDefault};
// Default memory type for all outputs
OrtMemType default_outputs_mem_type_{OrtMemTypeDefault};
// hash of kernel definition for lookup in minimal build
uint64_t hash_ = 0;
};
class KernelDefBuilder {
@ -283,6 +314,7 @@ class KernelDefBuilder {
Return the kernel definition, passing ownership of the KernelDef to the caller
*/
std::unique_ptr<KernelDef> Build() {
kernel_def_->CalculateHash();
return std::move(kernel_def_);
}

27
include/onnxruntime/core/framework/kernel_registry.h
View file

@ -19,6 +19,14 @@ class KernelRegistry {
Status Register(KernelCreateInfo&& create_info) ORT_MUST_USE_RESULT;
#if !defined(ORT_MINIMAL_BUILD)
static bool HasImplementationOf(const KernelRegistry& r, const onnxruntime::Node& node,
onnxruntime::ProviderType exec_provider) {
const KernelCreateInfo* info;
Status st = r.TryFindKernel(node, exec_provider, &info);
return st.IsOK();
}
// factory functions should always return a unique_ptr for maximum flexibility
// for its clients unless the factory is managing the lifecycle of the pointer
// itself.
@ -30,15 +38,16 @@ class KernelRegistry {
std::unique_ptr<OpKernel>& op_kernel) const ORT_MUST_USE_RESULT;
// Check if an execution provider can create kernel for a node and return the kernel if so
Status TryFindKernel(const onnxruntime::Node& node,
onnxruntime::ProviderType exec_provider, const KernelCreateInfo** out) const;
Status TryFindKernel(const onnxruntime::Node& node, onnxruntime::ProviderType exec_provider,
const KernelCreateInfo** out) const;
static bool HasImplementationOf(const KernelRegistry& r, const onnxruntime::Node& node,
onnxruntime::ProviderType exec_provider) {
const KernelCreateInfo* info;
Status st = r.TryFindKernel(node, exec_provider, &info);
return st.IsOK();
}
#endif
// Check if an execution provider can create kernel for a node and return the kernel if so.
// Kernel matching is via kernel_def_hash.
Status TryFindKernel(const onnxruntime::Node& node, onnxruntime::ProviderType exec_provider,
uint64_t kernel_def_hash,
const KernelCreateInfo** out) const;
bool IsEmpty() const { return kernel_creator_fn_map_.empty(); }
@ -50,6 +59,7 @@ class KernelRegistry {
#endif
private:
#if !defined(ORT_MINIMAL_BUILD)
// Check whether the types of inputs/outputs of the given node match the extra
// type-constraints of the given kernel. This serves two purposes: first, to
// select the right kernel implementation based on the types of the arguments
@ -67,6 +77,7 @@ class KernelRegistry {
static bool VerifyKernelDef(const onnxruntime::Node& node,
const KernelDef& kernel_def,
std::string& error_str);
#endif
static std::string GetMapKey(const std::string& op_name, const std::string& domain, const std::string& provider) {
std::string key(op_name);

2
include/onnxruntime/core/framework/op_kernel.h
View file

@ -17,8 +17,8 @@
#include "core/framework/sparse_tensor.h"
#include "core/graph/constants.h"
#include "core/graph/graph_viewer.h"
#include "core/graph/onnx_protobuf.h"
#include "gsl/gsl"
#include "onnx/defs/schema.h"
namespace onnxruntime {
class IExecutionFrame;

5
include/onnxruntime/core/graph/basic_types.h
View file

@ -12,6 +12,11 @@ class ValueInfoProto;
class TensorProto;
class TypeProto;
class AttributeProto;
// define types that would come from the ONNX library if we were building against it.
#if defined(ORT_MINIMAL_BUILD)
using OperatorSetVersion = int;
#endif
} // namespace ONNX_NAMESPACE
namespace onnxruntime {

7
include/onnxruntime/core/graph/graph.h
View file

@ -578,6 +578,13 @@ class Graph {
/** Returns true if an initializer value can be overridden by a graph input with the same name. */
bool CanOverrideInitializer() const noexcept { return ir_version_ >= 4; }
/** returns the initializer's TensorProto if 'name' is an initializer, is constant and
cannot be overridden at runtime. If the initializer is not found or is not constant, a nullptr is returned.
@param check_outer_scope If true and the graph is a subgraph,
check ancestor graph/s for 'name' if not found in 'graph'.
*/
const ONNX_NAMESPACE::TensorProto* GetConstantInitializer(const std::string& name, bool check_outer_scope) const;
/** Gets the Graph inputs excluding initializers.
These are the required inputs to the Graph as the initializers can be optionally overridden via graph inputs.
@remarks Contains no nullptr values. */

37
include/onnxruntime/core/graph/onnx_protobuf.h
View file

@ -18,22 +18,33 @@
#pragma warning(disable : 4100)
#pragma warning(disable : 4146) /*unary minus operator applied to unsigned type, result still unsigned*/
#pragma warning(disable : 4127)
#pragma warning(disable : 4244) /*'conversion' conversion from 'type1' to 'type2', possible loss of data*/
#pragma warning(disable : 4251) /*'identifier' : class 'type' needs to have dll-interface to be used by clients of class 'type2'*/
#pragma warning(disable : 4267) /*'var' : conversion from 'size_t' to 'type', possible loss of data*/
#pragma warning(disable : 4305) /*'identifier' : truncation from 'type1' to 'type2'*/
#pragma warning(disable : 4307) /*'operator' : integral constant overflow*/
#pragma warning(disable : 4309) /*'conversion' : truncation of constant value*/
#pragma warning(disable : 4334) /*'operator' : result of 32-bit shift implicitly converted to 64 bits (was 64-bit shift intended?)*/
#pragma warning(disable : 4355) /*'this' : used in base member initializer list*/
#pragma warning(disable : 4506) /*no definition for inline function 'function'*/
#pragma warning(disable : 4800) /*'type' : forcing value to bool 'true' or 'false' (performance warning)*/
#pragma warning(disable : 4996) /*The compiler encountered a deprecated declaration.*/
#pragma warning(disable : 6011) /*Dereferencing NULL pointer*/
#pragma warning(disable : 6387) /*'value' could be '0'*/
#pragma warning(disable : 4244) /*'conversion' conversion from 'type1' to 'type2', possible loss of data*/
#pragma warning(disable : 4251) /*'identifier' : class 'type' needs to have dll-interface to be used by clients of class 'type2'*/
#pragma warning(disable : 4267) /*'var' : conversion from 'size_t' to 'type', possible loss of data*/
#pragma warning(disable : 4305) /*'identifier' : truncation from 'type1' to 'type2'*/
#pragma warning(disable : 4307) /*'operator' : integral constant overflow*/
#pragma warning(disable : 4309) /*'conversion' : truncation of constant value*/
#pragma warning(disable : 4334) /*'operator' : result of 32-bit shift implicitly converted to 64 bits (was 64-bit shift intended?)*/
#pragma warning(disable : 4355) /*'this' : used in base member initializer list*/
#pragma warning(disable : 4506) /*no definition for inline function 'function'*/
#pragma warning(disable : 4800) /*'type' : forcing value to bool 'true' or 'false' (performance warning)*/
#pragma warning(disable : 4996) /*The compiler encountered a deprecated declaration.*/
#pragma warning(disable : 6011) /*Dereferencing NULL pointer*/
#pragma warning(disable : 6387) /*'value' could be '0'*/
#pragma warning(disable : 26495) /*Variable is uninitialized.*/
#endif
#if !defined(ORT_MINIMAL_BUILD)
#include "onnx/defs/schema.h"
#else
#include "onnx/defs/data_type_utils.h"
// stub definition of OpSchema to minimize other code changes
namespace ONNX_NAMESPACE {
class OpSchema {};
} // namespace ONNX_NAMESPACE
#endif
#include "onnx/onnx_pb.h"
#ifdef __GNUC__
#pragma GCC diagnostic pop

2
include/onnxruntime/core/optimizer/graph_transformer_utils.h
View file

@ -11,6 +11,7 @@
namespace onnxruntime {
struct FreeDimensionOverride;
class IExecutionProvider;
namespace optimizer_utils {
@ -25,6 +26,7 @@ std::vector<std::unique_ptr<RewriteRule>> GenerateRewriteRules(TransformerLevel
and the transformers_and_rules_to_enable. */
std::vector<std::unique_ptr<GraphTransformer>> GenerateTransformers(TransformerLevel level,
gsl::span<const FreeDimensionOverride> free_dimension_overrides,
const IExecutionProvider& execution_provider /*required by constant folding*/,
const std::vector<std::string>& rules_and_transformers_to_enable = {});
/** Given a TransformerLevel, this method generates a name for the rule-based graph transformer of that level. */

15
include/onnxruntime/core/session/environment.h
View file

@ -9,6 +9,7 @@
#include "core/common/status.h"
#include "core/platform/threadpool.h"
#include "core/common/logging/logging.h"
#include "core/framework/allocator.h"
struct OrtThreadingOptions;
namespace onnxruntime {
@ -54,6 +55,19 @@ class Environment {
return create_global_thread_pools_;
}
/**
* Registers an allocator for sharing between multiple sessions.
* Return an error if an allocator with the same OrtMemoryInfo is already registered.
*/
Status RegisterAllocator(AllocatorPtr allocator);
/**
* Returns the list of registered allocators in this env.
*/
const std::vector<AllocatorPtr>& GetRegisteredSharedAllocators() const {
return shared_allocators_;
}
private:
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(Environment);
@ -66,5 +80,6 @@ class Environment {
std::unique_ptr<onnxruntime::concurrency::ThreadPool> intra_op_thread_pool_;
std::unique_ptr<onnxruntime::concurrency::ThreadPool> inter_op_thread_pool_;
bool create_global_thread_pools_{false};
std::vector<AllocatorPtr> shared_allocators_;
};
} // namespace onnxruntime

23
include/onnxruntime/core/session/onnxruntime_c_api.h
View file

@ -144,6 +144,16 @@ typedef enum OrtErrorCode {
ORT_EP_FAIL,
} OrtErrorCode;
// This configures the arena based allocator used by ORT
// See ONNX_Runtime_Perf_Tuning.md for details on what these mean and how to choose these values
// Use -1 to allow ORT to choose defaults for all the options below
typedef struct OrtArenaCfg {
int max_mem;
int arena_extend_strategy; // 0 = kNextPowerOfTwo, 1 = kSameAsRequested
int initial_chunk_size_bytes;
int max_dead_bytes_per_chunk;
} OrtArenaCfg;
#define ORT_RUNTIME_CLASS(X) \
struct Ort##X; \
typedef struct Ort##X Ort##X;
@ -578,6 +588,7 @@ struct OrtApi {
// The returned pointer doesn't have to be freed.
// Always returns the same instance on every invocation.
// Please note that this is a non-arena based allocator.
ORT_API2_STATUS(GetAllocatorWithDefaultOptions, _Outptr_ OrtAllocator** out);
// Override symbolic dimensions (by specific denotation strings) with actual values if known at session initialization time to enable
@ -990,6 +1001,18 @@ struct OrtApi {
* This is a no-copy method whose pointer is only valid until the backing OrtValue is free'd.
*/
ORT_API2_STATUS(TensorAt, _Inout_ OrtValue* value, size_t* location_values, size_t location_values_count, _Outptr_ void** out);
/**
* Creates an allocator instance and registers it with the env to enable
* sharing between multiple sessions that use the same env instance.
* Lifetime of the created allocator will be valid for the duration of the environment.
* Returns an error if an allocator with the same OrtMemoryInfo is already registered.
* \param mem_info must be non-null.
* \param arena_cfg if nullptr defaults will be used.
* See docs/C_API.md for details.
*/
ORT_API2_STATUS(CreateAndRegisterAllocator, _Inout_ OrtEnv* env, _In_ const OrtMemoryInfo* mem_info,
_In_ const OrtArenaCfg* arena_cfg);
};
/*

2
include/onnxruntime/core/session/onnxruntime_cxx_api.h
View file

@ -164,6 +164,8 @@ struct Env : Base<OrtEnv> {
Env& EnableTelemetryEvents();
Env& DisableTelemetryEvents();
Env& CreateAndRegisterAllocator(const OrtMemoryInfo* mem_info, const OrtArenaCfg* arena_cfg);
static const OrtApi* s_api;
};

5
include/onnxruntime/core/session/onnxruntime_cxx_inline.h
View file

@ -308,6 +308,11 @@ inline Env& Env::DisableTelemetryEvents() {
return *this;
}
inline Env& Env::CreateAndRegisterAllocator(const OrtMemoryInfo* mem_info, const OrtArenaCfg* arena_cfg) {
ThrowOnError(GetApi().CreateAndRegisterAllocator(p_, mem_info, arena_cfg));
return *this;
}
inline CustomOpDomain::CustomOpDomain(const char* domain) {
ThrowOnError(GetApi().CreateCustomOpDomain(domain, &p_));
}

6
include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h
View file

@ -18,4 +18,8 @@
// Key for disable PrePacking,
// If the config value is set to "1" then the prepacking is disabled, otherwise prepacking is enabled (default value)
#define ORT_SESSION_OPTIONS_CONFIG_DISABLEPREPACKING "session_state.disable_prepacking"
#define ORT_SESSION_OPTIONS_CONFIG_DISABLEPREPACKING "session.disable_prepacking"
// A value of "1" means allocators registered in the env will be used. "0" means the allocators created in the session
// will be used. Use this to override the usage of env allocators on a per session level.
#define ORT_SESSION_OPTIONS_CONFIG_USE_ENV_ALLOCATORS "session.use_env_allocators"

4
onnxruntime/contrib_ops/cpu/expand_dims.cc
View file

@ -2,8 +2,6 @@
// Licensed under the MIT License.
#include "expand_dims.h"
#include "onnx/defs/schema.h"
#include "sample.h"
namespace onnxruntime {
namespace contrib {
@ -16,5 +14,5 @@ ONNX_CPU_OPERATOR_TYPED_MS_KERNEL(
.TypeConstraint("T", DataTypeImpl::AllTensorTypes())
.TypeConstraint("axis", DataTypeImpl::GetTensorType<int32_t>()),
contrib::ExpandDims);
} // namespace contrib
} // namespace contrib
} // namespace onnxruntime

3
onnxruntime/contrib_ops/cpu/sample.cc
View file

@ -2,7 +2,6 @@
// Licensed under the MIT License.
#include "sample.h"
#include "onnx/defs/schema.h"
namespace onnxruntime {
namespace contrib {
@ -13,5 +12,5 @@ ONNX_CPU_OPERATOR_TYPED_MS_KERNEL(
float,
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()).MayInplace(0, 0),
contrib::SampleOp<float>);
} // namespace contrib
} // namespace contrib
} // namespace onnxruntime

5
onnxruntime/contrib_ops/cpu/tokenizer.cc
View file

@ -2,12 +2,9 @@
// Licensed under the MIT License.
#include "core/common/common.h"
#include "core/common/utf8_util.h"
#include "core/framework/tensor.h"
#include "core/framework/op_kernel.h"
#include "core/graph/onnx_protobuf.h"
#include "onnx/defs/schema.h"
#include "core/common/utf8_util.h"
#include "re2/re2.h"
namespace onnxruntime {

6
onnxruntime/core/framework/allocatormgr.cc
View file

@ -22,7 +22,11 @@ AllocatorPtr CreateAllocator(const DeviceAllocatorRegistrationInfo& info,
onnxruntime::make_unique<MiMallocArena>(std::move(device_allocator), info.max_mem));
#else
return std::shared_ptr<IArenaAllocator>(
onnxruntime::make_unique<BFCArena>(std::move(device_allocator), info.max_mem, info.arena_extend_strategy));
onnxruntime::make_unique<BFCArena>(std::move(device_allocator),
info.max_mem,
info.arena_extend_strategy,
info.initial_chunk_size_bytes,
info.max_dead_bytes_per_chunk));
#endif
}

13
onnxruntime/core/framework/allocatormgr.h
View file

@ -11,21 +11,30 @@ namespace onnxruntime {
using DeviceAllocatorFactory = std::function<std::unique_ptr<IDeviceAllocator>(OrtDevice::DeviceId)>;
// TODO why does DeviceAllocatorRegistrationInfo have arena related configs?
// TODO even if it should, they should be inside their own struct (OrtArenaCfg) as opposed to
// littering them as individual members of DeviceAllocatorRegistrationInfo
struct DeviceAllocatorRegistrationInfo {
DeviceAllocatorRegistrationInfo(OrtMemType ort_mem_type,
DeviceAllocatorFactory alloc_factory,
size_t mem,
ArenaExtendStrategy strategy = ArenaExtendStrategy::kNextPowerOfTwo)
ArenaExtendStrategy strategy = BFCArena::DEFAULT_ARENA_EXTEND_STRATEGY,
int initial_chunk_size_bytes0 = BFCArena::DEFAULT_INITIAL_CHUNK_SIZE_BYTES,
int max_dead_bytes_per_chunk0 = BFCArena::DEFAULT_MAX_DEAD_BYTES_PER_CHUNK)
: mem_type(ort_mem_type),
factory(alloc_factory),
max_mem(mem),
arena_extend_strategy(strategy) {
arena_extend_strategy(strategy),
initial_chunk_size_bytes(initial_chunk_size_bytes0),
max_dead_bytes_per_chunk(max_dead_bytes_per_chunk0) {
}
OrtMemType mem_type;
DeviceAllocatorFactory factory;
size_t max_mem;
ArenaExtendStrategy arena_extend_strategy;
int initial_chunk_size_bytes;
int max_dead_bytes_per_chunk;
};
AllocatorPtr CreateAllocator(const DeviceAllocatorRegistrationInfo& info, OrtDevice::DeviceId device_id = 0,

27
onnxruntime/core/framework/bfc_arena.cc
View file

@ -2,11 +2,14 @@
// Licensed under the MIT License.
#include "core/framework/bfc_arena.h"
#include <type_traits>
namespace onnxruntime {
BFCArena::BFCArena(std::unique_ptr<IDeviceAllocator> resource_allocator,
size_t total_memory,
ArenaExtendStrategy arena_extend_strategy)
ArenaExtendStrategy arena_extend_strategy,
int initial_chunk_size_bytes,
int max_dead_bytes_per_chunk)
: IArenaAllocator(OrtMemoryInfo(resource_allocator->Info().name,
OrtAllocatorType::OrtArenaAllocator,
resource_allocator->Info().device,
@ -14,14 +17,17 @@ BFCArena::BFCArena(std::unique_ptr<IDeviceAllocator> resource_allocator,
resource_allocator->Info().mem_type)),
device_allocator_(std::move(resource_allocator)),
free_chunks_list_(kInvalidChunkHandle),
next_allocation_id_(1) {
LOGS_DEFAULT(INFO) << "Creating BFCArena for " << device_allocator_->Info().name;
// TODO - consider to make the initial chunk size and max 'fragmentation' (kMaxDeadBytesInChunk) values configurable.
// But first we need to add a mechanism to allow that sort of low level configuration to be done
// without adding separate parameters to SessionOptions for every single one of them.
curr_region_allocation_bytes_ = RoundedBytes(std::min(total_memory, size_t{1048576}));
next_allocation_id_(1),
initial_chunk_size_bytes_(initial_chunk_size_bytes),
max_dead_bytes_per_chunk_(max_dead_bytes_per_chunk) {
LOGS_DEFAULT(INFO) << "Creating BFCArena for " << device_allocator_->Info().name
<< " with following configs: initial_chunk_size_bytes: " << initial_chunk_size_bytes_
<< " max_dead_bytes_per_chunk: " << max_dead_bytes_per_chunk_
<< " memory limit: " << total_memory
<< " arena_extend_strategy " << static_cast<int32_t>(arena_extend_strategy);
// static_cast<std::underlying_type_t<ArenaExtendStrategy>>(arena_extend_strategy); doesn't work on this compiler
curr_region_allocation_bytes_ = RoundedBytes(std::min(total_memory, static_cast<size_t>(initial_chunk_size_bytes_)));
// Allocate the requested amount of memory.
memory_limit_ = total_memory;
stats_.bytes_limit = static_cast<int64_t>(total_memory);
@ -312,10 +318,9 @@ void* BFCArena::FindChunkPtr(BinNum bin_num, size_t rounded_bytes,
// If we can break the size of the chunk into two reasonably large
// pieces, do so. In any case don't waste more than
// kMaxDeadBytesInChunk bytes on padding this alloc.
const int64_t kMaxDeadBytesInChunk = 128 << 20; // 128mb
// max_dead_bytes_per_chunk bytes on padding this alloc.
if (chunk->size >= rounded_bytes * 2 ||
static_cast<int64_t>(chunk->size) - rounded_bytes >= kMaxDeadBytesInChunk) {
static_cast<int64_t>(chunk->size) - static_cast<int64_t>(rounded_bytes) >= max_dead_bytes_per_chunk_) {
SplitChunk(h, rounded_bytes);
chunk = ChunkFromHandle(h); // Update chunk pointer in case it moved
}

11
onnxruntime/core/framework/bfc_arena.h
View file

@ -54,9 +54,15 @@ enum class ArenaExtendStrategy : int32_t {
// all requests to allocate memory go through this interface.
class BFCArena : public IArenaAllocator {
public:
static const ArenaExtendStrategy DEFAULT_ARENA_EXTEND_STRATEGY = ArenaExtendStrategy::kNextPowerOfTwo;
static const int DEFAULT_INITIAL_CHUNK_SIZE_BYTES = 1048576;
static const int DEFAULT_MAX_DEAD_BYTES_PER_CHUNK = 128 * 1024 * 1024;
BFCArena(std::unique_ptr<IDeviceAllocator> resource_allocator,
size_t total_memory,
ArenaExtendStrategy arena_extend_strategy = ArenaExtendStrategy::kNextPowerOfTwo);
ArenaExtendStrategy arena_extend_strategy = DEFAULT_ARENA_EXTEND_STRATEGY,
int initial_chunk_size_bytes = DEFAULT_INITIAL_CHUNK_SIZE_BYTES,
int max_dead_bytes_per_chunk = DEFAULT_MAX_DEAD_BYTES_PER_CHUNK);
~BFCArena() override;
@ -443,6 +449,9 @@ class BFCArena : public IArenaAllocator {
std::unordered_map<void*, size_t> reserved_chunks_;
const int initial_chunk_size_bytes_;
const int max_dead_bytes_per_chunk_;
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(BFCArena);
};
#ifdef __GNUC__

23
onnxruntime/core/framework/execution_provider.cc
View file

@ -29,6 +29,7 @@ std::vector<std::unique_ptr<ComputeCapability>>
IExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph,
const std::vector<const KernelRegistry*>& kernel_registries) const {
std::vector<std::unique_ptr<ComputeCapability>> result;
#if !defined(ORT_MINIMAL_BUILD)
for (auto& node : graph.Nodes()) {
for (auto registry : kernel_registries) {
if (KernelRegistry::HasImplementationOf(*registry, node, Type())) {
@ -39,6 +40,9 @@ IExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph,
}
}
}
#else
ORT_NOT_IMPLEMENTED("IExecutionProvider::GetCapability is not supported in this build.");
#endif
return result;
}
@ -51,14 +55,25 @@ common::Status IExecutionProvider::OnRunEnd() { return Status::OK(); }
common::Status IExecutionProvider::OnSessionInitializationEnd() { return Status::OK(); }
// Update allocator in the provider if already present; ignore if not.
void IExecutionProvider::ReplaceAllocator(AllocatorPtr allocator) {
const auto& info = allocator->Info();
auto ite = mem_info_set_.find(info);
if (ite != mem_info_set_.end()) {
const int key = MakeKey(info.id, info.mem_type);
allocators_[key] = allocator;
}
}
void IExecutionProvider::InsertAllocator(AllocatorPtr allocator) {
const OrtMemoryInfo& info = allocator->Info();
const int key = MakeKey(info.id, info.mem_type);
auto iter = allocators_.find(key);
if (iter != allocators_.end()) {
auto ite = mem_info_set_.find(info);
if (ite != mem_info_set_.end()) {
ORT_THROW("duplicated allocator");
}
allocators_.insert(iter, {key, allocator});
const int key = MakeKey(info.id, info.mem_type);
allocators_.insert({key, allocator});
mem_info_set_.insert(ite, info);
allocator_list_.push_back(allocator);
}

9
onnxruntime/core/framework/execution_providers.h
View file

@ -56,6 +56,15 @@ class ExecutionProviders {
return exec_providers_[it->second].get();
}
IExecutionProvider* Get(onnxruntime::ProviderType provider_id) {
auto it = provider_idx_map_.find(provider_id);
if (it == provider_idx_map_.end()) {
return nullptr;
}
return exec_providers_[it->second].get();
}
bool Empty() const { return exec_providers_.empty(); }
size_t NumProviders() const { return exec_providers_.size(); }

119
onnxruntime/core/framework/kernel_registry.cc
View file

@ -9,6 +9,7 @@
using namespace ::onnxruntime::common;
namespace onnxruntime {
#if !defined(ORT_MINIMAL_BUILD)
namespace {
// Traverses the node's formal parameters and calls TraverseFn with the formal
// parameter and its associated TypeProto.
@ -199,33 +200,6 @@ bool KernelRegistry::VerifyKernelDef(const onnxruntime::Node& node,
return true;
}
Status KernelRegistry::Register(KernelDefBuilder& kernel_builder,
const KernelCreateFn& kernel_creator) {
return Register(KernelCreateInfo(kernel_builder.Build(), kernel_creator));
}
Status KernelRegistry::Register(KernelCreateInfo&& create_info) {
if (!create_info.kernel_def) {
return Status(ONNXRUNTIME, FAIL, "kernel def can't be NULL");
}
std::string key = GetMapKey(*create_info.kernel_def);
// Check op version conflicts.
auto range = kernel_creator_fn_map_.equal_range(key);
for (auto i = range.first; i != range.second; ++i) {
if (i->second.kernel_def &&
i->second.kernel_def->IsConflict(*create_info.kernel_def)) {
return Status(ONNXRUNTIME, FAIL,
"Failed to add kernel for " + key +
": Conflicting with a registered kernel with op versions.");
}
}
// Register the kernel.
// Ownership of the KernelDef is transferred to the map.
kernel_creator_fn_map_.emplace(key, std::move(create_info));
return Status::OK();
}
Status KernelRegistry::TryCreateKernel(const onnxruntime::Node& node,
const IExecutionProvider& execution_provider,
const std::unordered_map<int, OrtValue>& constant_initialized_tensors,
@ -253,37 +227,100 @@ static std::string ToString(const std::vector<std::string>& error_strs) {
return ostr.str();
}
Status KernelRegistry::TryFindKernel(const onnxruntime::Node& node,
onnxruntime::ProviderType exec_provider,
const KernelCreateInfo** out) const {
return TryFindKernel(node, exec_provider, uint64_t(0), out);
}
#endif // !defined(ORT_MINIMAL_BUILD)
// It's often this function returns a failed status, but it is totally expected.
// It just means this registry doesn't have such a kernel, please search it elsewhere.
// if this function is called before graph partition, then node.provider is not set.
// In this case, the kernel's provider must equal to exec_provider
// otherwise, kernel_def.provider must equal to node.provider. exec_provider is ignored.
Status KernelRegistry::TryFindKernel(const onnxruntime::Node& node,
onnxruntime::ProviderType exec_provider, const KernelCreateInfo** out) const {
onnxruntime::ProviderType exec_provider,
uint64_t kernel_def_hash,
const KernelCreateInfo** out) const {
const auto& node_provider = node.GetExecutionProviderType();
const auto& expected_provider = (node_provider.empty() ? exec_provider : node_provider);
auto range = kernel_creator_fn_map_.equal_range(GetMapKey(node.OpType(), node.Domain(), expected_provider));
std::vector<std::string> verify_kernel_def_error_strs;
for (auto i = range.first; i != range.second; ++i) {
std::string error_str;
if (VerifyKernelDef(node, *i->second.kernel_def, error_str)) {
*out = &i->second;
return Status::OK();
}
verify_kernel_def_error_strs.push_back(error_str);
}
*out = nullptr;
if (!verify_kernel_def_error_strs.empty()) {
// if we have a hash (ORT format model) use only that.
if (kernel_def_hash != 0) {
for (auto i = range.first; i != range.second; ++i) {
if (i->second.kernel_def->GetHash() == kernel_def_hash) {
*out = &i->second;
return Status::OK();
}
}
std::ostringstream oss;
oss << "Op with name (" << node.Name() << ")"
<< " and type (" << node.OpType() << ")"
<< " kernel is not supported in " << expected_provider << "."
<< " Encountered following errors: (" << ToString(verify_kernel_def_error_strs) << ")";
<< " kernel not found in " << expected_provider << "."
<< " No matching hash for " << kernel_def_hash;
return Status(ONNXRUNTIME, FAIL, oss.str());
}
#if !defined(ORT_MINIMAL_BUILD)
else {
std::vector<std::string> verify_kernel_def_error_strs;
for (auto i = range.first; i != range.second; ++i) {
std::string error_str;
if (VerifyKernelDef(node, *i->second.kernel_def, error_str)) {
*out = &i->second;
return Status::OK();
}
verify_kernel_def_error_strs.push_back(error_str);
}
if (!verify_kernel_def_error_strs.empty()) {
std::ostringstream oss;
oss << "Op with name (" << node.Name() << ")"
<< " and type (" << node.OpType() << ")"
<< " kernel is not supported in " << expected_provider << "."
<< " Encountered following errors: (" << ToString(verify_kernel_def_error_strs) << ")";
return Status(ONNXRUNTIME, FAIL, oss.str());
}
}
return Status(ONNXRUNTIME, FAIL, "Kernel not found");
#else
ORT_THROW("Kernel hash must be provided in minimal build.")
#endif
}
Status KernelRegistry::Register(KernelDefBuilder& kernel_builder,
const KernelCreateFn& kernel_creator) {
return Register(KernelCreateInfo(kernel_builder.Build(), kernel_creator));
}
Status KernelRegistry::Register(KernelCreateInfo&& create_info) {
if (!create_info.kernel_def) {
return Status(ONNXRUNTIME, FAIL, "kernel def can't be NULL");
}
std::string key = GetMapKey(*create_info.kernel_def);
// Check op version conflicts.
auto range = kernel_creator_fn_map_.equal_range(key);
for (auto i = range.first; i != range.second; ++i) {
if (i->second.kernel_def &&
i->second.kernel_def->IsConflict(*create_info.kernel_def)) {
return Status(ONNXRUNTIME, FAIL,
"Failed to add kernel for " + key +
": Conflicting with a registered kernel with op versions.");
}
}
// Register the kernel.
// Ownership of the KernelDef is transferred to the map.
kernel_creator_fn_map_.emplace(key, std::move(create_info));
return Status::OK();
}
} // namespace onnxruntime

22
onnxruntime/core/framework/kernel_registry_manager.cc
View file

@ -3,10 +3,13 @@
#include "core/framework/kernel_registry_manager.h"
#include "core/framework/kernel_registry.h"
#include "core/framework/customregistry.h"
#include "core/framework/execution_providers.h"
#include "core/framework/session_state.h"
#if !defined(ORT_MINIMAL_BUILD)
#include "core/framework/customregistry.h"
#endif
using namespace ONNX_NAMESPACE;
using namespace ::onnxruntime::common;
namespace onnxruntime {
@ -42,6 +45,7 @@ Status KernelRegistryManager::RegisterKernels(const ExecutionProviders& executio
return Status::OK();
}
#if !defined(ORT_MINIMAL_BUILD)
void KernelRegistryManager::RegisterKernelRegistry(std::shared_ptr<KernelRegistry> kernel_registry) {
if (nullptr == kernel_registry) {
return;
@ -58,6 +62,12 @@ bool KernelRegistryManager::HasImplementationOf(const KernelRegistryManager& r,
Status KernelRegistryManager::SearchKernelRegistry(const onnxruntime::Node& node,
/*out*/ const KernelCreateInfo** kernel_create_info) const {
return SearchKernelRegistry(node, uint64_t(0), kernel_create_info);
}
#endif
Status KernelRegistryManager::SearchKernelRegistry(const onnxruntime::Node& node, uint64_t kernel_def_hash,
/*out*/ const KernelCreateInfo** kernel_create_info) const {
Status status;
auto create_error_message = [&node, &status](const std::string& prefix) {
@ -74,10 +84,12 @@ Status KernelRegistryManager::SearchKernelRegistry(const onnxruntime::Node& node
return Status(ONNXRUNTIME, FAIL, create_error_message("The node is not placed on any Execution Provider. "));
}
#if !defined(ORT_MINIMAL_BUILD)
for (auto& registry : custom_kernel_registries_) {
status = registry->TryFindKernel(node, std::string(), kernel_create_info);
status = registry->TryFindKernel(node, std::string(), kernel_def_hash, kernel_create_info);
if (status.IsOK()) return status;
}
#endif
KernelRegistry* p = nullptr;
auto iter = provider_type_to_registry_.find(ptype);
@ -86,8 +98,10 @@ Status KernelRegistryManager::SearchKernelRegistry(const onnxruntime::Node& node
}
if (p != nullptr) {
status = p->TryFindKernel(node, std::string(), kernel_create_info);
if (status.IsOK()) return status;
status = p->TryFindKernel(node, std::string(), kernel_def_hash, kernel_create_info);
if (status.IsOK()) {
return status;
}
}
return Status(ONNXRUNTIME, NOT_IMPLEMENTED, create_error_message("Failed to find kernel for "));

17
onnxruntime/core/framework/kernel_registry_manager.h
View file

@ -8,7 +8,6 @@
#include <unordered_map>
#include "core/common/status.h"
#include "core/graph/graph_viewer.h"
#include "core/framework/customregistry.h"
#include "core/platform/ort_mutex.h"
namespace onnxruntime {
@ -33,6 +32,8 @@ class KernelRegistryManager {
// Register kernels from providers
Status RegisterKernels(const ExecutionProviders& execution_providers) ORT_MUST_USE_RESULT;
#if !defined(ORT_MINIMAL_BUILD)
// The registry passed in this function has highest priority than anything already in this KernelRegistryManager,
// and anything registered from RegisterKernels
// For example, if you do:
@ -47,10 +48,6 @@ class KernelRegistryManager {
Status SearchKernelRegistry(const onnxruntime::Node& node,
/*out*/ const KernelCreateInfo** kernel_create_info) const;
std::unique_ptr<OpKernel> CreateKernel(const onnxruntime::Node& node, const IExecutionProvider& execution_provider,
const SessionState& session_state,
const KernelCreateInfo& kernel_create_info) const ORT_MUST_USE_RESULT;
/**
* Whether this node can be run on this provider
*/
@ -71,6 +68,16 @@ class KernelRegistryManager {
if (iter != provider_type_to_registry_.end()) result.push_back(iter->second.get());
return result;
}
#endif
Status SearchKernelRegistry(const onnxruntime::Node& node,
uint64_t kernel_def_hash,
/*out*/ const KernelCreateInfo** kernel_create_info) const;
std::unique_ptr<OpKernel> CreateKernel(const onnxruntime::Node& node,
const IExecutionProvider& execution_provider,
const SessionState& session_state,
const KernelCreateInfo& kernel_create_info) const ORT_MUST_USE_RESULT;
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(KernelRegistryManager);

23
onnxruntime/core/framework/op_node_proto_helper.cc
View file

@ -1,16 +1,13 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/common/exceptions.h"
#include "core/common/status.h"
#include "core/common/logging/logging.h"
#include "core/graph/graph_viewer.h"
#include "core/graph/op.h"
#include "onnx/defs/schema.h"
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
#include "core/framework/tensorprotoutils.h"
#include "onnx/defs/schema.h"
#include "core/graph/onnx_protobuf.h"
#include "core/graph/op.h"
#include "gsl/gsl"
using namespace ONNX_NAMESPACE;
using namespace ::onnxruntime::common;
namespace onnxruntime {
@ -49,7 +46,7 @@ inline bool HasTyped<GraphProto>(const AttributeProto* attr) {
if (!attr) { \
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "No attribute with name:'", name, "'is defined."); \
} \
if (!HasTyped<T>(attr)) { \
if (!HasTyped<T>(attr)) { \
return Status(ONNXRUNTIME, FAIL, "Attibute name and type don't match"); \
} else { \
*value = static_cast<T>(attr->type()); \
@ -87,6 +84,7 @@ inline bool HasTyped<GraphProto>(const AttributeProto* attr) {
return Status::OK(); \
}
#if !defined(ORT_MINIMAL_BUILD)
#define ORT_DEFINE_GET_ATTR_SPECIALIZATIONS(type, list) \
ORT_DEFINE_GET_ATTR(ProtoHelperNodeContext, type, list) \
ORT_DEFINE_GET_ATTR(InferenceContext, type, list)
@ -94,6 +92,13 @@ inline bool HasTyped<GraphProto>(const AttributeProto* attr) {
#define ORT_DEFINE_GET_ATTRS_SPECIALIZATIONS(type, list) \
ORT_DEFINE_GET_ATTRS(ProtoHelperNodeContext, type, list) \
ORT_DEFINE_GET_ATTRS(InferenceContext, type, list)
#else
#define ORT_DEFINE_GET_ATTR_SPECIALIZATIONS(type, list) \
ORT_DEFINE_GET_ATTR(ProtoHelperNodeContext, type, list)
#define ORT_DEFINE_GET_ATTRS_SPECIALIZATIONS(type, list) \
ORT_DEFINE_GET_ATTRS(ProtoHelperNodeContext, type, list)
#endif
ORT_DEFINE_GET_ATTR_SPECIALIZATIONS(float, f)
ORT_DEFINE_GET_ATTR_SPECIALIZATIONS(int64_t, i)
@ -172,6 +177,8 @@ bool OpNodeProtoHelper<Impl_t>::HasPrimitiveAttribute(AttributeProto_AttributeTy
}
template class OpNodeProtoHelper<ProtoHelperNodeContext>;
#if !defined(ORT_MINIMAL_BUILD)
template class OpNodeProtoHelper<InferenceContext>;
#endif
} // namespace onnxruntime

21
onnxruntime/core/framework/provider_bridge_ort.cc
View file

@ -4,24 +4,27 @@
// This is the Onnxruntime side of the bridge to allow providers to be built as a DLL
// It implements onnxruntime::ProviderHost
#include "core/framework/data_types.h"
#include "core/framework/allocatormgr.h"
#include "core/framework/compute_capability.h"
#include "core/framework/data_types.h"
#include "core/framework/data_transfer_manager.h"
#include "core/framework/execution_provider.h"
#include "core/framework/kernel_registry.h"
#include "core/graph/model.h"
#include "core/platform/env.h"
#include "core/providers/dnnl/dnnl_provider_factory.h"
#include "core/session/inference_session.h"
#include "core/providers/tensorrt/tensorrt_provider_factory.h"
#include "core/session/abi_session_options_impl.h"
#include "core/session/ort_apis.h"
#ifdef USE_TENSORRT
#include "core/providers/tensorrt/tensorrt_provider_factory.h"
#include "core/providers/cuda/cuda_allocator.h"
#include "core/providers/cuda/gpu_data_transfer.h"
#include "core/providers/cuda/math/unary_elementwise_ops_impl.h"
#include "core/providers/cuda/cuda_common.h"
#endif
#include "core/session/abi_session_options_impl.h"
#include "core/session/ort_apis.h"
#include "core/platform/env.h"
#include "core/graph/model.h"
#include "core/framework/data_transfer_manager.h"
#include "core/framework/compute_capability.h"
#include "core/framework/execution_provider.h"
#define PROVIDER_BRIDGE_ORT
#include "core/providers/shared_library/provider_interfaces.h"
#include "onnx/common/stl_backports.h"

2
onnxruntime/core/graph/contrib_ops/contrib_defs.cc
View file

@ -7,8 +7,8 @@
#include "core/graph/contrib_ops/contrib_defs.h"
#include "core/graph/contrib_ops/nchwc_schema_defs.h"
#include "core/graph/contrib_ops/range_schema_defs.h"
#include "core/graph/onnx_protobuf.h"
#include "core/graph/op.h"
#include "onnx/defs/schema.h"
#include "onnx/defs/shape_inference.h"
#include "onnx/defs/tensor_proto_util.h"
#include "core/mlas/inc/mlas.h"

4
onnxruntime/core/graph/featurizers_ops/featurizers_defs.cc
View file

@ -3,11 +3,9 @@
#include "core/graph/constants.h"
#include "core/graph/featurizers_ops/featurizers_defs.h"
#include "core/graph/onnx_protobuf.h"
#include "core/graph/op.h"
#include "onnx/defs/schema.h"
#include "onnx/defs/shape_inference.h"
#define MS_FEATURIZERS_OPERATOR_SCHEMA(name) MS_FEATURIZERS_OPERATOR_SCHEMA_UNIQ_HELPER(__COUNTER__, name)
#define MS_FEATURIZERS_OPERATOR_SCHEMA_UNIQ_HELPER(Counter, name) MS_FEATURIZERS_OPERATOR_SCHEMA_UNIQ(Counter, name)

29
onnxruntime/core/graph/graph.cc
View file

@ -21,12 +21,12 @@
#include "core/graph/graph_viewer.h"
#include "core/graph/indexed_sub_graph.h"
#include "core/graph/model.h"
#include "core/graph/schema_registry.h"
#include "core/graph/op.h"
#if !defined(ORT_MINIMAL_BUILD)
#include "core/graph/function.h"
#include "core/graph/function_impl.h"
#include "core/graph/schema_registry.h"
#include "onnx/checker.h"
using namespace ONNX_NAMESPACE::checker;
#endif
@ -1580,7 +1580,7 @@ class InferenceContextImpl : public ONNX_NAMESPACE::InferenceContext {
// only return data if it's for a constant initializer. checks for outer scope initializers
// if this is a subgraph and the name isn't found locally.
const TensorProto* initializer = graph_utils::GetConstantInitializer(graph_, def->Name(), true);
const TensorProto* initializer = graph_.GetConstantInitializer(def->Name(), true);
return initializer;
}
@ -2392,6 +2392,31 @@ void Graph::CleanAllInitializedTensors() noexcept {
}
}
const ONNX_NAMESPACE::TensorProto* Graph::GetConstantInitializer(const std::string& initializer_name,
bool check_outer_scope) const {
const ONNX_NAMESPACE::TensorProto* initializer = nullptr;
if (GetInitializedTensor(initializer_name, initializer)) {
if (CanOverrideInitializer()) {
const auto& graph_inputs = GetInputsIncludingInitializers();
bool is_constant = std::none_of(graph_inputs.cbegin(), graph_inputs.cend(),
[&initializer_name](const NodeArg* input) {
return input->Name() == initializer_name;
});
if (!is_constant) {
initializer = nullptr;
}
}
} else if (check_outer_scope && IsSubgraph()) {
// make sure there's not a local value with the same name. if there is it shadows any initializer in outer scope.
if (IsOuterScopeValue(initializer_name)) {
initializer = parent_graph_->GetConstantInitializer(initializer_name, check_outer_scope);
}
}
return initializer;
}
#if !defined(ORT_MINIMAL_BUILD)
void Graph::AddValueInfo(const NodeArg* new_value_info) {
for (const auto* info : value_info_) {

22
onnxruntime/core/graph/graph_utils.cc
View file

@ -489,27 +489,7 @@ bool IsGraphInput(const Graph& graph, const NodeArg* input) {
const ONNX_NAMESPACE::TensorProto* GetConstantInitializer(const Graph& graph, const std::string& initializer_name,
bool check_outer_scope) {
const ONNX_NAMESPACE::TensorProto* initializer = nullptr;
if (graph.GetInitializedTensor(initializer_name, initializer)) {
if (graph.CanOverrideInitializer()) {
const auto& graph_inputs = graph.GetInputsIncludingInitializers();
bool is_constant = std::none_of(graph_inputs.cbegin(), graph_inputs.cend(),
[&initializer_name](const NodeArg* input) {
return input->Name() == initializer_name;
});
if (!is_constant) {
initializer = nullptr;
}
}
} else if (check_outer_scope && graph.IsSubgraph()) {
// make sure there's not a local value with the same name. if there is it shadows any initializer in outer scope.
if (graph.IsOuterScopeValue(initializer_name)) {
initializer = GetConstantInitializer(*graph.ParentGraph(), initializer_name, check_outer_scope);
}
}
return initializer;
return graph.GetConstantInitializer(initializer_name, check_outer_scope);
}
bool IsInitializer(const Graph& graph, const std::string& name, bool check_outer_scope) {

4
onnxruntime/core/graph/graph_viewer.cc
View file

@ -8,8 +8,6 @@
#include "core/graph/graph_viewer.h"
#include "core/graph/graph_utils.h"
namespace onnxruntime {
bool NodeCompare::operator()(const Node* n1, const Node* n2) const {
@ -115,7 +113,7 @@ bool GraphViewer::IsSubgraph() const {
}
bool GraphViewer::IsConstantInitializer(const std::string& name, bool check_outer_scope) const {
return graph_utils::IsConstantInitializer(*graph_, name, check_outer_scope);
return graph_->GetConstantInitializer(name, check_outer_scope) != nullptr;
}
} // namespace onnxruntime

3
onnxruntime/core/graph/model.cc
View file

@ -20,7 +20,10 @@
#include "gsl/gsl"
#include "core/platform/env.h"
#if !defined(ORT_MINIMAL_BUILD)
#include "core/graph/schema_registry.h"
#endif
using namespace ONNX_NAMESPACE;
using namespace onnxruntime;

13
onnxruntime/core/optimizer/constant_folding.cc
View file

@ -12,6 +12,14 @@ using namespace onnxruntime::common;
namespace onnxruntime {
ConstantFolding::ConstantFolding(const IExecutionProvider& execution_provider,
const std::unordered_set<std::string>& compatible_execution_providers,
const std::unordered_set<std::string>& excluded_initializers) noexcept
: GraphTransformer("ConstantFolding", compatible_execution_providers),
excluded_initializers_(excluded_initializers),
execution_provider_(execution_provider) {
}
// We need to handle a Shape node separately as the input doesn't need to be a constant initializer for
// Shape to be able to be constant folded.
static bool ConstantFoldShapeNode(Graph& graph, Node& node) {
@ -96,11 +104,8 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level,
}
// Create execution frame for executing constant nodes.
std::unique_ptr<CPUExecutionProvider> cpu_execution_provider =
onnxruntime::make_unique<CPUExecutionProvider>(CPUExecutionProviderInfo());
// Create execution frame for executing constant nodes.
OptimizerExecutionFrame::Info info({node}, constant_inputs, std::move(cpu_execution_provider));
OptimizerExecutionFrame::Info info({node}, constant_inputs, execution_provider_);
std::vector<int> fetch_mlvalue_idxs;
for (const auto* node_out : node->OutputDefs()) {

15
onnxruntime/core/optimizer/constant_folding.h
View file

@ -5,6 +5,8 @@
#include "core/optimizer/graph_transformer.h"
#include "core/framework/ml_value.h"
#include <memory>
#include "core/framework/execution_provider.h"
namespace onnxruntime {
@ -16,16 +18,19 @@ it statically computes parts of the graph that rely only on constant initializer
*/
class ConstantFolding : public GraphTransformer {
public:
/** Constant folding will not be applied to nodes that have one of initializers from excluded_initializers as input.
For pre-training, the trainable weights are those initializers to be excluded. */
ConstantFolding(const std::unordered_set<std::string>& compatible_execution_providers = {},
const std::unordered_set<std::string>& excluded_initializers = {}) noexcept
: GraphTransformer("ConstantFolding", compatible_execution_providers), excluded_initializers_(excluded_initializers) {}
/*! Constant folding will not be applied to nodes that have one of initializers from excluded_initializers as input.
For pre-training, the trainable weights are those initializers to be excluded.
\param execution_provider Execution provider instance to execute constant folding.
*/
ConstantFolding(const IExecutionProvider& execution_provider,
const std::unordered_set<std::string>& compatible_execution_providers = {},
const std::unordered_set<std::string>& excluded_initializers = {}) noexcept;
private:
Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override;
const std::unordered_set<std::string> excluded_initializers_;
const IExecutionProvider& execution_provider_;
};
} // namespace onnxruntime

4
onnxruntime/core/optimizer/graph_transformer_utils.cc
View file

@ -36,6 +36,7 @@
#include "core/optimizer/unsqueeze_elimination.h"
namespace onnxruntime {
class IExecutionProvider;
namespace optimizer_utils {
@ -107,6 +108,7 @@ std::unique_ptr<RuleBasedGraphTransformer> GenerateRuleBasedGraphTransformer(Tra
std::vector<std::unique_ptr<GraphTransformer>> GenerateTransformers(TransformerLevel level,
gsl::span<const FreeDimensionOverride> free_dimension_overrides,
const IExecutionProvider& execution_provider, /*required by constant folding*/
const std::vector<std::string>& transformers_and_rules_to_enable) {
std::vector<std::unique_ptr<GraphTransformer>> transformers;
std::unique_ptr<RuleBasedGraphTransformer> rule_transformer = nullptr;
@ -115,7 +117,7 @@ std::vector<std::unique_ptr<GraphTransformer>> GenerateTransformers(TransformerL
std::unordered_set<std::string> l1_execution_providers = {};
transformers.emplace_back(onnxruntime::make_unique<CommonSubexpressionElimination>(l1_execution_providers));
transformers.emplace_back(onnxruntime::make_unique<ConstantFolding>(l1_execution_providers));
transformers.emplace_back(onnxruntime::make_unique<ConstantFolding>(execution_provider, l1_execution_providers));
transformers.emplace_back(onnxruntime::make_unique<MatMulAddFusion>(l1_execution_providers));
transformers.emplace_back(onnxruntime::make_unique<ReshapeFusion>(l1_execution_providers));
transformers.emplace_back(onnxruntime::make_unique<FreeDimensionOverrideTransformer>(free_dimension_overrides));

12
onnxruntime/core/optimizer/optimizer_execution_frame.cc
View file

@ -19,11 +19,9 @@ namespace onnxruntime {
OptimizerExecutionFrame::Info::Info(const std::vector<const Node*>& nodes,
const InitializedTensorSet& initialized_tensor_set,
std::unique_ptr<CPUExecutionProvider> cpu_execution_provider) {
ORT_ENFORCE(cpu_execution_provider, "Provided CPU execution provider is a nullptr");
cpu_execution_provider_ = std::move(cpu_execution_provider);
allocator_ptr_ = cpu_execution_provider_->GetAllocator(device_id_, mem_type_);
const IExecutionProvider& execution_provider)
: execution_provider_(execution_provider) {
allocator_ptr_ = execution_provider_.GetAllocator(device_id_, mem_type_);
ORT_ENFORCE(allocator_ptr_, "Failed to get allocator for optimizer");
data_transfer_mgr_.RegisterDataTransfer(onnxruntime::make_unique<CPUDataTransfer>());
@ -67,8 +65,8 @@ OptimizerExecutionFrame::Info::Info(const std::vector<const Node*>& nodes,
std::unique_ptr<const OpKernel> OptimizerExecutionFrame::Info::CreateKernel(const Node* node) const {
std::unique_ptr<OpKernel> op_kernel;
std::shared_ptr<KernelRegistry> kernel_registry = cpu_execution_provider_->GetKernelRegistry();
auto status = kernel_registry->TryCreateKernel(*node, *cpu_execution_provider_, initializers_,
std::shared_ptr<KernelRegistry> kernel_registry = execution_provider_.GetKernelRegistry();
auto status = kernel_registry->TryCreateKernel(*node, execution_provider_, initializers_,
ort_value_name_idx_map_, FuncManager(), data_transfer_mgr_,
op_kernel);

6
onnxruntime/core/optimizer/optimizer_execution_frame.h
View file

@ -21,14 +21,14 @@ class OptimizerExecutionFrame final : public IExecutionFrame {
class Info {
public:
Info(const std::vector<const Node*>& nodes, const InitializedTensorSet& initialized_tensor_set,
std::unique_ptr<CPUExecutionProvider> cpu_execution_provider);
const IExecutionProvider& execution_provider);
~Info() {
for (auto& kvp : deleter_for_initialized_tensors_) {
kvp.second.f(kvp.second.param);
}
}
AllocatorPtr GetAllocator(const OrtMemoryInfo& info) const {
return cpu_execution_provider_->GetAllocator(info.id, info.mem_type);
return execution_provider_.GetAllocator(info.id, info.mem_type);
}
AllocatorPtr GetAllocator() const {
@ -55,7 +55,6 @@ class OptimizerExecutionFrame final : public IExecutionFrame {
private:
// The optimizer is running on CPU execution provider by default.
std::unique_ptr<CPUExecutionProvider> cpu_execution_provider_;
const int device_id_{0};
const OrtMemType mem_type_{OrtMemTypeDefault};
AllocatorPtr allocator_ptr_;
@ -69,6 +68,7 @@ class OptimizerExecutionFrame final : public IExecutionFrame {
// munmap memory region and close file descriptor
std::unordered_map<int, OrtCallback> deleter_for_initialized_tensors_;
std::unique_ptr<NodeIndexInfo> node_index_info_;
const IExecutionProvider& execution_provider_;
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(Info);
};

12
onnxruntime/core/providers/cpu/nn/string_normalizer.cc
View file

@ -2,14 +2,13 @@
// Licensed under the MIT License.
#include "string_normalizer.h"
#include "onnx/defs/schema.h"
#include "core/common/common.h"
#include "core/framework/tensor.h"
#ifdef _MSC_VER
#include <codecvt>
#include <locale.h>
#elif defined (__APPLE__) or defined (__ANDROID__)
#elif defined(__APPLE__) or defined(__ANDROID__)
#include <codecvt>
#else
#include <limits>
@ -77,7 +76,7 @@ using Utf8Converter = std::wstring_convert<std::codecvt_utf8<wchar_t>>;
const std::string default_locale("en-US");
#else // MS_VER
#else // MS_VER
class Locale {
public:
@ -112,7 +111,6 @@ using Utf8Converter = std::wstring_convert<std::codecvt_utf8<wchar_t>>;
// All others (Linux)
class Utf8Converter {
public:
Utf8Converter(const std::string&, const std::wstring&) {}
std::wstring from_bytes(const std::string& s) const {
@ -182,11 +180,11 @@ class Utf8Converter {
}
};
#endif // __APPLE__
#endif // __APPLE__
const std::string default_locale("en_US.UTF-8"); // All non-MS
const std::string default_locale("en_US.UTF-8"); // All non-MS
#endif // MS_VER
#endif // MS_VER
template <class ForwardIter>
Status CopyCaseAction(ForwardIter first, ForwardIter end, OpKernelContext* ctx,

1
onnxruntime/core/providers/cpu/nn/tfidfvectorizer.cc
View file

@ -2,7 +2,6 @@
// Licensed under the MIT License.
#include "tfidfvectorizer.h"
#include "onnx/defs/schema.h"
#include "core/common/common.h"
#include "core/framework/tensor.h"
#include "core/platform/threadpool.h"

1
onnxruntime/core/providers/cpu/tensor/reverse_sequence.cc
View file

@ -2,7 +2,6 @@
// Licensed under the MIT License.
#include "reverse_sequence.h"
#include "onnx/defs/schema.h"
// there's no way to use a raw pointer as the copy destination with std::copy_n
// (which gsl::copy uses with span::data() which returns a raw pointer) with the 14.11 toolset

3
onnxruntime/core/session/allocator_impl.h
View file

@ -4,6 +4,8 @@
#pragma once
#include "core/session/onnxruntime_c_api.h"
#include "core/framework/allocator.h"
#include "core/framework/arena.h"
#include "core/session/device_allocator.h"
namespace onnxruntime {
class AllocatorWrapper : public IAllocator {
@ -19,4 +21,5 @@ class AllocatorWrapper : public IAllocator {
private:
OrtAllocator* impl_;
};
} // namespace onnxruntime

6
onnxruntime/core/session/custom_ops.cc
View file

@ -7,7 +7,6 @@
#include "core/graph/onnx_protobuf.h"
#include "core/session/inference_session.h"
#include "core/session/ort_apis.h"
#include "core/framework/customregistry.h"
#include "core/framework/data_types.h"
#include "core/framework/op_kernel_info.h"
#include "core/framework/op_kernel_context_internal.h"
@ -68,6 +67,9 @@ ORT_API_STATUS_IMPL(OrtApis::KernelInfoGetAttribute_string, _In_ const OrtKernel
return onnxruntime::ToOrtStatus(status);
}
#if !defined(ORT_MINIMAL_BUILD)
#include "core/framework/customregistry.h"
namespace onnxruntime {
struct CustomOpKernel : OpKernel {
@ -149,3 +151,5 @@ common::Status CreateCustomRegistry(const std::vector<OrtCustomOpDomain*>& op_do
}
} // namespace onnxruntime
#endif // !defined(ORT_MINIMAL_BUILD)

117
onnxruntime/core/session/device_allocator.cc
View file

@ -1,47 +1,15 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/framework/allocator.h"
#include "core/framework/utils.h"
#include "device_allocator.h"
#include "core/session/inference_session.h"
#include "core/session/onnxruntime_cxx_api.h"
#include "core/session/ort_apis.h"
#include <assert.h>
struct OrtAllocatorForDevice : public OrtAllocator {
explicit OrtAllocatorForDevice(onnxruntime::AllocatorPtr&& dev_allocator)
: device_allocator_(std::move(dev_allocator)) {
OrtAllocator::version = ORT_API_VERSION;
OrtAllocator::Alloc = [](OrtAllocator* this_, size_t size) { return static_cast<OrtAllocatorForDevice*>(this_)->Alloc(size); };
OrtAllocator::Free = [](OrtAllocator* this_, void* p) { static_cast<OrtAllocatorForDevice*>(this_)->Free(p); };
OrtAllocator::Info = [](const OrtAllocator* this_) { return static_cast<const OrtAllocatorForDevice*>(this_)->Info(); };
}
~OrtAllocatorForDevice() = default;
void* Alloc(size_t size) const {
return device_allocator_->Alloc(size);
}
void Free(void* p) const {
device_allocator_->Free(p);
}
const OrtMemoryInfo* Info() const {
return &device_allocator_->Info();
}
OrtAllocatorForDevice(const OrtAllocatorForDevice&) = delete;
OrtAllocatorForDevice& operator=(const OrtAllocatorForDevice&) = delete;
private:
onnxruntime::AllocatorPtr device_allocator_;
};
#include "core/session/ort_env.h"
#include "core/session/allocator_impl.h"
#define API_IMPL_BEGIN try {
#define API_IMPL_END \
} \
catch (const std::exception& ex) { \
catch (const std::exception& ex) { \
return OrtApis::CreateStatus(ORT_RUNTIME_EXCEPTION, ex.what()); \
}
@ -52,11 +20,82 @@ ORT_API_STATUS_IMPL(OrtApis::CreateAllocator, const OrtSession* sess, const OrtM
if (!allocator_ptr) {
return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "No requested allocator available");
}
*out = new OrtAllocatorForDevice(std::move(allocator_ptr));
*out = new onnxruntime::OrtAllocatorForDevice(std::move(allocator_ptr));
return nullptr;
API_IMPL_END
}
ORT_API(void, OrtApis::ReleaseAllocator, _Frees_ptr_opt_ OrtAllocator* allocator) {
delete reinterpret_cast<OrtAllocatorForDevice*>(allocator);
ORT_API_STATUS_IMPL(OrtApis::CreateAndRegisterAllocator, _Inout_ OrtEnv* env, _In_ const OrtMemoryInfo* mem_info,
_In_ const OrtArenaCfg* arena_cfg) {
using namespace onnxruntime;
if (!env) {
return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Env is null");
}
if (!mem_info) {
return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "OrtMemoryInfo is null");
}
// TODO should we allow sharing of non-CPU allocators?
if (mem_info->device.Type() != OrtDevice::CPU) {
return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Only CPU devices are supported for now.");
}
// determine if arena should be used
bool create_arena = mem_info->alloc_type == OrtArenaAllocator;
#ifdef USE_JEMALLOC
#if defined(USE_MIMALLOC_ARENA_ALLOCATOR) || defined(USE_MIMALLOC_STL_ALLOCATOR)
#error jemalloc and mimalloc should not both be enabled
#endif
//JEMalloc already has memory pool, so just use device allocator.
create_arena = false;
#elif !(defined(__amd64__) || defined(_M_AMD64))
//Disable Arena allocator for x86_32 build because it may run into infinite loop when integer overflow happens
create_arena = false;
#endif
AllocatorPtr allocator_ptr;
size_t max_mem = std::numeric_limits<size_t>::max();
// create appropriate DeviceAllocatorRegistrationInfo and allocator based on create_arena
if (create_arena) {
ArenaExtendStrategy arena_extend_strategy = BFCArena::DEFAULT_ARENA_EXTEND_STRATEGY;
int initial_chunk_size_bytes = BFCArena::DEFAULT_INITIAL_CHUNK_SIZE_BYTES;
int max_dead_bytes_per_chunk = BFCArena::DEFAULT_MAX_DEAD_BYTES_PER_CHUNK;
if (arena_cfg) {
if (arena_cfg->max_mem != -1) max_mem = arena_cfg->max_mem;
if (arena_cfg->arena_extend_strategy == 0) {
arena_extend_strategy = ArenaExtendStrategy::kNextPowerOfTwo;
} else if (arena_cfg->arena_extend_strategy == 1) {
arena_extend_strategy = ArenaExtendStrategy::kSameAsRequested;
}
if (arena_cfg->initial_chunk_size_bytes != -1) initial_chunk_size_bytes = arena_cfg->initial_chunk_size_bytes;
if (arena_cfg->max_dead_bytes_per_chunk != -1) max_dead_bytes_per_chunk = arena_cfg->max_dead_bytes_per_chunk;
}
DeviceAllocatorRegistrationInfo device_info{
OrtMemTypeDefault,
[mem_info](int) { return onnxruntime::make_unique<TAllocator>(*mem_info); },
max_mem,
arena_extend_strategy,
initial_chunk_size_bytes,
max_dead_bytes_per_chunk};
allocator_ptr = CreateAllocator(device_info, 0, create_arena);
} else {
DeviceAllocatorRegistrationInfo device_info{OrtMemTypeDefault,
[](int) { return onnxruntime::make_unique<TAllocator>(); },
max_mem};
allocator_ptr = CreateAllocator(device_info, 0, create_arena);
}
auto st = env->RegisterAllocator(allocator_ptr);
if (!st.IsOK()) {
return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, st.ErrorMessage().c_str());
}
return nullptr;
}
ORT_API(void, OrtApis::ReleaseAllocator, _Frees_ptr_opt_ OrtAllocator* allocator) {
delete reinterpret_cast<onnxruntime::OrtAllocatorForDevice*>(allocator);
}

45
onnxruntime/core/session/device_allocator.h Normal file
View file

@ -0,0 +1,45 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/framework/allocator.h"
#include "core/framework/utils.h"
#include "core/session/onnxruntime_cxx_api.h"
#include "core/session/ort_apis.h"
#include <assert.h>
#include "core/framework/allocatormgr.h"
namespace onnxruntime {
struct OrtAllocatorForDevice : public OrtAllocator {
explicit OrtAllocatorForDevice(onnxruntime::AllocatorPtr&& dev_allocator)
: device_allocator_(std::move(dev_allocator)) {
OrtAllocator::version = ORT_API_VERSION;
OrtAllocator::Alloc = [](OrtAllocator* this_, size_t size) { return static_cast<OrtAllocatorForDevice*>(this_)->Alloc(size); };
OrtAllocator::Free = [](OrtAllocator* this_, void* p) { static_cast<OrtAllocatorForDevice*>(this_)->Free(p); };
OrtAllocator::Info = [](const OrtAllocator* this_) { return static_cast<const OrtAllocatorForDevice*>(this_)->Info(); };
}
~OrtAllocatorForDevice() = default;
void* Alloc(size_t size) const {
return device_allocator_->Alloc(size);
}
void Free(void* p) const {
device_allocator_->Free(p);
}
const OrtMemoryInfo* Info() const {
return &device_allocator_->Info();
}
OrtAllocatorForDevice(const OrtAllocatorForDevice&) = delete;
OrtAllocatorForDevice& operator=(const OrtAllocatorForDevice&) = delete;
onnxruntime::IAllocator* GetAllocator() {
return device_allocator_.get();
}
private:
onnxruntime::AllocatorPtr device_allocator_;
};
} // namespace onnxruntime

18
onnxruntime/core/session/environment.cc
View file

@ -5,9 +5,13 @@
#include "core/framework/allocatormgr.h"
#include "core/graph/constants.h"
#include "core/graph/op.h"
#if !defined(ORT_MINIMAL_BUILD)
#include "onnx/defs/operator_sets.h"
#include "onnx/defs/operator_sets_ml.h"
#if defined(ENABLE_TRAINING)
#include "onnx/defs/operator_sets_training.h"
#endif
#endif
#ifndef DISABLE_CONTRIB_OPS
#include "core/graph/contrib_ops/contrib_defs.h"
#endif
@ -20,6 +24,7 @@
#include "core/platform/env.h"
#include "core/util/thread_utils.h"
#include "core/session/allocator_impl.h"
#ifdef ONNXRUNTIME_ENABLE_INSTRUMENT
#include "core/platform/tracing.h"
@ -48,6 +53,19 @@ Status Environment::Create(std::unique_ptr<logging::LoggingManager> logging_mana
return status;
}
Status Environment::RegisterAllocator(AllocatorPtr allocator) {
const auto& mem_info = allocator->Info();
// We don't expect millions of allocators getting registered. Hence linear search should be fine.
auto ite = std::find_if(std::begin(shared_allocators_),
std::end(shared_allocators_),
[&mem_info](const AllocatorPtr& alloc_ptr) { return alloc_ptr->Info() == mem_info; });
if (ite != shared_allocators_.end()) {
return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Allocator with this OrtMemoryInfo is already registered.");
}
shared_allocators_.insert(ite, allocator);
return Status::OK();
}
Status Environment::Initialize(std::unique_ptr<logging::LoggingManager> logging_manager,
const OrtThreadingOptions* tp_options,
bool create_global_thread_pools) {

85
onnxruntime/core/session/inference_session.cc
View file

@ -12,43 +12,46 @@
#include <thread>
#include "core/common/logging/logging.h"
#include "core/platform/threadpool.h"
#include "core/graph/graph_viewer.h"
#include "core/graph/graph_utils.h"
#include "core/graph/model.h"
#include "core/framework/allocatormgr.h"
#include "core/framework/customregistry.h"
#include "core/session/environment.h"
#include "core/framework/error_code_helper.h"
#include "core/framework/execution_frame.h"
#include "core/framework/feeds_fetches_manager.h"
#include "core/framework/graph_partitioner.h"
#include "core/framework/kernel_def_builder.h"
#include "core/framework/kernel_registry.h"
#include "core/framework/ort_value_pattern_planner.h"
#include "core/framework/mldata_type_utils.h"
#include "core/framework/op_kernel_context_internal.h"
#include "core/framework/TensorSeq.h"
#include "core/framework/tensorprotoutils.h"
#include "core/framework/tensor_type_and_shape.h"
#include "core/framework/op_kernel_context_internal.h"
#include "core/framework/ort_value_pattern_planner.h"
#include "core/framework/utils.h"
#include "core/graph/graph_viewer.h"
#include "core/graph/graph_utils.h"
#include "core/graph/model.h"
#include "core/optimizer/transformer_memcpy.h"
#include "core/optimizer/graph_transformer.h"
#include "core/optimizer/insert_cast_transformer.h"
#include "core/optimizer/rule_based_graph_transformer.h"
#include "core/optimizer/graph_transformer_utils.h"
#include "core/platform/Barrier.h"
#include "core/platform/ort_mutex.h"
#include "core/platform/threadpool.h"
#include "core/providers/cpu/controlflow/utils.h"
#include "core/providers/cpu/cpu_execution_provider.h"
#ifdef USE_DML // TODO: This is necessary for the workaround in TransformGraph
#include "core/providers/dml/DmlExecutionProvider/src/GraphTransformer.h"
#endif
#include "core/session/environment.h"
#include "core/session/IOBinding.h"
#include "core/session/custom_ops.h"
#include "core/util/protobuf_parsing_utils.h"
#include "core/optimizer/rule_based_graph_transformer.h"
#include "core/optimizer/graph_transformer_utils.h"
#include "core/util/thread_utils.h"
#include "core/session/inference_session_utils.h"
#include "core/platform/ort_mutex.h"
#include "core/platform/Barrier.h"
#include "core/util/protobuf_parsing_utils.h"
#include "core/util/thread_utils.h"
#if !defined(ORT_MINIMAL_BUILD)
#include "core/framework/customregistry.h"
#include "core/session/custom_ops.h"
#endif
using namespace ONNX_NAMESPACE;
@ -233,7 +236,8 @@ InferenceSession::InferenceSession(const SessionOptions& session_options, const
graph_transformation_mgr_(session_options.max_num_graph_transformation_steps),
insert_cast_transformer_("CastFloat16Transformer"),
#endif
logging_manager_(session_env.GetLoggingManager()) {
logging_manager_(session_env.GetLoggingManager()),
environment_(session_env) {
// Initialize assets of this session instance
ConstructorCommon(session_options, session_env);
}
@ -244,7 +248,8 @@ InferenceSession::InferenceSession(const SessionOptions& session_options, const
: model_location_(ToWideString(model_uri)),
graph_transformation_mgr_(session_options.max_num_graph_transformation_steps),
insert_cast_transformer_("CastFloat16Transformer"),
logging_manager_(session_env.GetLoggingManager()) {
logging_manager_(session_env.GetLoggingManager()),
environment_(session_env) {
auto status = Model::Load(model_location_, model_proto_);
ORT_ENFORCE(status.IsOK(), "Given model could not be parsed while creating inference session. Error message: ",
status.ErrorMessage());
@ -259,7 +264,8 @@ InferenceSession::InferenceSession(const SessionOptions& session_options,
const std::wstring& model_uri)
: graph_transformation_mgr_(session_options.max_num_graph_transformation_steps),
insert_cast_transformer_("CastFloat16Transformer"),
logging_manager_(session_env.GetLoggingManager()) {
logging_manager_(session_env.GetLoggingManager()),
environment_(session_env) {
model_location_ = ToWideString(model_uri);
auto status = Model::Load(model_location_, model_proto_);
ORT_ENFORCE(status.IsOK(), "Given model could not be parsed while creating inference session. Error message: ",
@ -274,7 +280,8 @@ InferenceSession::InferenceSession(const SessionOptions& session_options, const
std::istream& model_istream)
: graph_transformation_mgr_(session_options.max_num_graph_transformation_steps),
insert_cast_transformer_("CastFloat16Transformer"),
logging_manager_(session_env.GetLoggingManager()) {
logging_manager_(session_env.GetLoggingManager()),
environment_(session_env) {
Status st = Model::Load(model_istream, &model_proto_);
ORT_ENFORCE(st.IsOK(), "Could not parse model successfully while constructing the inference session");
is_model_proto_parsed_ = true;
@ -286,7 +293,8 @@ InferenceSession::InferenceSession(const SessionOptions& session_options, const
const void* model_data, int model_data_len)
: graph_transformation_mgr_(session_options.max_num_graph_transformation_steps),
insert_cast_transformer_("CastFloat16Transformer"),
logging_manager_(session_env.GetLoggingManager()) {
logging_manager_(session_env.GetLoggingManager()),
environment_(session_env) {
const bool result = model_proto_.ParseFromArray(model_data, model_data_len);
ORT_ENFORCE(result, "Could not parse model successfully while constructing the inference session");
is_model_proto_parsed_ = true;
@ -675,7 +683,7 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph,
if (!node.Name().empty())
oss << node.Name() << ":";
oss << node.OpType() << "(" << node.SinceVersion() << ")";
return Status(common::ONNXRUNTIME, common::NOT_IMPLEMENTED, oss.str());
} else {
if (is_verbose_mode) { // TODO: should we disable this if the number of nodes are above a certain threshold?
@ -804,6 +812,24 @@ common::Status InferenceSession::Initialize() {
// re-acquire mutex
std::lock_guard<onnxruntime::OrtMutex> l(session_mutex_);
// At this time we know all the providers that will be part of this session.
// Read shared allocators from the environment and update them in the respective providers.
//
// The reason for updating the providers is so that when the session state is created the allocators
// are setup appropariately keyed by OrtMemoryInfo with delegates going to the respective providers.
// Secondly, the GetAllocator() method inside IExecutionProvider is still used in various places, hence
// it doesn't make sense to just update the allocator map inside session state with these shared allocators; doing
// so would cause inconsistency between the allocator map inside session sate and that inside the providers.
// TODO: we could refactor the allocators to not require the call to GetAllocator but that change is much bigger
// since we've to take into account the per-thread cuda allocators.
// TODO (contd.) We could also possibly absorb the per-thread logic in a new allocator decorator that derives
// from IAllocator to keep things clean.
std::string use_env_allocators = GetSessionConfigOrDefault(session_options_,
ORT_SESSION_OPTIONS_CONFIG_USE_ENV_ALLOCATORS, "0");
if (use_env_allocators == "1") {
UpdateProvidersWithSharedAllocators();
}
#ifdef ONNXRUNTIME_ENABLE_INSTRUMENT
TraceLoggingWriteStart(session_activity, "OrtInferenceSessionActivity");
session_activity_started_ = true;
@ -914,6 +940,19 @@ common::Status InferenceSession::Initialize() {
return status;
}
// This method should be called from within Initialize() only and before the creation of the session state.
// This ensures all providers have been registered in the session and the session state is consistent with the providers.
void InferenceSession::UpdateProvidersWithSharedAllocators() {
using namespace std;
const auto& provider_ids = execution_providers_.GetIds();
for (const auto& one_shared_alloc : environment_.GetRegisteredSharedAllocators()) {
for (const auto& id : provider_ids) {
auto* provider_ptr = execution_providers_.Get(id);
provider_ptr->ReplaceAllocator(one_shared_alloc);
}
}
}
int InferenceSession::GetCurrentNumRuns() const {
return current_num_runs_.load();
}
@ -1473,7 +1512,9 @@ void InferenceSession::AddPredefinedTransformers(GraphTransformerManager& transf
auto add_transformers = [&](TransformerLevel level) {
// Generate and register transformers for level
auto transformers_to_register =
optimizer_utils::GenerateTransformers(level, session_options_.free_dimension_overrides, custom_list);
optimizer_utils::GenerateTransformers(level, session_options_.free_dimension_overrides,
*execution_providers_.Get(onnxruntime::kCpuExecutionProvider),
custom_list);
for (auto& entry : transformers_to_register) {
transformer_manager.Register(std::move(entry), level);
}

11
onnxruntime/core/session/inference_session.h
View file

@ -335,6 +335,7 @@ class InferenceSession {
* Get all the providers' options this session was initialized with.
*/
const ProviderOptionsMap& GetAllProviderOptions() const;
/**
* Start profiling on this inference session. This simply turns on profiling events to be
* recorded. A corresponding EndProfiling has to follow to write profiling data to a file.
@ -424,6 +425,9 @@ class InferenceSession {
// The file path of where the model was loaded. e.g. /tmp/test_squeezenet/model.onnx
std::basic_string<ORTCHAR_T> model_location_;
// The list of execution providers.
ExecutionProviders execution_providers_;
private:
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(InferenceSession);
@ -467,6 +471,9 @@ class InferenceSession {
template <typename T>
void StartProfiling(const std::basic_string<T>& file_prefix);
// Updates all providers with the allocators from the env based on OrtMemoryInfo
void UpdateProvidersWithSharedAllocators();
#if !defined(ORT_MINIMAL_BUILD)
virtual void AddPredefinedTransformers(GraphTransformerManager& transformer_manager,
TransformerLevel graph_optimization_level,
@ -499,9 +506,6 @@ class InferenceSession {
// Profiler for this session.
profiling::Profiler session_profiler_;
// The list of execution providers.
ExecutionProviders execution_providers_;
// Immutable state for each op in the model. Shared by all executors.
// It has a dependency on execution_providers_.
std::unique_ptr<SessionState> session_state_;
@ -588,6 +592,7 @@ class InferenceSession {
// Flag indicating if ModelProto has been parsed in an applicable ctor
bool is_model_proto_parsed_ = false;
const Environment& environment_;
};
struct SessionIOBinding {

2
onnxruntime/core/session/onnxruntime_c_api.cc
View file

@ -1904,8 +1904,8 @@ static constexpr OrtApi ort_api_1_to_5 = {
&OrtApis::GetBoundOutputValues,
&OrtApis::ClearBoundInputs,
&OrtApis::ClearBoundOutputs,
&OrtApis::TensorAt,
&OrtApis::CreateAndRegisterAllocator,
};
// Assert to do a limited check to ensure Version 1 of OrtApi never changes (will detect an addition or deletion but not if they cancel out each other)

1
onnxruntime/core/session/ort_apis.h
View file

@ -230,4 +230,5 @@ ORT_API_STATUS_IMPL(AddSessionConfigEntry, _Inout_ OrtSessionOptions* options,
ORT_API_STATUS_IMPL(TensorAt, _Inout_ OrtValue* value, size_t* location_values, size_t location_values_count, _Outptr_ void** out);
ORT_API_STATUS_IMPL(CreateAndRegisterAllocator, _Inout_ OrtEnv* env, _In_ const OrtMemoryInfo* mem_info, _In_ const OrtArenaCfg* arena_cfg);
} // namespace OrtApis

6
onnxruntime/core/session/ort_env.cc
View file

@ -8,6 +8,7 @@
#include "ort_env.h"
#include "core/session/ort_apis.h"
#include "core/session/environment.h"
#include "core/session/allocator_impl.h"
#include "core/common/logging/logging.h"
#ifdef __ANDROID__
#include "core/platform/android/logging/android_log_sink.h"
@ -100,4 +101,9 @@ onnxruntime::logging::LoggingManager* OrtEnv::GetLoggingManager() const {
void OrtEnv::SetLoggingManager(std::unique_ptr<onnxruntime::logging::LoggingManager> logging_manager) {
value_->SetLoggingManager(std::move(logging_manager));
}
onnxruntime::Status OrtEnv::RegisterAllocator(AllocatorPtr allocator) {
auto status = value_->RegisterAllocator(allocator);
return status;
}

7
onnxruntime/core/session/ort_env.h
View file

@ -8,6 +8,7 @@
#include "core/common/logging/isink.h"
#include "core/platform/ort_mutex.h"
#include "core/common/status.h"
#include "core/framework/allocator.h"
namespace onnxruntime {
class Environment;
@ -55,6 +56,12 @@ struct OrtEnv {
onnxruntime::logging::LoggingManager* GetLoggingManager() const;
void SetLoggingManager(std::unique_ptr<onnxruntime::logging::LoggingManager> logging_manager);
/**
* Registers an allocator for sharing between multiple sessions.
* Returns an error if an allocator with the same OrtMemoryInfo is already registered.
*/
onnxruntime::Status RegisterAllocator(onnxruntime::AllocatorPtr allocator);
private:
static OrtEnv* p_instance_;
static onnxruntime::OrtMutex m_;

13
onnxruntime/python/onnxruntime_pybind_state.cc
View file

@ -9,16 +9,17 @@
#define PY_ARRAY_UNIQUE_SYMBOL onnxruntime_python_ARRAY_API
#include <numpy/arrayobject.h>
#include "core/framework/data_transfer_utils.h"
#include "core/framework/data_types_internal.h"
#include "core/framework/tensorprotoutils.h"
#include "core/graph/graph_viewer.h"
#include "core/common/logging/logging.h"
#include "core/common/logging/severity.h"
#include "core/framework/TensorSeq.h"
#include "core/framework/bfc_arena.h"
#include "core/framework/data_transfer_utils.h"
#include "core/framework/data_types_internal.h"
#include "core/framework/kernel_registry.h"
#include "core/framework/random_seed.h"
#include "core/framework/session_options.h"
#include "core/framework/bfc_arena.h"
#include "core/framework/tensorprotoutils.h"
#include "core/framework/TensorSeq.h"
#include "core/graph/graph_viewer.h"
#include "core/session/IOBinding.h"
#if USE_CUDA

9
onnxruntime/test/framework/allocation_planner_test.cc
View file

@ -5,14 +5,17 @@
#include <unordered_map>
#include <unordered_set>
#include <sstream>
#include "core/framework/session_state.h"
#include "core/graph/model.h"
#include "gtest/gtest.h"
#include "core/framework/session_state.h"
#include "core/framework/kernel_registry.h"
#include "core/framework/op_kernel.h"
#include "test/framework/model_builder_utils.h"
#include "core/framework/allocation_planner.h"
#include "core/util/thread_utils.h"
#include "core/graph/model.h"
#include "core/providers/cpu/cpu_execution_provider.h"
#include "core/util/thread_utils.h"
#include "test/test_environment.h"
#include "test/util/include/asserts.h"

1
onnxruntime/test/framework/allocator_test.cc
View file

@ -3,6 +3,7 @@
#include "core/framework/allocatormgr.h"
#include "core/framework/allocator.h"
#include "test_utils.h"
#include "gtest/gtest.h"

106
onnxruntime/test/framework/inference_session_test.cc
View file

@ -33,6 +33,8 @@
#endif
#include "core/session/environment.h"
#include "core/session/IOBinding.h"
#include "core/session/device_allocator.h"
#include "core/session/allocator_impl.h"
#include "dummy_provider.h"
#include "test_utils.h"
#include "test/capturing_sink.h"
@ -2387,5 +2389,109 @@ TEST(InferenceSessionTests, InvalidSessionEnvCombination) {
}
}
// Tests for sharing allocators between sessions
class InferenceSessionTestSharingAllocator : public InferenceSession {
public:
InferenceSessionTestSharingAllocator(const SessionOptions& session_options,
const Environment& env)
: InferenceSession(session_options, env) {
}
const SessionState& GetSessionState() { return InferenceSession::GetSessionState(); }
};
// Ensure sessions use the same allocator. It uses ORT created allocator.
TEST(InferenceSessionTests, AllocatorSharing_EnsureSessionsUseSameOrtCreatedAllocator) {
auto logging_manager = onnxruntime::make_unique<logging::LoggingManager>(
std::unique_ptr<ISink>(new CLogSink()), logging::Severity::kVERBOSE, false,
LoggingManager::InstanceType::Temporal);
std::unique_ptr<Environment> env;
auto st = Environment::Create(std::move(logging_manager), env);
ASSERT_TRUE(st.IsOK());
// create allocator to register with the env
bool use_arena = true;
#if !(defined(__amd64__) || defined(_M_AMD64))
use_arena = false;
#endif
OrtMemoryInfo mem_info{onnxruntime::CPU, use_arena ? OrtArenaAllocator : OrtDeviceAllocator};
size_t max_mem = std::numeric_limits<size_t>::max();
DeviceAllocatorRegistrationInfo device_info{
OrtMemTypeDefault,
[mem_info](int) { return onnxruntime::make_unique<TAllocator>(mem_info); },
max_mem};
AllocatorPtr allocator_ptr = CreateAllocator(device_info, 0, use_arena);
st = env->RegisterAllocator(allocator_ptr);
ASSERT_STATUS_OK(st);
// create sessions to share the allocator
SessionOptions so1;
AddSessionConfigEntryImpl(so1, ORT_SESSION_OPTIONS_CONFIG_USE_ENV_ALLOCATORS, "1");
InferenceSessionTestSharingAllocator sess1(so1, *env);
ASSERT_STATUS_OK(sess1.Load(MODEL_URI));
ASSERT_STATUS_OK(sess1.Initialize());
SessionOptions so2;
AddSessionConfigEntryImpl(so2, ORT_SESSION_OPTIONS_CONFIG_USE_ENV_ALLOCATORS, "1");
InferenceSessionTestSharingAllocator sess2(so2, *env);
ASSERT_STATUS_OK(sess2.Load(MODEL_URI));
ASSERT_STATUS_OK(sess2.Initialize());
// This line ensures the allocator in the session is the same as that in the env
ASSERT_EQ(sess1.GetSessionState().GetAllocator(mem_info).get(),
allocator_ptr.get());
// This line ensures the underlying IAllocator* is the same across 2 sessions.
ASSERT_EQ(sess1.GetSessionState().GetAllocator(mem_info).get(),
sess2.GetSessionState().GetAllocator(mem_info).get());
}
// Ensure sessions don't use the same allocator. It uses ORT created allocator.
TEST(InferenceSessionTests, AllocatorSharing_EnsureSessionsDontUseSameOrtCreatedAllocator) {
auto logging_manager = onnxruntime::make_unique<logging::LoggingManager>(
std::unique_ptr<ISink>(new CLogSink()), logging::Severity::kVERBOSE, false,
LoggingManager::InstanceType::Temporal);
std::unique_ptr<Environment> env;
auto st = Environment::Create(std::move(logging_manager), env);
ASSERT_TRUE(st.IsOK());
// create allocator to register with the env
bool use_arena = true;
#if !(defined(__amd64__) || defined(_M_AMD64))
use_arena = false;
#endif
OrtMemoryInfo mem_info{onnxruntime::CPU, use_arena ? OrtArenaAllocator : OrtDeviceAllocator};
size_t max_mem = std::numeric_limits<size_t>::max();
DeviceAllocatorRegistrationInfo device_info{
OrtMemTypeDefault,
[mem_info](int) { return onnxruntime::make_unique<TAllocator>(mem_info); },
max_mem};
AllocatorPtr allocator_ptr = CreateAllocator(device_info, 0, use_arena);
st = env->RegisterAllocator(allocator_ptr);
ASSERT_STATUS_OK(st);
// create sessions to share the allocator
SessionOptions so1;
AddSessionConfigEntryImpl(so1, ORT_SESSION_OPTIONS_CONFIG_USE_ENV_ALLOCATORS, "1");
InferenceSessionTestSharingAllocator sess1(so1, *env);
ASSERT_STATUS_OK(sess1.Load(MODEL_URI));
ASSERT_STATUS_OK(sess1.Initialize());
SessionOptions so2;
AddSessionConfigEntryImpl(so2, ORT_SESSION_OPTIONS_CONFIG_USE_ENV_ALLOCATORS, "0");
InferenceSessionTestSharingAllocator sess2(so2, *env);
ASSERT_STATUS_OK(sess2.Load(MODEL_URI));
ASSERT_STATUS_OK(sess2.Initialize());
// This line ensures the allocator in the session is the same as that in the env
ASSERT_EQ(sess1.GetSessionState().GetAllocator(mem_info).get(),
allocator_ptr.get());
// This line ensures the underlying OrtAllocator* is the same across 2 sessions.
ASSERT_NE(sess1.GetSessionState().GetAllocator(mem_info).get(),
sess2.GetSessionState().GetAllocator(mem_info).get());
}
} // namespace test
} // namespace onnxruntime

3
onnxruntime/test/framework/session_state_test.cc
View file

@ -5,6 +5,7 @@
#include "core/framework/execution_providers.h"
#include "core/framework/graph_partitioner.h"
#include "core/framework/kernel_registry.h"
#include "core/framework/op_kernel.h"
#include "core/framework/session_state.h"
#include "core/graph/graph_utils.h"
@ -259,7 +260,7 @@ TEST_P(SessionStatePrepackingTest, PrePackingTest) {
SessionOptions sess_options;
bool use_prepacking = GetParam();
sess_options.session_configurations[ORT_SESSION_OPTIONS_CONFIG_DISABLEPREPACKING] = use_prepacking ? "0" : "1";
ASSERT_STATUS_OK(session_state.FinalizeSessionState(std::basic_string<PATH_CHAR_TYPE>(),
ASSERT_STATUS_OK(session_state.FinalizeSessionState(std::basic_string<PATH_CHAR_TYPE>(),
kernel_registry_manager,
sess_options));

3
onnxruntime/test/optimizer/cse_test.cc
View file

@ -231,10 +231,11 @@ TEST(CseTests, MergeConstants) {
GraphTransformerManager graph_transformation_mgr(1);
// In current implementation, equal constants are not merged. So CSE must precede constant folding, otherwise we end up
// with multiple copies of the same constant.
std::unique_ptr<CPUExecutionProvider> e = onnxruntime::make_unique<CPUExecutionProvider>(CPUExecutionProviderInfo());
ASSERT_TRUE(
graph_transformation_mgr.Register(onnxruntime::make_unique<CommonSubexpressionElimination>(), TransformerLevel::Level1).IsOK());
ASSERT_TRUE(
graph_transformation_mgr.Register(onnxruntime::make_unique<ConstantFolding>(), TransformerLevel::Level1).IsOK());
graph_transformation_mgr.Register(onnxruntime::make_unique<ConstantFolding>(*e.get()), TransformerLevel::Level1).IsOK());
ASSERT_TRUE(
graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, DefaultLoggingManager().DefaultLogger()).IsOK());

23
onnxruntime/test/optimizer/graph_transform_test.cc
View file

@ -139,9 +139,10 @@ TEST_F(GraphTransformationTests, ConstantFolding) {
Graph& graph = model->MainGraph();
std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
ASSERT_TRUE(op_to_count["Unsqueeze"] == 2);
std::unique_ptr<CPUExecutionProvider> e =
onnxruntime::make_unique<CPUExecutionProvider>(CPUExecutionProviderInfo());
onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
graph_transformation_mgr.Register(onnxruntime::make_unique<ConstantFolding>(), TransformerLevel::Level1);
graph_transformation_mgr.Register(onnxruntime::make_unique<ConstantFolding>(*e.get()), TransformerLevel::Level1);
ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));
@ -156,9 +157,10 @@ TEST_F(GraphTransformationTests, ConstantFoldingNodesOnDifferentEP) {
Graph& graph = model->MainGraph();
std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
ASSERT_TRUE(op_to_count["Unsqueeze"] == 2);
std::unique_ptr<CPUExecutionProvider> e =
onnxruntime::make_unique<CPUExecutionProvider>(CPUExecutionProviderInfo());
onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
graph_transformation_mgr.Register(onnxruntime::make_unique<ConstantFolding>(), TransformerLevel::Level1);
graph_transformation_mgr.Register(onnxruntime::make_unique<ConstantFolding>(*e.get()), TransformerLevel::Level1);
// assign all nodes to CUDA. the constant folding should override this to perform the constant folding on cpu
for (auto& node : graph.Nodes()) {
@ -236,9 +238,10 @@ TEST_F(GraphTransformationTests, ConstantFoldingSubgraph) {
std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
ASSERT_TRUE(op_to_count["Add"] == 2); // one in each subgraph
std::unique_ptr<CPUExecutionProvider> e =
onnxruntime::make_unique<CPUExecutionProvider>(CPUExecutionProviderInfo());
onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
graph_transformation_mgr.Register(onnxruntime::make_unique<ConstantFolding>(), TransformerLevel::Level1);
graph_transformation_mgr.Register(onnxruntime::make_unique<ConstantFolding>(*e.get()), TransformerLevel::Level1);
ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));
@ -261,7 +264,9 @@ TEST_F(GraphTransformationTests, ConstantFoldingWithShapeToInitializer) {
std::unordered_set<std::string> excluded_initializers;
excluded_initializers.insert("matmul_weight");
onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
graph_transformation_mgr.Register(onnxruntime::make_unique<ConstantFolding>(compatible_eps, excluded_initializers), TransformerLevel::Level1);
std::unique_ptr<CPUExecutionProvider> e =
onnxruntime::make_unique<CPUExecutionProvider>(CPUExecutionProviderInfo());
graph_transformation_mgr.Register(onnxruntime::make_unique<ConstantFolding>(*e.get(), compatible_eps, excluded_initializers), TransformerLevel::Level1);
ASSERT_TRUE(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_).IsOK());
@ -283,7 +288,9 @@ TEST_F(GraphTransformationTests, ConstantFoldingWithScalarShapeToInitializer) {
std::unordered_set<std::string> compatible_eps;
onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
graph_transformation_mgr.Register(onnxruntime::make_unique<ConstantFolding>(compatible_eps), TransformerLevel::Level1);
std::unique_ptr<CPUExecutionProvider> e =
onnxruntime::make_unique<CPUExecutionProvider>(CPUExecutionProviderInfo());
graph_transformation_mgr.Register(onnxruntime::make_unique<ConstantFolding>(*e.get(), compatible_eps), TransformerLevel::Level1);
ASSERT_TRUE(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_).IsOK());

13
onnxruntime/test/optimizer/graph_transform_utils_test.cc
View file

@ -41,8 +41,9 @@ TEST(GraphTransformerUtilsTests, TestGenerateGraphTransformers) {
std::string l1_transformer = "ConstantFolding";
std::string l2_transformer = "ConvActivationFusion";
std::vector<std::string> custom_list = {l1_rule1, l1_transformer, l2_transformer};
auto transformers = optimizer_utils::GenerateTransformers(TransformerLevel::Level1, {}, custom_list);
std::unique_ptr<CPUExecutionProvider> cpu_execution_provider =
onnxruntime::make_unique<CPUExecutionProvider>(CPUExecutionProviderInfo());
auto transformers = optimizer_utils::GenerateTransformers(TransformerLevel::Level1, {}, *cpu_execution_provider.get(), custom_list);
ASSERT_TRUE(transformers.size() == 2);
auto l1_rule_transformer_name = optimizer_utils::GenerateRuleBasedTransformerName(TransformerLevel::Level1);
@ -54,7 +55,7 @@ TEST(GraphTransformerUtilsTests, TestGenerateGraphTransformers) {
}
ASSERT_TRUE(rule_transformer && rule_transformer->RulesCount() == 1);
transformers = optimizer_utils::GenerateTransformers(TransformerLevel::Level2, {}, custom_list);
transformers = optimizer_utils::GenerateTransformers(TransformerLevel::Level2, {}, *cpu_execution_provider.get(), custom_list);
#ifndef DISABLE_CONTRIB_OPS
ASSERT_TRUE(transformers.size() == 1);
#else
@ -65,15 +66,17 @@ TEST(GraphTransformerUtilsTests, TestGenerateGraphTransformers) {
TEST(GraphTransformerUtilsTests, TestCustomOnlyTransformers) {
// Transformers that are disabled by default. They can only be enabled by custom list.
std::string l2_transformer = "GeluApproximation";
std::unique_ptr<CPUExecutionProvider> cpu_execution_provider =
onnxruntime::make_unique<CPUExecutionProvider>(CPUExecutionProviderInfo());
std::vector<std::string> default_list = {};
auto default_transformers = optimizer_utils::GenerateTransformers(TransformerLevel::Level2, {}, default_list);
auto default_transformers = optimizer_utils::GenerateTransformers(TransformerLevel::Level2, {}, *cpu_execution_provider.get(), default_list);
for (auto& transformer : default_transformers) {
ASSERT_TRUE(transformer->Name() != l2_transformer);
}
std::vector<std::string> custom_list = {l2_transformer};
auto custom_transformers = optimizer_utils::GenerateTransformers(TransformerLevel::Level2, {}, custom_list);
auto custom_transformers = optimizer_utils::GenerateTransformers(TransformerLevel::Level2, {}, *cpu_execution_provider.get(), custom_list);
#ifndef DISABLE_CONTRIB_OPS
ASSERT_TRUE(custom_transformers.size() == 1);
ASSERT_TRUE(custom_transformers[0]->Name() == l2_transformer);

2
onnxruntime/test/optimizer/optimizer_test.cc
View file

@ -66,7 +66,7 @@ TEST(OptimizerTest, Basic) {
std::unique_ptr<CPUExecutionProvider> cpu_execution_provider =
onnxruntime::make_unique<CPUExecutionProvider>(CPUExecutionProviderInfo());
OptimizerExecutionFrame::Info info(nodes, initialized_tensor_set, std::move(cpu_execution_provider));
OptimizerExecutionFrame::Info info(nodes, initialized_tensor_set, *cpu_execution_provider.get());
std::vector<int> fetch_mlvalue_idxs{info.GetMLValueIndex("out")};
OptimizerExecutionFrame frame(info, fetch_mlvalue_idxs);
const logging::Logger& logger = DefaultLoggingManager().DefaultLogger();

55
onnxruntime/test/shared_lib/test_inference.cc
View file

@ -826,3 +826,58 @@ TEST(CApiTest, get_available_providers_cpp) {
ASSERT_TRUE(providers.size() > 0);
ASSERT_TRUE(providers[0] == std::string("CPUExecutionProvider"));
}
// This test uses the CreateAndRegisterAllocator API to register an allocator with the env,
// creates 2 sessions and then runs those 2 sessions one after another
TEST(CApiTest, TestSharedAllocatorUsingCreateAndRegisterAllocator) {
// simple inference test
// prepare inputs
std::vector<Input> inputs(1);
Input& input = inputs.back();
input.name = "X";
input.dims = {3, 2};
input.values = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
// prepare expected inputs and outputs
std::vector<int64_t> expected_dims_y = {3, 2};
std::vector<float> expected_values_y = {1.0f, 4.0f, 9.0f, 16.0f, 25.0f, 36.0f};
OrtEnv* env_ptr = (OrtEnv*)(*ort_env);
OrtMemoryInfo* mem_info = nullptr;
const auto& api = Ort::GetApi();
std::unique_ptr<OrtMemoryInfo, decltype(api.ReleaseMemoryInfo)> rel_info(mem_info, api.ReleaseMemoryInfo);
ASSERT_TRUE(api.CreateCpuMemoryInfo(OrtArenaAllocator, OrtMemTypeDefault, &mem_info) == nullptr);
OrtArenaCfg arena_cfg{-1, -1, -1, -1};
ASSERT_TRUE(api.CreateAndRegisterAllocator(env_ptr, mem_info, &arena_cfg) == nullptr);
// test for duplicates
std::unique_ptr<OrtStatus, decltype(api.ReleaseStatus)> status_releaser(api.CreateAndRegisterAllocator(env_ptr, mem_info,
&arena_cfg),
api.ReleaseStatus);
ASSERT_FALSE(status_releaser.get() == nullptr);
Ort::SessionOptions session_options;
auto default_allocator = onnxruntime::make_unique<MockedOrtAllocator>();
session_options.AddConfigEntry(ORT_SESSION_OPTIONS_CONFIG_USE_ENV_ALLOCATORS, "1");
// create session 1
Ort::Session session1(*ort_env, MODEL_URI, session_options);
RunSession<float>(default_allocator.get(),
session1,
inputs,
"Y",
expected_dims_y,
expected_values_y,
nullptr);
// create session 2
Ort::Session session2(*ort_env, MODEL_URI, session_options);
RunSession<float>(default_allocator.get(),
session2,
inputs,
"Y",
expected_dims_y,
expected_values_y,
nullptr);
}

3
orttraining/orttraining/core/optimizer/graph_transformer_utils.cc
View file

@ -51,6 +51,7 @@ std::vector<std::unique_ptr<GraphTransformer>> GeneratePreTrainingTransformers(
TransformerLevel level,
const std::unordered_set<std::string>& weights_to_train,
const TrainingSession::TrainingConfiguration::GraphTransformerConfiguration& config,
const IExecutionProvider& execution_provider,
const std::vector<std::string>& transformers_and_rules_to_enable) {
std::vector<std::unique_ptr<GraphTransformer>> transformers;
std::unique_ptr<RuleBasedGraphTransformer> rule_transformer = nullptr;
@ -88,7 +89,7 @@ std::vector<std::unique_ptr<GraphTransformer>> GeneratePreTrainingTransformers(
transformers.emplace_back(onnxruntime::make_unique<GeluApproximation>(compatible_eps));
}
transformers.emplace_back(onnxruntime::make_unique<ConstantFolding>(compatible_eps, weights_to_train));
transformers.emplace_back(onnxruntime::make_unique<ConstantFolding>(execution_provider, compatible_eps, weights_to_train));
auto horizontal_parallel_size = training::DistributedRunContext::GroupSize(training::WorkerGroupType::HorizontalParallel);
if (horizontal_parallel_size > 1) {
LOGS_DEFAULT(WARNING) << horizontal_parallel_size << "-way horizontal model parallel is enabled";

1
orttraining/orttraining/core/optimizer/graph_transformer_utils.h
View file

@ -19,6 +19,7 @@ std::vector<std::unique_ptr<GraphTransformer>> GeneratePreTrainingTransformers(
TransformerLevel level,
const std::unordered_set<std::string>& weights_to_train,
const TrainingSession::TrainingConfiguration::GraphTransformerConfiguration& config,
const IExecutionProvider& execution_provider, // required for constant folding
const std::vector<std::string>& rules_and_transformers_to_enable = {});
/** Generates all predefined (both rule-based and non-rule-based) transformers for this level.

16
orttraining/orttraining/core/session/training_session.cc
View file

@ -7,6 +7,7 @@
#include "core/graph/model.h"
#include "core/session/IOBinding.h"
#include "core/providers/cpu/controlflow/utils.h"
#include "core/providers/cpu/cpu_execution_provider.h"
#include "orttraining/core/graph/loss_function_builder.h"
#include "orttraining/core/graph/optimizer_builder.h"
#include "orttraining/core/framework/checkpointing.h"
@ -492,7 +493,14 @@ static Status AddGradientAccumulationNodes(Graph& graph,
Status TrainingSession::ApplyTransformationsToMainGraph(const std::unordered_set<std::string>& weights_to_train,
const TrainingConfiguration::GraphTransformerConfiguration& config) {
GraphTransformerManager graph_transformation_mgr{1};
AddPreTrainingTransformers(graph_transformation_mgr, weights_to_train, config);
// TODO: ideally we can just reuse the CPU EP registered with the session, but in the training session case
// the EPs are registered after ConfigureForTraining and before Initialize is called. Hence we don't have access
// to the registered CPU EP at this stage. Hence creating the EP here again. This is still much better than
// creating an EP instance for every single node in ConstantFolding.
// Create execution frame for executing constant nodes.
std::unique_ptr<CPUExecutionProvider> cpu_execution_provider =
onnxruntime::make_unique<CPUExecutionProvider>(CPUExecutionProviderInfo());
AddPreTrainingTransformers(*cpu_execution_provider, graph_transformation_mgr, weights_to_train, config);
// apply transformers
Graph& graph = model_->MainGraph();
@ -504,15 +512,17 @@ Status TrainingSession::ApplyTransformationsToMainGraph(const std::unordered_set
}
// Registers all the pre transformers with transformer manager
void TrainingSession::AddPreTrainingTransformers(GraphTransformerManager& transformer_manager,
void TrainingSession::AddPreTrainingTransformers(const IExecutionProvider& execution_provider,
GraphTransformerManager& transformer_manager,
const std::unordered_set<std::string>& weights_to_train,
const TrainingConfiguration::GraphTransformerConfiguration& config,
TransformerLevel graph_optimization_level,
const std::vector<std::string>& custom_list) {
auto add_transformers = [&](TransformerLevel level) {
// Generate and register transformers for level
auto transformers_to_register = transformer_utils::GeneratePreTrainingTransformers(
level, weights_to_train, config, custom_list);
level, weights_to_train, config, execution_provider, custom_list);
for (auto& entry : transformers_to_register) {
transformer_manager.Register(std::move(entry), level);
}

3
orttraining/orttraining/core/session/training_session.h
View file

@ -400,7 +400,8 @@ class TrainingSession : public InferenceSession {
const TrainingConfiguration::GraphTransformerConfiguration& config);
/** configure initial transformers for training */
void AddPreTrainingTransformers(GraphTransformerManager& transformer_manager,
void AddPreTrainingTransformers(const IExecutionProvider& execution_provider, // for constant folding
GraphTransformerManager& transformer_manager,
const std::unordered_set<std::string>& weights_to_train,
const TrainingConfiguration::GraphTransformerConfiguration& config,
TransformerLevel graph_optimization_level = TransformerLevel::MaxLevel,

9
orttraining/orttraining/python/experimental/optim/config.py
View file

@ -53,7 +53,7 @@ class _OptimizerConfig(object):
" and additional entries for custom hyper parameter values")
for k, _ in group.items():
if k != 'params':
assert k in defaults, f"'params' has 'k' hyper parameter not present at 'defaults'"
assert k in defaults or k.replace("_coef", "") in defaults, f"'params' has {k} hyper parameter not present at 'defaults'"
self.name = name
self.lr = float(defaults['lr'])
@ -75,6 +75,9 @@ class _OptimizerConfig(object):
if name not in param_group:
param_group.setdefault(name, value)
if "lambda_coef" in param_group:
param_group["lambda"] = param_group.pop("lambda_coef")
self.params.append(param_group)
@ -160,7 +163,7 @@ class AdamConfig(_OptimizerConfig):
defaults = {'lr': lr,
'alpha': alpha,
'beta': beta,
'lambda_coef': lambda_coef,
'lambda': lambda_coef,
'epsilon': epsilon,
'do_bias_correction': do_bias_correction,
'weight_decay_mode': weight_decay_mode}
@ -223,7 +226,7 @@ class LambConfig(_OptimizerConfig):
defaults = {'lr': lr,
'alpha': alpha,
'beta': beta,
'lambda_coef': lambda_coef,
'lambda': lambda_coef,
'ratio_min': ratio_min,
'ratio_max': ratio_max,
'epsilon': epsilon,

19
orttraining/orttraining/python/experimental/orttrainer.py
View file

@ -525,12 +525,27 @@ class ORTTrainer(object):
trainable_params.add(initializer.name)
optimizer_attributes_map[initializer.name] = {}
optimizer_int_attributes_map[initializer.name] = {}
not_in_param_groups = True
for param_group in self.optim_config.params:
if initializer.name not in param_group['params']:
continue # keep looking for a matching param_group
not_in_param_groups = False
for k, v in param_group.items():
if k == 'params':
continue # 'params' is not a hyper parameter, skip it
# 'params' is not a hyper parameter, skip it. 'lr' per weight is not supported
if k == 'params' or k == 'lr':
continue
if isinstance(v, float):
optimizer_attributes_map[initializer.name][k] = v
elif isinstance(v, int):
optimizer_int_attributes_map[initializer.name][k] = v
else:
raise ValueError("Optimizer attributes must be either float or int.")
# set default values for params not found in groups
if not_in_param_groups:
for k, v in self.optim_config.defaults.items():
if k == 'lr':
continue
if isinstance(v, float):
optimizer_attributes_map[initializer.name][k] = v
elif isinstance(v, int):

162
orttraining/orttraining/test/python/orttraining_test_orttrainer_bert_toy_onnx.py
View file

@ -70,6 +70,22 @@ def bert_model_description(dynamic_shape=True):
return model_desc
def optimizer_parameters_mutiple_groups(model):
'''A method to assign different hyper parameters for different model parameter groups'''
no_decay_keys = ["bias", "gamma", "beta", "LayerNorm"]
no_decay_param_group = []
decay_param_group = []
for initializer in model.graph.initializer:
if any(key in initializer.name for key in no_decay_keys):
no_decay_param_group.append(initializer.name)
else:
decay_param_group.append(initializer.name)
params = [{'params': no_decay_param_group, "alpha": 0.9, "beta": 0.999, "lambda_coef": 0.0, "epsilon": 1e-6},
{'params': decay_param_group, "alpha": 0.9, "beta": 0.999, "lambda_coef": 0.01, "epsilon": 1e-6}]
return params
def optimizer_parameters(model):
'''A method to assign different hyper parameters for different model parameter groups'''
@ -78,7 +94,7 @@ def optimizer_parameters(model):
for initializer in model.graph.initializer:
if any(key in initializer.name for key in no_decay_keys):
no_decay_param_group.append(initializer.name)
params = [{'params': no_decay_param_group, "alpha": 0.9, "beta": 0.999, "lambda_coef": 0.0, "epsilon": 1e-6}]
params = [{'params': no_decay_param_group, "alpha": 0.9, "beta": 0.999, "lambda_coef": 0.0, "epsilon": 1e-6, "do_bias_correction":False}]
return params
@ -122,14 +138,6 @@ def legacy_model_params(lr, device = torch.device("cuda", 0)):
learning_rate = torch.tensor([lr]).to(device)
return (legacy_model_desc, learning_rate_description, learning_rate)
no_decay_keys = ["bias", "gamma", "beta", "LayerNorm"]
no_decay = any(no_decay_key in name for no_decay_key in no_decay_keys)
if no_decay:
return {"alpha": 0.9, "beta": 0.999, "lambda": 0.0, "epsilon": 1e-6}
else:
return {"alpha": 0.9, "beta": 0.999, "lambda": 0.01, "epsilon": 1e-6}
def legacy_ort_trainer_learning_rate_description():
return Legacy_IODescription('Learning_Rate', [1, ], torch.float32)
@ -147,6 +155,23 @@ def legacy_bert_model_description():
next_sentence_labels_desc], [loss_desc])
def legacy_optim_params_a(name):
return {"alpha": 0.9, "beta": 0.999, "lambda": 0.01, "epsilon": 1e-6, "do_bias_correction": False}
def legacy_optim_params_b(name):
params = ['bert.embeddings.LayerNorm.bias', 'bert.embeddings.LayerNorm.weight']
if name in params:
return {"alpha": 0.9, "beta": 0.999, "lambda": 0.0, "epsilon": 1e-6, "do_bias_correction": False}
return {"alpha": 0.9, "beta": 0.999, "lambda": 0.01, "epsilon": 1e-6, "do_bias_correction": False}
def legacy_optim_params_c(name):
params_group = optimizer_parameters(load_bert_onnx_model())
if name in params_group[0]['params']:
return {"alpha": 0.9, "beta": 0.999, "lambda": 0.0, "epsilon": 1e-6, "do_bias_correction": False}
return {"alpha": 0.9, "beta": 0.999, "lambda": 0.01, "epsilon": 1e-6, "do_bias_correction": False}
###############################################################################
# Testing starts here #########################################################
###############################################################################
@ -209,34 +234,51 @@ def testToyBERTDeterministicCheck(expected_losses):
@pytest.mark.parametrize("initial_lr, lr_scheduler, expected_learning_rates, expected_losses", [
(1.0, optim.lr_scheduler.ConstantWarmupLRScheduler,\
[0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 1.0, 1.0, 1.0, 1.0],
[10.988012313842773, 10.99213981628418, 18.11327362060547, 35.164119720458984, 40.68449401855469,\
24.585613250732422, 90.66416931152344, 97.23558044433594, 123.73894500732422, 185.28268432617188]),
[0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
[10.988012313842773, 10.99213981628418, 120.79301452636719, 36.11647033691406, 95.83200073242188,\
221.2766571044922, 208.40316772460938, 279.5332946777344, 402.46380615234375, 325.79254150390625]),
(0.5, optim.lr_scheduler.ConstantWarmupLRScheduler,\
[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.5, 0.5, 0.5, 0.5],
[10.988012313842773, 10.99213981628418, 14.942025184631348, 32.938804626464844, 15.188745498657227,\
31.730798721313477, 38.70293045043945, 29.556488037109375, 34.914039611816406, 43.36301040649414]),
[0.0, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
[10.988012313842773, 10.99213981628418, 52.69743347167969, 19.741533279418945, 83.88340759277344,\
126.39848327636719, 91.53898620605469, 63.62016296386719, 102.21206665039062, 180.1424560546875]),
(1.0, optim.lr_scheduler.CosineWarmupLRScheduler,\
[0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 0.9045084971874737, 0.6545084971874737, 0.34549150281252633, 0.09549150281252633],
[10.988012313842773, 10.99213981628418, 18.11327362060547, 35.164119720458984, 40.68449401855469,\
24.585613250732422, 90.66416931152344, 95.54305267333984, 77.43124389648438, 111.71074676513672]),
[0.0, 0.9931806517013612, 0.9397368756032445, 0.8386407858128706, 0.7008477123264848, 0.5412896727361662,\
0.37725725642960045, 0.22652592093878665, 0.10542974530180327, 0.02709137914968268],
[10.988012313842773, 10.99213981628418, 120.6441650390625, 32.152557373046875, 89.63705444335938,\
138.8782196044922, 117.57748413085938, 148.01927185058594, 229.60403442382812, 110.2930908203125]),
(1.0, optim.lr_scheduler.LinearWarmupLRScheduler,\
[0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 0.8, 0.6, 0.4, 0.2],\
[10.988012313842773, 10.99213981628418, 18.11327362060547, 35.164119720458984, 40.68449401855469, 24.585613250732422,\
90.66416931152344, 94.88548278808594, 77.90852355957031, 116.56438446044922]),
[0.0, 0.9473684210526315, 0.8421052631578947, 0.7368421052631579, 0.631578947368421, 0.5263157894736842,\
0.42105263157894735, 0.3157894736842105, 0.21052631578947367, 0.10526315789473684],
[10.988012313842773, 10.99213981628418, 112.89633178710938, 31.114538192749023, 80.94029235839844,\
131.34490966796875, 111.4329605102539, 133.74252319335938, 219.37344360351562, 109.67041015625]),
(1.0, optim.lr_scheduler.PolyWarmupLRScheduler,\
[0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 0.80000002, 0.60000004, 0.40000006000000005, 0.20000007999999997],
[10.988012313842773, 10.99213981628418, 18.11327362060547, 35.164119720458984, 40.68449401855469, 24.585613250732422,\
90.66416931152344, 94.88548278808594, 77.90852355957031, 116.56439971923828])
[0.0, 0.9473684263157895, 0.8421052789473684, 0.7368421315789474, 0.6315789842105263, 0.5263158368421054,
0.42105268947368424, 0.31578954210526317, 0.21052639473684212, 0.10526324736842106],
[10.988012313842773, 10.99213981628418, 112.89633178710938, 31.114538192749023, 80.9402847290039,\
131.3447265625, 111.43253326416016, 133.7415008544922, 219.37147521972656, 109.66986083984375])
])
def testToyBERTModelLRScheduler(initial_lr, lr_scheduler, expected_learning_rates, expected_losses):
# Common setup
device = 'cuda'
total_steps = 10
seed = 1
warmup = 0.05
cycles = 0.5
power = 1.
lr_end = 1e-7
torch.manual_seed(seed)
onnxruntime.set_seed(seed)
# Setup LR Schedulers
if lr_scheduler == optim.lr_scheduler.ConstantWarmupLRScheduler or lr_scheduler == optim.lr_scheduler.LinearWarmupLRScheduler:
lr_scheduler = lr_scheduler(total_steps=total_steps, warmup=warmup)
elif lr_scheduler == optim.lr_scheduler.CosineWarmupLRScheduler:
lr_scheduler = lr_scheduler(total_steps=total_steps, warmup=warmup, cycles=cycles)
elif lr_scheduler == optim.lr_scheduler.PolyWarmupLRScheduler:
lr_scheduler = lr_scheduler(total_steps=total_steps, warmup=warmup, power=power, lr_end=lr_end)
else:
raise RuntimeError("Invalid lr_scheduler")
# Modeling
model_desc = bert_model_description()
model = load_bert_onnx_model()
@ -248,7 +290,7 @@ def testToyBERTModelLRScheduler(initial_lr, lr_scheduler, expected_learning_rate
'device': {
'id': device,
},
'lr_scheduler' : lr_scheduler(total_steps=total_steps, warmup=0.5)
'lr_scheduler' : lr_scheduler
})
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
@ -267,11 +309,11 @@ def testToyBERTModelLRScheduler(initial_lr, lr_scheduler, expected_learning_rate
# Dynamic Loss Scaler implemented implicitly
@pytest.mark.parametrize("loss_scaler, expected_losses", [
(None, [10.98803424835205, 10.99240493774414, 11.090575218200684, 11.042827606201172, 10.988829612731934,
11.105679512023926, 10.981969833374023, 11.08173656463623, 10.997121810913086, 11.10731315612793]),
(amp.DynamicLossScaler(), [10.98803424835205, 10.99240493774414, 11.090575218200684, 11.042827606201172,
10.988829612731934, 11.105679512023926, 10.981969833374023, 11.081737518310547, 10.99714183807373, 11.107304573059082]),
(CustomLossScaler(), [10.98803424835205, 10.99240493774414, 11.090554237365723, 11.042823791503906, 10.98877239227295,
(None, [10.98803424835205, 10.99240493774414, 11.090575218200684, 11.042827606201172, 10.988829612731934,\
11.105679512023926, 10.981968879699707, 11.081787109375, 10.997162818908691, 11.107288360595703]),
(amp.DynamicLossScaler(), [10.98803424835205, 10.99240493774414, 11.090575218200684, 11.042827606201172,\
10.988829612731934, 11.105679512023926, 10.981969833374023, 11.081744194030762, 10.997139930725098, 11.107272148132324]),
(CustomLossScaler(), [10.98803424835205, 10.99240493774414, 11.090554237365723, 11.042823791503906, 10.98877239227295,\
11.105667114257812, 10.981982231140137, 11.081765174865723, 10.997125625610352, 11.107298851013184])
])
def testToyBERTModelMixedPrecisionLossScaler(loss_scaler, expected_losses):
@ -594,7 +636,7 @@ def testToyBERTSaveAsONNX():
os.remove(onnx_file_name)
# Create a new trainer from persisted ONNX model and compare with original ONNX model
trainer_from_onnx = orttrainer.ORTTrainer(reload_onnx_model, model_desc, optim_config)#, options=opts)
trainer_from_onnx = orttrainer.ORTTrainer(reload_onnx_model, model_desc, optim_config, options=opts)
assert trainer_from_onnx._onnx_model is not None
assert (id(trainer_from_onnx._onnx_model) != id(trainer._onnx_model))
for initializer, loaded_initializer in zip(trainer._onnx_model.graph.initializer, trainer_from_onnx._onnx_model.graph.initializer):
@ -849,3 +891,61 @@ def testToyBERTModelGradientAccumulationLegacyExperimental(gradient_accumulation
# Check results
_test_helpers.assert_model_outputs(experimental_losses, legacy_losses, rtol=1e-6)
@pytest.mark.parametrize("params, legacy_optim_map", [
# Change the hyper parameters for all parameters
([], legacy_optim_params_a),
# Change the hyperparameters for a subset of hardcoded parameters
([{'params':['bert.embeddings.LayerNorm.bias', 'bert.embeddings.LayerNorm.weight'], "alpha": 0.9,
"beta": 0.999, "lambda_coef": 0.0, "epsilon": 1e-6, "do_bias_correction":False}], legacy_optim_params_b),
# Change the hyperparameters for a generated set of paramers
(optimizer_parameters(load_bert_onnx_model()), legacy_optim_params_c)
])
def testToyBERTModelLegacyExperimentalCustomOptimParameters(params, legacy_optim_map):
# Common setup
total_steps = 10
device = "cuda"
seed = 1
# EXPERIMENTAL API
model_desc = bert_model_description()
model = load_bert_onnx_model()
optim_config = optim.LambConfig(params, alpha= 0.9, beta= 0.999, lambda_coef= 0.01, epsilon= 1e-6, do_bias_correction=False)
opts = orttrainer.ORTTrainerOptions({
'debug' : {
'deterministic_compute': True
},
'device': {
'id': device,
},
})
torch.manual_seed(seed)
onnxruntime.set_seed(seed)
trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, options=opts)
experimental_losses = []
for i in range(total_steps):
sample_input = generate_random_input_from_model_desc(model_desc, i)
experimental_losses.append(trainer.train_step(*sample_input).cpu().item())
# LEGACY IMPLEMENTATION
device = torch.device(device)
legacy_model_desc, learning_rate_description, learning_rate = legacy_model_params(trainer.optim_config.lr)
torch.manual_seed(seed)
onnxruntime.set_seed(seed)
legacy_trainer = Legacy_ORTTrainer(model, None, legacy_model_desc, "LambOptimizer",
legacy_optim_map,
learning_rate_description,
device,
_use_deterministic_compute=True)
legacy_losses = []
for i in range(total_steps):
sample_input = generate_random_input_from_model_desc(model_desc, i)
legacy_sample_input = [*sample_input, learning_rate]
legacy_losses.append(legacy_trainer.train_step(legacy_sample_input).cpu().item())
# Check results
_test_helpers.assert_model_outputs(experimental_losses, legacy_losses)

5
orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py
View file

@ -1,13 +1,12 @@
from functools import partial
import inspect
import math
from numpy.testing import assert_allclose
import onnx
import os
import pytest
import torch
from numpy.testing import assert_allclose
from onnxruntime import set_seed
from onnxruntime.capi.ort_trainer import IODescription as Legacy_IODescription,\
@ -966,8 +965,6 @@ def testORTTrainerLegacyAndExperimentalGradientAccumulation(seed, device, gradie
_test_helpers.assert_model_outputs(legacy_loss, experimental_loss, rtol=1e-6)
@pytest.mark.parametrize("seed,device,optimizer_config,lr_scheduler, get_lr_this_step", [
(0, 'cuda', optim.AdamConfig, optim.lr_scheduler.ConstantWarmupLRScheduler, _test_commons.legacy_constant_lr_scheduler),
(0, 'cuda', optim.LambConfig, optim.lr_scheduler.ConstantWarmupLRScheduler, _test_commons.legacy_constant_lr_scheduler),

2
tools/ci_build/build.py
View file

@ -1093,8 +1093,6 @@ def run_training_python_frontend_tests(cwd):
run_subprocess([
sys.executable, 'orttraining_test_transformers.py',
'BertModelTest.test_for_pretraining_full_precision_list_and_dict_input'], cwd=cwd)
run_subprocess([sys.executable, 'orttraining_test_orttrainer_frontend.py'], cwd=cwd)
run_subprocess([sys.executable, 'orttraining_test_orttrainer_bert_toy_onnx.py'], cwd=cwd)
def run_training_python_frontend_e2e_tests(cwd):

2
tools/ci_build/github/linux/docker/scripts/install_deps.sh
View file

@ -110,7 +110,7 @@ if [ $DEVICE_TYPE = "Normal" ]; then
elif [ $DEVICE_TYPE = "gpu" ]; then
${PYTHON_EXE} -m pip install sympy==1.1.1
if [[ $BUILD_EXTR_PAR = *--enable_training* ]]; then
${PYTHON_EXE} -m pip install --upgrade --pre torch==1.6.0.dev20200610 torchvision==0.7.0.dev20200610 -f https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html
${PYTHON_EXE} -m pip install --upgrade --pre torch==1.6.0.dev20200610 torchvision==0.7.0.dev20200610 torchtext==0.6.0.dev20200610 -f https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html
${PYTHON_EXE} -m pip install transformers==v2.10.0
# transformers requires sklearn
${PYTHON_EXE} -m pip install sklearn