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

Allow Kernels refer to some attribute data directly in the protobuf (#5624)

Browse source 
* Introduce OpKernelInfo GetAttrAsSpan() for floats and ints attribute proto arrays
  and GetAttrsStringRefs() to return a vector of string references.
  These new APIs allow kernels not copy attribute arrays especially if they are large
  and save on memory.
  but refer directly to data that is in AttributeProto.
  Modify TfIdfVectorizer to take advantage of the new API.

Signed-off-by: Dmitri Smirnov <dmitrism@microsoft.com>
This commit is contained in:
Dmitri Smirnov 2020-10-29 16:12:54 -07:00 committed by GitHub
parent 1fa1c51544
commit 742ffb860c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 125 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
include/onnxruntime/core/framework/op_node_proto_helper.h
View file

@ -81,6 +81,23 @@ class OpNodeProtoHelper {
template <typename T>
MUST_USE_RESULT Status GetAttrs(const std::string& name, gsl::span<T> values) const;
/// <summary>
/// Return a gsl::span that points to an array of primitive types held by AttributeProto
/// This function allows to avoid copying big attributes locally into a kernel and operate on
/// AttributeProto data directly.
///
/// Does not apply to strings, Tensors and Sparse Tensors that require special treatment.
/// </summary>
/// <typeparam name="T">Primitive type contained in the array</typeparam>
/// <param name="name">Attribute name</param>
/// <param name="values">Attribute data in a span, out parameter</param>
/// <returns>Status</returns>
template <typename T>
MUST_USE_RESULT Status GetAttrsAsSpan(const std::string& name, gsl::span<const T>& values) const;
MUST_USE_RESULT Status GetAttrsStringRefs(const std::string& name,
std::vector<std::reference_wrapper<const std::string>>& refs) const;
uint32_t GetPrimitiveAttrElementCount(ONNX_NAMESPACE::AttributeProto_AttributeType type,
const std::string& name) const noexcept;

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

@ -37,6 +37,43 @@ inline bool HasTyped<GraphProto>(const AttributeProto* attr) {
return utils::HasGraph(*attr);
}
template <typename T>
inline bool HasTypedList(const AttributeProto* attr);
template <>
inline bool HasTypedList<float>(const AttributeProto* attr) {
return utils::HasFloats(*attr);
}
template <>
inline bool HasTypedList<int64_t>(const AttributeProto* attr) {
return utils::HasInts(*attr);
}
template <>
inline bool HasTypedList<std::string>(const AttributeProto* attr) {
return utils::HasStrings(*attr);
}
template <typename T>
inline constexpr int ArrayTypeToAttributeType();
template <>
inline constexpr int ArrayTypeToAttributeType<float>() {
return AttributeProto_AttributeType_FLOATS;
}
template <>
inline constexpr int ArrayTypeToAttributeType<int64_t>() {
return AttributeProto_AttributeType_INTS;
}
template <>
inline constexpr int ArrayTypeToAttributeType<std::string>() {
return AttributeProto_AttributeType_STRINGS;
}
#define ORT_DEFINE_GET_ATTR(IMPL_T, T, type) \
template <> \
template <> \
@ -84,6 +121,28 @@ inline bool HasTyped<GraphProto>(const AttributeProto* attr) {
return Status::OK(); \
}
// Will not work for std::strings
#define ORT_DEFINE_GET_ATTRS_AS_SPAN(IMPL_T, T, list) \
template <> \
template <> \
Status OpNodeProtoHelper<IMPL_T>::GetAttrsAsSpan<T>( \
const std::string& name, gsl::span<const T>& values) const { \
const AttributeProto* attr = TryGetAttribute(name); \
if (!attr) { \
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "No attribute with name: ", name, " is defined."); \
} \
if (!HasTypedList<T>(attr)) { \
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Attribute: ", name, \
" expected to be of type: ", \
AttributeProto::AttributeType_Name(ArrayTypeToAttributeType<T>()), \
" but is of type: ", \
AttributeProto::AttributeType_Name(attr->type())); \
} \
values = gsl::make_span<const T>(reinterpret_cast<const T*>(attr->list().data()), \
attr->list##_size()); \
return Status::OK(); \
}
#if !defined(ORT_MINIMAL_BUILD)
#define ORT_DEFINE_GET_ATTR_SPECIALIZATIONS(type, list) \
ORT_DEFINE_GET_ATTR(ProtoHelperNodeContext, type, list) \
@ -92,12 +151,19 @@ 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)
#define ORT_DEFINE_GET_ATTRS_SPAN_SPECIALIZATION(type, list) \
ORT_DEFINE_GET_ATTRS_AS_SPAN(ProtoHelperNodeContext, 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)
#define ORT_DEFINE_GET_ATTRS_SPAN_SPECIALIZATION(type, list) \
ORT_DEFINE_GET_ATTRS_AS_SPAN(ProtoHelperNodeContext, type, list)
#endif
ORT_DEFINE_GET_ATTR_SPECIALIZATIONS(float, f)
@ -111,6 +177,33 @@ ORT_DEFINE_GET_ATTRS_SPECIALIZATIONS(std::string, strings)
ORT_DEFINE_GET_ATTRS_SPECIALIZATIONS(TensorProto, tensors)
ORT_DEFINE_GET_ATTRS_SPECIALIZATIONS(GraphProto, graphs)
ORT_DEFINE_GET_ATTRS_SPAN_SPECIALIZATION(float, floats)
ORT_DEFINE_GET_ATTRS_SPAN_SPECIALIZATION(int64_t, ints)
template <typename Impl_t>
MUST_USE_RESULT Status OpNodeProtoHelper<Impl_t>::GetAttrsStringRefs(
const std::string& name,
std::vector<std::reference_wrapper<const std::string>>& refs) const {
const AttributeProto* attr = TryGetAttribute(name);
if (!attr) {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "No attribute with name: ", name, " is defined.");
}
if (!HasTypedList<std::string>(attr)) {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Requested attribute: ",
name, " is expected to have type: ",
AttributeProto::AttributeType_Name(AttributeProto_AttributeType_STRINGS),
" but is of type: ",
AttributeProto::AttributeType_Name(attr->type()));
}
std::vector<std::reference_wrapper<const std::string>> result;
if (attr->strings_size() > 0) {
result.reserve(attr->strings_size());
std::copy(attr->strings().cbegin(), attr->strings().cend(), std::back_inserter(result));
}
refs.swap(result);
return Status::OK();
}
size_t ProtoHelperNodeContext::getNumInputs() const {
return node_.InputDefs().size();
}

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

