mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-17 21:10:43 +00:00
eeef157888
winml/ was previously excluded from lintrunner config. This change includes the directory and adds the clang-format config file specific to winml/ that fits existing style. --------- Signed-off-by: Justin Chu <justinchu@microsoft.com>
94 lines
3.4 KiB
C++
94 lines
3.4 KiB
C++
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
//
|
|
// Implements a custom operator kernel which counts the number of calls to Compute(), but otherwise is a no-op.
|
|
//
|
|
|
|
#pragma once
|
|
|
|
#include "test.h"
|
|
|
|
template <typename T>
|
|
struct NullShapeInferrer : winrt::implements<NullShapeInferrer<T>, IMLOperatorShapeInferrer> {
|
|
STDMETHOD(InferOutputShapes)(IMLOperatorShapeInferenceContext* context) noexcept {
|
|
WINML_EXPECT_NO_THROW(OperatorHelper::ShapeInferenceFunction<T>(context));
|
|
return S_OK;
|
|
}
|
|
};
|
|
|
|
struct NullOperator : winrt::implements<NullOperator, IMLOperatorKernel> {
|
|
NullOperator(std::atomic<uint32_t>* callCount) : m_callCount(callCount) {}
|
|
|
|
STDMETHOD(Compute)(IMLOperatorKernelContext* context) {
|
|
winrt::com_ptr<IMLOperatorTensor> outputTensor;
|
|
WINML_EXPECT_HRESULT_SUCCEEDED(context->GetOutputTensor(0, outputTensor.put()));
|
|
|
|
++(*m_callCount);
|
|
return S_OK;
|
|
}
|
|
|
|
private:
|
|
std::atomic<uint32_t>* m_callCount;
|
|
};
|
|
|
|
struct NullOperatorFactory : winrt::implements<NullOperatorFactory, IMLOperatorKernelFactory> {
|
|
NullOperatorFactory(std::atomic<uint32_t>* callCount) : m_callCount(callCount) {}
|
|
|
|
STDMETHOD(CreateKernel)(IMLOperatorKernelCreationContext* context, IMLOperatorKernel** kernel) {
|
|
ORT_UNUSED_PARAMETER(context);
|
|
auto op = winrt::make<NullOperator>(m_callCount);
|
|
op.copy_to(kernel);
|
|
return S_OK;
|
|
}
|
|
|
|
static MLOperatorEdgeDescription CreateEdgeDescriptor(MLOperatorEdgeType type, MLOperatorTensorDataType dataType) {
|
|
ORT_UNUSED_PARAMETER(type);
|
|
MLOperatorEdgeDescription desc;
|
|
desc.edgeType = MLOperatorEdgeType::Tensor;
|
|
desc.tensorDataType = dataType;
|
|
return desc;
|
|
}
|
|
|
|
static void RegisterKernel(
|
|
const char* name,
|
|
const char* domain,
|
|
int versionSince,
|
|
winrt::com_ptr<IMLOperatorRegistry> registry,
|
|
winrt::com_ptr<IMLOperatorShapeInferrer> shapeInferrer,
|
|
std::atomic<uint32_t>* callCount
|
|
) {
|
|
MLOperatorKernelDescription kernelDescription;
|
|
kernelDescription.domain = domain;
|
|
kernelDescription.name = name;
|
|
kernelDescription.minimumOperatorSetVersion = versionSince;
|
|
kernelDescription.executionType = MLOperatorExecutionType::D3D12;
|
|
|
|
MLOperatorEdgeTypeConstrant typeConstraint;
|
|
typeConstraint.typeLabel = "T";
|
|
std::vector<MLOperatorEdgeDescription> allowedEdges{
|
|
CreateEdgeDescriptor(MLOperatorEdgeType::Tensor, MLOperatorTensorDataType::Double),
|
|
CreateEdgeDescriptor(MLOperatorEdgeType::Tensor, MLOperatorTensorDataType::Float),
|
|
CreateEdgeDescriptor(MLOperatorEdgeType::Tensor, MLOperatorTensorDataType::Float16)};
|
|
typeConstraint.allowedTypes = allowedEdges.data();
|
|
typeConstraint.allowedTypeCount = static_cast<uint32_t>(allowedEdges.size());
|
|
|
|
std::vector<MLOperatorEdgeTypeConstrant> typeConstraints{typeConstraint};
|
|
kernelDescription.typeConstraints = typeConstraints.data();
|
|
kernelDescription.typeConstraintCount = static_cast<uint32_t>(typeConstraints.size());
|
|
|
|
kernelDescription.defaultAttributes = nullptr;
|
|
kernelDescription.defaultAttributeCount = 0;
|
|
kernelDescription.options = MLOperatorKernelOptions::None;
|
|
kernelDescription.executionOptions = 0;
|
|
|
|
auto factory = winrt::make<NullOperatorFactory>(callCount);
|
|
|
|
WINML_EXPECT_HRESULT_SUCCEEDED(
|
|
registry->RegisterOperatorKernel(&kernelDescription, factory.get(), shapeInferrer.get())
|
|
);
|
|
}
|
|
|
|
private:
|
|
std::atomic<uint32_t>* m_callCount;
|
|
};
|