onnxruntime/winml/adapter/abi_custom_registry_impl.cpp
Dwayne Robinson 8b0c2e1f3d Merged PR 6101363: Int64 prototype work for ONNX runtime DML EP
Add `support64BitTensorsViaEmulation` to the internal registration info, that informs the graph partitioner that int64 is supported via emulation, even if the device doesn't support it natively.

See further description in the corresponding WindowsAI DML PR:
https://dev.azure.com/microsoft/WindowsAI/_git/WindowsAI/pullrequest/6101182

Note a later PR will most likely *delete* this newly added flag and simplify much of the existing logic, even deleting the strides hack completely ^__^...

Related work items: #28761231
2021-06-11 07:54:30 +00:00

95 lines
3 KiB
C++

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "pch.h"
#ifdef USE_DML
#include "abi_custom_registry_impl.h"
namespace Windows::AI::MachineLearning::Adapter {
HRESULT STDMETHODCALLTYPE AbiCustomRegistryImpl::RegisterOperatorSetSchema(
const MLOperatorSetId* opSetId,
int baseline_version,
const MLOperatorSchemaDescription* const* schema,
uint32_t schemaCount,
_In_opt_ IMLOperatorTypeInferrer* typeInferrer,
_In_opt_ IMLOperatorShapeInferrer* shapeInferrer) const noexcept try {
#ifdef LAYERING_DONE
for (uint32_t i = 0; i < schemaCount; ++i) {
telemetry_helper.RegisterOperatorSetSchema(
schema[i]->name,
schema[i]->inputCount,
schema[i]->outputCount,
schema[i]->typeConstraintCount,
schema[i]->attributeCount,
schema[i]->defaultAttributeCount);
}
#endif
// Delegate to base class
return AbiCustomRegistry::RegisterOperatorSetSchema(
opSetId,
baseline_version,
schema,
schemaCount,
typeInferrer,
shapeInferrer);
}
CATCH_RETURN();
HRESULT STDMETHODCALLTYPE AbiCustomRegistryImpl::RegisterOperatorKernel(
const MLOperatorKernelDescription* opKernel,
IMLOperatorKernelFactory* operatorKernelFactory,
_In_opt_ IMLOperatorShapeInferrer* shapeInferrer) const noexcept {
return RegisterOperatorKernel(opKernel, operatorKernelFactory, shapeInferrer, nullptr, false, false, false);
}
HRESULT STDMETHODCALLTYPE AbiCustomRegistryImpl::RegisterOperatorKernel(
const MLOperatorKernelDescription* opKernel,
IMLOperatorKernelFactory* operatorKernelFactory,
_In_opt_ IMLOperatorShapeInferrer* shapeInferrer,
_In_opt_ IMLOperatorSupportQueryPrivate* supportQuery,
bool isInternalOperator,
bool canAliasFirstInput,
bool supportsGraph,
const uint32_t* requiredInputCountForGraph,
bool supportedWith64BitTensorsVia32BitStrides,
bool supportedWith64BitTensorsVia32BitStridesFromAnyEp,
bool prefer64BitTensorsDirectly,
bool support64BitTensorsViaEmulation,
_In_reads_(constantCpuInputCount) const uint32_t* requiredConstantCpuInputs,
uint32_t constantCpuInputCount) const noexcept try {
#ifdef LAYERING_DONE
// Log a custom op telemetry if the operator is not a built-in DML operator
if (!isInternalOperator) {
telemetry_helper.LogRegisterOperatorKernel(
opKernel->name,
opKernel->domain,
static_cast<int>(opKernel->executionType));
}
#endif
// Delegate to base class
return AbiCustomRegistry::RegisterOperatorKernel(
opKernel,
operatorKernelFactory,
shapeInferrer,
supportQuery,
isInternalOperator,
canAliasFirstInput,
supportsGraph,
requiredInputCountForGraph,
supportedWith64BitTensorsVia32BitStrides,
supportedWith64BitTensorsVia32BitStridesFromAnyEp,
prefer64BitTensorsDirectly,
support64BitTensorsViaEmulation,
requiredConstantCpuInputs,
constantCpuInputCount);
}
CATCH_RETURN();
} // namespace Windows::AI::MachineLearning::Adapter
#endif USE_DML