10 #include <unordered_map>
18 namespace onnxruntime {
21 typedef std::map<size_t, OrtMemType>
MemTypeMap;
26 constexpr
static inline bool MemTypeOnCpuExplicitly(
OrtMemType mem_type) {
42 *start = op_since_version_start_;
43 *end = op_since_version_end_;
46 #ifdef onnxruntime_PYBIND_EXPORT_OPSCHEMA
48 return std::pair<int, int>(op_since_version_start_, op_since_version_end_);
53 return provider_type_;
57 const std::unordered_map<std::string, std::vector<MLDataType>>&
TypeConstraints()
const {
58 return type_constraints_;
61 const std::vector<std::pair<int, int>>&
MayInplace()
const {
65 const std::vector<std::pair<int, int>>&
Alias()
const {
70 return variadic_alias_offsets_;
74 auto it = input_memory_type_args_.find(input_index);
75 if (it == input_memory_type_args_.end())
76 return default_inputs_mem_type_;
88 #ifdef ENABLE_STRIDED_TENSORS
89 const std::vector<int>& MayStridedInput()
const {
return may_strided_inputs_; }
90 const std::vector<std::pair<int, int>>& MayStridedOutput()
const {
return may_strided_output_map_; }
94 auto it = output_memory_type_args_.find(output_index);
95 if (it == output_memory_type_args_.end())
96 return default_outputs_mem_type_;
101 return exec_queue_id_;
115 int op_since_version_start_ = 1;
116 int op_since_version_end_ = INT_MAX;
128 std::unordered_map<std::string, std::vector<MLDataType>> type_constraints_;
131 std::vector<std::pair<int, int>> inplace_map_;
134 std::vector<std::pair<int, int>> alias_map_;
138 std::optional<std::pair<int, int>> variadic_alias_offsets_;
141 bool allocate_inputs_contiguously_ =
false;
144 bool external_outputs_ =
false;
146 #ifdef ENABLE_STRIDED_TENSORS
148 std::vector<int> may_strided_inputs_;
151 std::vector<std::pair<int, int>> may_strided_output_map_;
159 int exec_queue_id_ = 0;
168 static std::unique_ptr<KernelDefBuilder>
Create() {
return std::make_unique<KernelDefBuilder>(); }
171 : kernel_def_(std::make_unique<
KernelDef>()) {}
183 kernel_def_->op_since_version_start_ = since_version;
195 kernel_def_->op_since_version_start_ = since_version_start;
196 kernel_def_->op_since_version_end_ = since_version_end;
251 kernel_def_->allocate_inputs_contiguously_ =
true;
260 kernel_def_->external_outputs_ =
true;
264 #ifdef ENABLE_STRIDED_TENSORS
282 kernel_def_->input_memory_type_args_.insert(std::make_pair(input_index, type));
291 for (
auto input_index : input_indexes) {
292 kernel_def_->input_memory_type_args_.insert(std::make_pair(input_index, type));
302 kernel_def_->output_memory_type_args_.insert(std::make_pair(output_index, type));
311 for (
auto output_index : output_indexes) {
312 kernel_def_->output_memory_type_args_.insert(std::make_pair(output_index, type));
321 kernel_def_->exec_queue_id_ = queue_id;
329 kernel_def_->default_inputs_mem_type_ = mem_type;
337 kernel_def_->default_outputs_mem_type_ = mem_type;
344 std::unique_ptr<KernelDef>
Build() {
345 return std::move(kernel_def_);
350 std::unique_ptr<KernelDef> kernel_def_;
KernelDefBuilder & Provider(ProviderType provider_type)
static std::unique_ptr< KernelDefBuilder > Create()
KernelDefBuilder & Alias(const std::vector< std::pair< int, int >> &aliases)
const std::string & ProviderType
Base class for MLDataType.
KernelDefBuilder & ExternalOutputs()
const std::vector< std::pair< int, int > > & MayInplace() const
GLsizei const GLchar *const * string
OrtMemType InputMemoryType(size_t input_index) const
const std::string & Domain() const
const std::vector< std::pair< int, int > > & Alias() const
bool HasExternalOutputs() const
KernelDefBuilder & SetDomain(const std::string &domain)
onnxruntime::ProviderType Provider() const
KernelDefBuilder & ExecQueueId(int queue_id)
KernelDefBuilder & VariadicAlias(int input_offset, int output_offset)
CPU accessible memory outputted by non-CPU execution provider, i.e. CUDA_PINNED.
bool AllocateInputsContiguously() const
bool IsConflict(const KernelDef &other) const
void SinceVersion(int *start, int *end) const
KernelDefBuilder & OutputMemoryType(OrtMemType type, const std::vector< int > &output_indexes)
KernelDefBuilder & SetDefaultOutputMemoryType(OrtMemType mem_type)
The default allocator for execution provider.
KernelDefBuilder & AllocateInputsContiguously()
KernelDefBuilder & SetName(const std::string &op_name)
const std::unordered_map< std::string, std::vector< MLDataType > > & TypeConstraints() const
KernelDefBuilder & SinceVersion(int since_version)
bool IsOutputOnCpu(size_t output_index) const
std::map< size_t, OrtMemType > MemTypeMap
KernelDefBuilder & TypeConstraint(const std::string &arg_name, std::vector< MLDataType > types)
Any CPU memory used by non-CPU execution provider.
const std::string & OpName() const
OrtMemType OutputMemoryType(size_t output_index) const
KernelDefBuilder & MayInplace(const std::vector< std::pair< int, int >> &inplaces)
KernelDefBuilder & InputMemoryType(OrtMemType type, int input_index)
KernelDefBuilder & InputMemoryType(OrtMemType type, const std::vector< int > &input_indexes)
std::unique_ptr< KernelDef > Build()
KernelDefBuilder & SetDefaultInputsMemoryType(OrtMemType mem_type)
GLsizei GLenum GLenum * types
const std::optional< std::pair< int, int > > & VariadicAlias() const
KernelDefBuilder & OutputMemoryType(OrtMemType type, int output_index)
KernelDefBuilder & SinceVersion(int since_version_start, int since_version_end)
OrtMemType
Memory types for allocated memory, execution provider specific types should be extended in each provi...
bool IsInputOnCpu(size_t input_index) const