[VitisAI] add OpSchema, VitisAI use IKernelLookup to check supported ops, VitisAI def_builder adds TypeConstraint related processing (#21688)

### Description
<!-- Describe your changes. -->
1. add OpSchema
2. VitisAI use IKernelLookup to check supported ops
3. VitisAI def_builder adds TypeConstraint related processing

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

---------

Co-authored-by: Zhenze Wang <zhenzew@xilinx.com>
This commit is contained in:
zz002 2024-08-27 12:16:44 +08:00 committed by GitHub
parent af18824f43
commit 422e6e6fb0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 75 additions and 28 deletions

View file

@ -123,6 +123,7 @@ struct ValueInfoProto;
struct ValueInfoProtos; // RepeatedPtrField
struct FunctionProto;
struct InferenceContext;
struct OpSchema;
class GraphInferencer;
using InferenceFunction = std::function<void(InferenceContext&)>;
} // namespace ONNX_NAMESPACE

View file

@ -174,6 +174,7 @@ MLDataType DataTypeImpl::GetType<UInt4x2>() { return Provider_GetHost()->DataTyp
template <>
MLDataType DataTypeImpl::GetType<std::string>() { return Provider_GetHost()->DataTypeImpl__GetType_string(); }
MLDataType DataTypeImpl::GetTensorTypeFromOnnxType(int onnx_type) { return Provider_GetHost()->DataTypeImpl__GetTensorTypeFromOnnxType(onnx_type); }
template <>
MLDataType DataTypeImpl::GetTensorType<bool>() { return Provider_GetHost()->DataTypeImpl__GetTensorType_bool(); }
template <>

View file

@ -27,6 +27,8 @@ enum OperatorStatus : int;
// String pointer as unique TypeProto identifier.
using DataType = const std::string*;
using DataTypeSet = std::unordered_set<DataType>;
using TypeConstraintMap = std::unordered_map<std::string, std::pair<DataTypeSet, std::string>>;
} // namespace ONNX_NAMESPACE
@ -566,6 +568,12 @@ struct ProviderHost {
virtual ONNX_NAMESPACE::StringStringEntryProto* FunctionProto__add_metadata_props(ONNX_NAMESPACE::FunctionProto* p) = 0;
virtual void RegisterSchema(const std::string& domain, const OrtCustomOp* op, int type) = 0;
virtual const ONNX_NAMESPACE::OpSchema* GetSchema(const std::string& name, const int maxInclusiveVersion, const std::string& domain) = 0;
virtual const std::string& OpSchema__inputs__GetName(const ONNX_NAMESPACE::OpSchema* p, const size_t i) = 0;
virtual const std::string& OpSchema__inputs__GetTypeStr(const ONNX_NAMESPACE::OpSchema* p, const size_t i) = 0;
virtual const std::string& OpSchema__outputs__GetName(const ONNX_NAMESPACE::OpSchema* p, const size_t i) = 0;
virtual const std::string& OpSchema__outputs__GetTypeStr(const ONNX_NAMESPACE::OpSchema* p, const size_t i) = 0;
virtual const ONNX_NAMESPACE::TypeConstraintMap& OpSchema__typeConstraintMap(const ONNX_NAMESPACE::OpSchema* p) const = 0;
// ConfigOptions
virtual std::optional<std::string> ConfigOptions__GetConfigEntry(const ConfigOptions* p, const std::string& config_key) = 0;
@ -694,6 +702,7 @@ struct ProviderHost {
virtual MLDataType DataTypeImpl__GetType_Int4x2() = 0;
virtual MLDataType DataTypeImpl__GetType_UInt4x2() = 0;
virtual MLDataType DataTypeImpl__GetTensorTypeFromOnnxType(int) = 0;
virtual MLDataType DataTypeImpl__GetTensorType_bool() = 0;
virtual MLDataType DataTypeImpl__GetTensorType_int8() = 0;
virtual MLDataType DataTypeImpl__GetTensorType_uint8() = 0;

View file

@ -451,6 +451,15 @@ struct FunctionProto final {
FunctionProto(const FunctionProto&) = delete;
void operator=(const FunctionProto&) = delete;
};
struct OpSchema final {
const TypeConstraintMap& typeConstraintMap() const { return g_host->OpSchema__typeConstraintMap(this); }
const std::string& inputs__GetName(const size_t i) const { return g_host->OpSchema__inputs__GetName(this, i); };
const std::string& inputs__GetTypeStr(const size_t i) const { return g_host->OpSchema__inputs__GetTypeStr(this, i); };
const std::string& outputs__GetName(const size_t i) const { return g_host->OpSchema__outputs__GetName(this, i); };
const std::string& outputs__GetTypeStr(const size_t i) const { return g_host->OpSchema__outputs__GetTypeStr(this, i); };
PROVIDER_DISALLOW_ALL(OpSchema)
};
} // namespace ONNX_NAMESPACE
namespace onnxruntime {
@ -702,6 +711,7 @@ class DataTypeImpl final {
#endif
static MLDataType GetTypeFromOnnxType(int);
static MLDataType GetTensorTypeFromOnnxType(int);
bool IsTensorType() const { return g_host->DataTypeImpl__IsTensorType(this); }
bool IsTensorSequenceType() const { return g_host->DataTypeImpl__IsTensorSequenceType(this); }

View file

@ -42,7 +42,7 @@ std::unique_ptr<ComputeCapability> XirSubgraphToComputeCapability1(const onnxrun
std::vector<std::unique_ptr<ComputeCapability>>
GetComputeCapabilityOps(const onnxruntime::GraphViewer& graph,
vaip_core::DllSafe<std::vector<std::unique_ptr<vaip_core::ExecutionProvider>>>* eps,
const std::set<std::string>& all_support_optypes_by_eps) {
const onnxruntime::IExecutionProvider::IKernelLookup& kernel_lookup) {
std::set<NodeIndex> all_nodes_included_eps;
for (auto& ep : **eps) {
auto nodes = node_names_to_nodes(graph, *ep->get_meta_def_nodes());
@ -52,11 +52,8 @@ GetComputeCapabilityOps(const onnxruntime::GraphViewer& graph,
std::vector<NodeIndex> node_indexs = graph.GetNodesInTopologicalOrder();
node_indexs.erase(std::remove_if(node_indexs.begin(), node_indexs.end(), [&](NodeIndex index) { return all_nodes_included_eps.count(index) > 0; }), node_indexs.end());
node_indexs.erase(std::remove_if(node_indexs.begin(), node_indexs.end(),
[&](NodeIndex index) {
auto node = graph.GetNode(index);
return all_support_optypes_by_eps.count(node->Domain() + ":" + node->OpType()) == 0; }),
[&](NodeIndex index) { return kernel_lookup.LookUpKernel(*graph.GetNode(index)) == nullptr; }),
node_indexs.end());
std::vector<std::unique_ptr<ComputeCapability>> result;
for (auto& n : node_indexs) {
auto indexed_subgraph = IndexedSubGraph::Create();

View file

@ -133,17 +133,46 @@ void create_kernel_registry(std::vector<OrtCustomOpDomain*> domains) {
s_kernel_registry_vitisaiep = KernelRegistry::Create();
for (const auto& domain : domains) {
for (const auto* op : domain->custom_ops_) {
const size_t input_count = op->GetInputTypeCount(op);
const size_t output_count = op->GetOutputTypeCount(op);
auto def_builder = KernelDefBuilder::Create();
def_builder->SetName(op->GetName(op));
def_builder->SetDomain(domain->domain_.c_str());
def_builder->SinceVersion(op->GetStartVersion(op), op->GetEndVersion(op));
if (op->version > 12) {
auto input_count = op->GetInputTypeCount(op);
for (auto i = 0u; i < input_count; i++) {
def_builder->InputMemoryType(op->GetInputMemoryType(op, i), i);
}
}
def_builder->Provider(onnxruntime::kVitisAIExecutionProvider);
def_builder->Provider(op->GetExecutionProviderType(op));
auto schema = Provider_GetHost()->GetSchema(op->GetName(op), op->GetStartVersion(op), domain->domain_);
for (size_t i = 0; i < input_count; i++) {
const auto input_type = op->GetInputType(op, i);
auto input_name = schema->inputs__GetName(i);
if (schema->typeConstraintMap().count(schema->inputs__GetTypeStr(i))) {
input_name = schema->inputs__GetTypeStr(i);
}
if (input_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED) {
def_builder->TypeConstraint(input_name.c_str(), DataTypeImpl::AllTensorTypes());
} else {
def_builder->TypeConstraint(input_name.c_str(), DataTypeImpl::GetTensorTypeFromOnnxType(input_type));
}
}
for (size_t i = 0; i < output_count; i++) {
const auto output_type = op->GetOutputType(op, i);
auto output_name = schema->outputs__GetName(i);
if (schema != nullptr) {
if (schema->typeConstraintMap().count(schema->outputs__GetTypeStr(i))) {
output_name = schema->outputs__GetTypeStr(i);
}
}
if (output_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED) {
def_builder->TypeConstraint(output_name.c_str(), DataTypeImpl::AllTensorTypes());
} else {
def_builder->TypeConstraint(output_name.c_str(), DataTypeImpl::GetTensorTypeFromOnnxType(output_type));
}
}
KernelCreateFn kernel_create_fn =
[op](FuncManager&, const OpKernelInfo& info, std::unique_ptr<OpKernel>& out) -> Status {
out = std::make_unique<MyCustomOpKernel>(info, *op);

View file

@ -12,13 +12,15 @@ namespace vaip {
void register_xir_ops(const std::vector<OrtCustomOpDomain*>& domains) {
for (auto domain : domains) {
for (auto op : domain->custom_ops_) {
auto name = op->GetName(op);
if ((std::string)name == "super_layer") {
Provider_GetHost()->RegisterSchema(domain->domain_, op, 1);
} else if ((std::string)name == "FixNeuron") {
Provider_GetHost()->RegisterSchema(domain->domain_, op, 2);
} else {
Provider_GetHost()->RegisterSchema(domain->domain_, op, 3);
if (Provider_GetHost()->GetSchema(op->GetName(op), op->GetStartVersion(op), domain->domain_) == nullptr) {
auto name = op->GetName(op);
if ((std::string)name == "super_layer") {
Provider_GetHost()->RegisterSchema(domain->domain_, op, 1);
} else if ((std::string)name == "FixNeuron") {
Provider_GetHost()->RegisterSchema(domain->domain_, op, 2);
} else {
Provider_GetHost()->RegisterSchema(domain->domain_, op, 3);
}
}
}
}

View file

@ -8,6 +8,6 @@ namespace vaip {
using namespace ::onnxruntime;
std::unique_ptr<ComputeCapability> XirSubgraphToComputeCapability1(const onnxruntime::GraphViewer& graph, vaip_core::ExecutionProvider* ep, size_t index);
std::vector<std::unique_ptr<ComputeCapability>> GetComputeCapabilityOps(const onnxruntime::GraphViewer& graph,
vaip_core::DllSafe<std::vector<std::unique_ptr<vaip_core::ExecutionProvider>>>* ep, const std::set<std::string>& all_not_support_optypes);
vaip_core::DllSafe<std::vector<std::unique_ptr<vaip_core::ExecutionProvider>>>* ep, const onnxruntime::IExecutionProvider::IKernelLookup& kernel_lookup);
} // namespace vaip

View file

@ -25,8 +25,6 @@ constexpr const char* VITISAI = "VITISAI";
VitisAIExecutionProvider::VitisAIExecutionProvider(
const ProviderOptions& info)
: IExecutionProvider{onnxruntime::kVitisAIExecutionProvider}, info_(info) {
CreateKernelRegistry();
auto it = info_.find("ep_context_enable");
ep_ctx_enabled_ = it != info_.end() && it->second == "1";
it = info_.find("ep_context_embed_mode");
@ -39,14 +37,6 @@ VitisAIExecutionProvider::VitisAIExecutionProvider(
LOGS_DEFAULT(VERBOSE) << "User specified EP context cache path: " << ep_ctx_model_path_cfg_;
}
void VitisAIExecutionProvider::CreateKernelRegistry() {
for (const auto& domain : get_domains_vitisaiep()) {
for (const auto* op : domain->custom_ops_) {
vitisai_optypes_.insert(domain->domain_ + ":" + op->GetName(op));
}
}
}
std::shared_ptr<KernelRegistry> VitisAIExecutionProvider::GetKernelRegistry() const { return get_kernel_registry_vitisaiep(); }
// This method is called after both `GetComputeCapabilityOps()` and `Compile()`.
@ -60,7 +50,7 @@ const InlinedVector<const Node*> VitisAIExecutionProvider::GetEpContextNodes() c
return ep_context_node_ptrs;
}
std::vector<std::unique_ptr<ComputeCapability>> VitisAIExecutionProvider::GetCapability(
const onnxruntime::GraphViewer& graph_viewer, const IKernelLookup& /*kernel_lookup*/) const {
const onnxruntime::GraphViewer& graph_viewer, const IKernelLookup& kernel_lookup) const {
if (graph_viewer.IsSubgraph()) {
// VITIS AI EP not support sungraph. Assigned to CPU.
return {};
@ -70,7 +60,7 @@ std::vector<std::unique_ptr<ComputeCapability>> VitisAIExecutionProvider::GetCap
return {};
}
execution_providers_ = std::make_unique<my_ep_t>(compile_onnx_model(graph_viewer, *GetLogger(), info_));
auto result = vaip::GetComputeCapabilityOps(graph_viewer, execution_providers_.get(), vitisai_optypes_);
auto result = vaip::GetComputeCapabilityOps(graph_viewer, execution_providers_.get(), kernel_lookup);
size_t index = 0u;
for (auto& ep : **execution_providers_) {
result.emplace_back(vaip::XirSubgraphToComputeCapability1(graph_viewer, ep.get(), index));

View file

@ -41,7 +41,6 @@ class VitisAIExecutionProvider : public IExecutionProvider {
const InlinedVector<const Node*> GetEpContextNodes() const override;
private:
void CreateKernelRegistry();
using my_ep_t = vaip_core::DllSafe<std::vector<std::unique_ptr<vaip_core::ExecutionProvider>>>;
using my_ep_uptr_t = std::shared_ptr<my_ep_t>;
// we have to hide the implementation by forward declaration.

View file

@ -813,6 +813,14 @@ struct ProviderHostImpl : ProviderHost {
}
ONNX_NAMESPACE::RegisterSchema(schema, ORT_API_VERSION);
}
const ONNX_NAMESPACE::OpSchema* GetSchema(const std::string& name, const int maxInclusiveVersion, const std::string& domain) override {
return ONNX_NAMESPACE::OpSchemaRegistry::Instance()->GetSchema(name, maxInclusiveVersion, domain);
}
const std::string& OpSchema__inputs__GetName(const ONNX_NAMESPACE::OpSchema* p, const size_t i) override { return p->inputs()[i].GetName(); }
const std::string& OpSchema__inputs__GetTypeStr(const ONNX_NAMESPACE::OpSchema* p, const size_t i) override { return p->inputs()[i].GetTypeStr(); }
const std::string& OpSchema__outputs__GetName(const ONNX_NAMESPACE::OpSchema* p, const size_t i) override { return p->outputs()[i].GetName(); }
const std::string& OpSchema__outputs__GetTypeStr(const ONNX_NAMESPACE::OpSchema* p, const size_t i) override { return p->outputs()[i].GetTypeStr(); }
const ONNX_NAMESPACE::TypeConstraintMap& OpSchema__typeConstraintMap(const ONNX_NAMESPACE::OpSchema* p) const override { return p->typeConstraintMap(); }
// ConfigOptions (wrapped)
std::optional<std::string> ConfigOptions__GetConfigEntry(const ConfigOptions* p, const std::string& config_key) override {
@ -946,6 +954,7 @@ struct ProviderHostImpl : ProviderHost {
MLDataType DataTypeImpl__GetType_Int4x2() override { return DataTypeImpl::GetType<Int4x2>(); }
MLDataType DataTypeImpl__GetType_UInt4x2() override { return DataTypeImpl::GetType<UInt4x2>(); }
MLDataType DataTypeImpl__GetTensorTypeFromOnnxType(int onnx_type) override { return DataTypeImpl::TensorTypeFromONNXEnum(onnx_type)->AsTensorType(); }
MLDataType DataTypeImpl__GetTensorType_bool() override { return DataTypeImpl::GetTensorType<bool>(); }
MLDataType DataTypeImpl__GetTensorType_int8() override { return DataTypeImpl::GetTensorType<int8_t>(); }
MLDataType DataTypeImpl__GetTensorType_uint8() override { return DataTypeImpl::GetTensorType<uint8_t>(); }