text: // Copyright (c) Microsoft Corporation. All rights reserved.
  // Licensed under the MIT License.
  
  #pragma once
  
  #include <unordered_map>
  #include <functional>
  #include <variant>
  #include <optional>
  
  #include "core/framework/op_kernel.h"
  
  struct AbstractOperatorDesc;
  interface IMLOperatorTensor;
  
  namespace onnxruntime
  {
      class KernelDef;
      class Node;
  }
  
  namespace Windows::AI::MachineLearning::Adapter
  {
      interface __declspec(uuid("5b19a18a-5ed5-4df2-a363-21b89380a698"))
      IWinmlExecutionProvider : public IUnknown
      {
      public:
          // Hold a reference to an object until preceding work in the queue is complete.  This
          // only needs to be handled by providers which hide the asynchronous nature of 
          // computation, and involve resoures which cannot be automatically by work in the
          // the provider's underlying queues.
          virtual void QueueReference(IUnknown *object) = 0;
  
          virtual void GetShadowCopyIfRequired(
              bool isInternalOperator,
              IUnknown* data,
              IUnknown** dataCopy) const = 0;
  
          virtual void GetABIDataInterface(
              bool isInternalOperator,
              IUnknown* data,
              IUnknown** abiData) const = 0;
          
          virtual uint64_t TryGetPooledAllocationId(
              IUnknown* data,
              bool isInternalOperator) = 0;
  
          virtual void GetABIExecutionInterface(
              bool isInternalOperator,
              IUnknown** abiExecutionObject) const = 0;
  
          // Whether TransitionResourcesForOperator should be called before and after executing
          // an operator registered to this provider with the specified flags
          virtual bool TransitionsRequiredForOperator(bool isInternalOperator) = 0;
  
          // If TransitionsRequiredForOperator returns true, should be called before and after executing
          // an operator to transition its resources to and from the appropriate state.
          virtual void TransitionResourcesForOperator(
              bool isBeforeOp,
              uint32_t resourceCount,
              IUnknown** resources) = 0;
  
          // Waits for flushed work, discards unflushed work, and discards associated references to 
          // prevent circular references.  Must be the last call on the object before destruction.
          virtual void Close() = 0;
      };
  
      using MLOperatorTensorGetter = std::function<Microsoft::WRL::ComPtr<IMLOperatorTensor>(uint32_t index)>;
  
      struct DmlOperatorParams
      {
          Microsoft::WRL::ComPtr<IDMLOperator> op;
          std::unique_ptr<AbstractOperatorDesc> desc;
      };
  
      // This is the counterpart to the MLOperatorKernelDmlProperties ABI struct which owns its memory and uses containers.
      struct DmlGraphNodeCreateInfo
      {
          bool initialized = false;
  
          // Mapping between DML in/out indices and kernel in/out indices
          std::vector<uint32_t> kernelInputIndices;
          std::vector<uint32_t> kernelOutputIndices;
  
          Microsoft::WRL::ComPtr<IDMLOperator> op;
          std::unique_ptr<AbstractOperatorDesc> desc;
  
          bool allowHalfPrecisionComputation = false;
      };
  
      using GraphNodeFactory = std::function<void(
          const onnxruntime::Node& node, 
          MLOperatorTensorGetter& constantInputGetter,
          const void* executionHandle,
          DmlGraphNodeCreateInfo* graphNodeCreateInfo
          )>;
  
      struct GraphNodeFactoryRegistration
      {
          GraphNodeFactory factory;
          std::optional<uint32_t> requiredInputCount;
      };
  
      using KernelSupportQuery = std::function<bool(const onnxruntime::Node& node)>;
  
      struct InternalRegistrationInfo
      {
          std::vector<uint32_t> requiredConstantCpuInputs;
          std::optional<GraphNodeFactoryRegistration> graphNodeFactoryRegistration;
          KernelSupportQuery supportQuery;
      };
  
      using InternalRegistrationInfoMap = std::unordered_map<onnxruntime::KernelDef*, std::shared_ptr<InternalRegistrationInfo>>;
  }