mmdeploy/csrc/core/operator.h

114 lines
3.8 KiB
C++

// Copyright (c) OpenMMLab. All rights reserved.
#ifndef MMDEPLOY_SRC_EXPERIMENTAL_PIPELINE_OPERATOR_H_
#define MMDEPLOY_SRC_EXPERIMENTAL_PIPELINE_OPERATOR_H_
#include "core/value.h"
namespace mmdeploy::graph {
using std::string;
using std::tuple;
using std::vector;
MMDEPLOY_API Result<void> Gather(const Value::Array& array, const vector<int>& idxs,
Value::Array& output);
MMDEPLOY_API Result<void> Gather(Value::Array&& array, const vector<int>& idxs,
Value::Array& output);
MMDEPLOY_API Result<void> Gather(const Value::Object& object, const vector<std::string>& keys,
Value::Array& output);
MMDEPLOY_API Result<void> Gather(Value::Object&& object, const vector<std::string>& keys,
Value::Array& output);
MMDEPLOY_API Result<void> Scatter(Value::Array array, const vector<int>& idxs,
Value::Array& output);
MMDEPLOY_API Result<void> Scatter(Value::Array array, const vector<std::string>& keys,
Value::Object& output);
inline Result<Value::Array> Gather(const Value::Array& array, const vector<int>& idxs) {
Value::Array output;
OUTCOME_TRY(Gather(array, idxs, output));
return output;
}
inline Result<Value::Array> Gather(Value::Array&& array, const vector<int>& idxs) {
Value::Array output;
OUTCOME_TRY(Gather(std::move(array), idxs, output));
return output;
}
inline Result<Value::Array> Gather(const Value::Object& object, const vector<std::string>& keys) {
Value::Array output;
OUTCOME_TRY(Gather(object, keys, output));
return output;
}
inline Result<Value::Array> Gather(Value::Object&& object, const vector<std::string>& keys) {
Value::Array output;
OUTCOME_TRY(Gather(std::move(object), keys, output));
return output;
}
inline Result<Value::Array> Scatter(Value::Array array, const vector<int>& idxs) {
Value::Array output(idxs.size(), Value::kNull);
OUTCOME_TRY(Scatter(std::move(array), idxs, output));
return output;
}
inline Result<Value::Object> Scatter(Value::Array array, const vector<std::string>& keys) {
Value::Object output;
OUTCOME_TRY(Scatter(std::move(array), keys, output));
return output;
}
template <class V, std::enable_if_t<is_value_v<std::decay_t<V> >, bool> = true>
Result<tuple<Value, vector<int> > > Flatten(V&& input) {
if (!input.is_array()) {
return Status(eInvalidArgument);
}
Value output = ValueType::kArray;
std::vector<int> idxs;
for (int i = 0; i < input.size(); ++i) {
auto inner = std::forward<V>(input)[i];
if (!inner.is_array()) {
return Status(eInvalidArgument);
}
for (auto& item : inner) {
output.push_back(std::move(item));
idxs.push_back(i);
}
}
idxs.push_back(input.size());
return {output, idxs};
}
template <class V, std::enable_if_t<is_value_v<std::decay_t<V> >, bool> = true>
Result<Value> Unflatten(V&& input, const vector<int>& idxs) {
if (!input.is_array()) {
return Status(eInvalidArgument);
}
Value output = ValueType::kArray;
for (int i = 0; i < idxs.back(); ++i) {
output.push_back(ValueType::kArray);
}
for (int i = 0; i < input.size(); ++i) {
if (idxs[i] >= output.size()) {
return Status(eInvalidArgument);
}
output[idxs[i]].push_back(std::forward<V>(input)[i]);
}
return output;
}
// object of arrays -> array of objects, all arrays must be of same length
MMDEPLOY_API Result<Value> DistribOA(const Value& oa);
// array of objects -> object of arrays, all objects must be isomorphic
MMDEPLOY_API Result<Value> DistribAO(const Value& ao);
// array of arrays -> array of arrays, this is equivalent to transpose
MMDEPLOY_API Result<Value> DistribAA(const Value& a);
} // namespace mmdeploy::graph
#endif // MMDEPLOY_SRC_EXPERIMENTAL_PIPELINE_OPERATOR_H_