@ -121,13 +121,12 @@ struct TfIdfVectorizer::Impl {
// and so on in pool. For example, if ngram_counts is [0, 17, 36],
// the first index (zero-based) of 1-gram/2-gram/3-gram
// in pool are 0/17/36.
std::vector<int64_t> ngram_counts_;
gsl::span<const int64_t> ngram_counts_;
// Contains output indexes
// represents ngram_indexes output
std::vector<int64_t> ngram_indexes_;
std::vector<float> weights_;
gsl::span<const int64_t> ngram_indexes_;
gsl::span<const float> weights_;
std::vector<std::string> pool_strings_;
// This map contains references to pool_string_ entries
// of pool_strings attribute
StrMap str_map_;
@ -179,7 +178,7 @@ TfIdfVectorizer::TfIdfVectorizer(const OpKernelInfo& info) : OpKernel(info), imp
ORT_ENFORCE(status.IsOK(), "max_skip_count is required");
ORT_ENFORCE(impl_->max_skip_count_ >= 0, "max_skip_count must be non-negative: ", std::to_string(impl_->max_skip_count_));
status = info.GetAttrs(std::string("ngram_counts"), impl_->ngram_counts_);
status = info.GetAttrsAsSpan("ngram_counts", impl_->ngram_counts_);
ORT_ENFORCE(status.IsOK() && !impl_->ngram_counts_.empty(), "Non-empty ngram_counts is required");
ORT_ENFORCE(size_t(impl_->min_gram_length_) <= impl_->ngram_counts_.size(),
"min_gram_length must be inbounds of ngram_counts: ",
@ -188,7 +187,7 @@ TfIdfVectorizer::TfIdfVectorizer(const OpKernelInfo& info) : OpKernel(info), imp
"max_gram_length must be inbounds of ngram_counts: ",
std::to_string(impl_->max_gram_length_), " <= ", std::to_string(impl_->ngram_counts_.size()));
status = info.GetAttrs("ngram_indexes", impl_->ngram_indexes_);
status = info.GetAttrsAsSpan("ngram_indexes", impl_->ngram_indexes_);
ORT_ENFORCE(status.IsOK() && !impl_->ngram_indexes_.empty(), "Non-empty ngram_indexes is required");
{
// Check that all are positive
@ -200,7 +199,7 @@ TfIdfVectorizer::TfIdfVectorizer(const OpKernelInfo& info) : OpKernel(info), imp
impl_->output_size_ = *greatest_hit + 1;
}
status = info.GetAttrs("weights", impl_->weights_);
status = info.GetAttrsAsSpan("weights", impl_->weights_);
if (status.IsOK()) {
ORT_ENFORCE(impl_->weights_.size() == impl_->ngram_indexes_.size(),
"Got weights of size: ", std::to_string(impl_->weights_.size()),
@ -208,17 +207,18 @@ TfIdfVectorizer::TfIdfVectorizer(const OpKernelInfo& info) : OpKernel(info), imp
" must be of equal size");
}
std::vector<int64_t> pool_int64s;
status = info.GetAttrs("pool_strings", impl_->pool_strings_);
gsl::span<const int64_t> pool_int64s;
std::vector<std::reference_wrapper<const std::string>> pool_strings;
status = info.GetAttrsStringRefs("pool_strings", pool_strings);
if (status.IsOK()) {
ORT_ENFORCE(!impl_->pool_strings_.empty(), "pool_strings must not be empty if specified");
ORT_ENFORCE(!pool_strings.empty(), "pool_strings must not be empty if specified");
} else {
status = info.GetAttrs("pool_int64s", pool_int64s);
status = info.GetAttrsAsSpan("pool_int64s", pool_int64s);
ORT_ENFORCE(status.IsOK() && !pool_int64s.empty(), "non-empty pool_int64s is required if pool_strings not provided");
}
// Iterator via the pool. Insert 1 item for 1-grams, 2 items for 2-grams, etc.
const auto total_items = (impl_->pool_strings_.empty()) ? pool_int64s.size() : impl_->pool_strings_.size();
const auto total_items = (pool_strings.empty()) ? pool_int64s.size() : pool_strings.size();
size_t ngram_id = 1; // start with 1, 0 - means no n-gram
// Load into dictionary only required gram sizes
const size_t min_gram_length = impl_->min_gram_length_;
@ -236,10 +236,10 @@ TfIdfVectorizer::TfIdfVectorizer(const OpKernelInfo& info) : OpKernel(info), imp
auto ngrams = items / ngram_size;
// Skip loading into hash_set ngrams that are not in the range of [min_gram_length-max_gram_length]
if (ngram_size >= min_gram_length && ngram_size <= max_gram_length) {
if (impl_->pool_strings_.empty()) {
if (pool_strings.empty()) {
ngram_id = PopulateGrams<int64_t>(pool_int64s.begin() + start_idx, ngrams, ngram_size, ngram_id, impl_->int64_map_);
} else {
ngram_id = PopulateGrams<std::string>(impl_->pool_strings_.begin() + start_idx, ngrams, ngram_size, ngram_id, impl_->str_map_);
ngram_id = PopulateGrams<std::string>(pool_strings.begin() + start_idx, ngrams, ngram_size, ngram_id, impl_->str_map_);
}
} else {
ngram_id += ngrams;
@ -315,7 +315,7 @@ void TfIdfVectorizer::ComputeImpl(OpKernelContext* ctx, ptrdiff_t row_num, size_
auto X = ctx->Input<Tensor>(0);
const auto elem_size = X->DataType()->Size();
const void* row_begin = AdvanceElementPtr(X->DataRaw(), row_num * row_size, elem_size);
const void* const row_begin = AdvanceElementPtr(X->DataRaw(), row_num * row_size, elem_size);
const void* const row_end = AdvanceElementPtr(row_begin, row_size, elem_size);
const auto& impl = *impl_;