onnxruntime/onnxruntime/core/framework/func_kernel.h

#pragma once
#include "core/framework/op_kernel.h"
#include "core/framework/func_api.h"
#include "core/framework/op_kernel_context_internal.h"
#include "core/graph/function.h"

namespace onnxruntime {

void* allocate_helper_func(void* allocator, size_t alignment, size_t size);

void release_helper_func(void* allocator, void* p);

//A kernel that wrapper the ComputeFunction call generated by execution provider when fuse the sub-graph
class FunctionKernel : public OpKernel {
 public:
  //The original design is we load the dll, find the entry point and wrapper it.
  //Here for quick prototype, we keep the entry pointer in the node.
  explicit FunctionKernel(const OpKernelInfo& info) : OpKernel(info) {
    num_inputs_ = info.node().InputDefs().size();
    num_outputs_ = info.node().OutputDefs().size();
    auto status = info.GetFusedFuncs(compute_info_);
    ORT_ENFORCE(status.IsOK(), status.ErrorMessage());
    if (compute_info_->create_state_func) {
      //TODO: we are only provide host allocate method in compute context.
      //Do we need to hold the ref-counting here?
      host_allocator_ = info.GetAllocator(0, OrtMemType::OrtMemTypeDefault);
      ComputeContext context = {allocate_helper_func, release_helper_func, host_allocator_.get(),
                                info.node().Name().c_str()};
      ORT_ENFORCE(compute_info_->create_state_func(&context, &func_state_) == 0);
    }
  }

  ~FunctionKernel() override {
    if (compute_info_->release_state_func && func_state_) {
      compute_info_->release_state_func(func_state_);
    }
  }

  virtual Status Compute(OpKernelContext* context) const override {
    auto* context_internal = static_cast<OpKernelContextInternal*>(context);
    return compute_info_->compute_func(func_state_, OrtGetApiBase()->GetApi(ORT_API_VERSION),
                                       reinterpret_cast<OrtKernelContext*>(context_internal));
  }

 private:
  NodeComputeInfo* compute_info_{nullptr};
  FunctionState func_state_{nullptr};
  size_t num_inputs_;
  size_t num_outputs_;
  AllocatorPtr host_allocator_;
};
}  // namespace onnxruntime