Repository: microsoft/onnxruntime
File: onnxruntime/core/providers/nv_tensorrt_rtx/nv_provider_factory.cc (C++)

Code before:
    AllocOnStream = nullptr;  // AllocOnStreamImpl
    Free = FreeImpl;
    Info = InfoImpl;
    Reserve = AllocImpl;  // no special behavior for Reserve so use AllocImpl
    GetStats = nullptr;   // GetStatsImpl. The CUDA allocators don't have stats currently so we can skip.

    const OrtEpApi& ep_api = *api.GetEpApi();
    const OrtMemoryDevice* mem_device = ep_api.MemoryInfo_GetMemoryDevice(mem_info);
    uint32_t device_id = ep_api.MemoryDevice_GetDeviceId(mem_device);
    const char* name = nullptr;
    auto* status = api.MemoryInfoGetName(mem_info, &name);
    static_cast<void>(status);  // GetName never fails

    if (ep_api.MemoryDevice_GetMemoryType(mem_device) == OrtDeviceMemoryType_HOST_ACCESSIBLE) {
      allocator_ = std::make_unique<CUDAPinnedAllocator>(device_id, name);
    } else {
      allocator_ = std::make_unique<CUDAAllocator>(device_id, name);
    }
  }

  static void* ORT_API_CALL AllocImpl(struct OrtAllocator* this_, size_t size) {
    auto& impl = *static_cast<NvTrtRtxOrtAllocator*>(this_);
    return impl.allocator_->Alloc(size);
  }

  // TODO missing a bridge for this
  // static void* ORT_API_CALL AllocOnStreamImpl(struct OrtAllocator* this_, size_t size, OrtSyncStream* stream) {
  //   auto& impl = *static_cast<NvTrtRtxOrtAllocator*>(this_);
  //   Stream *internal_stream = static_cast<Stream*>(stream);
  //   return impl.allocator_->AllocOnStream(size, stream);
  // }

  static void ORT_API_CALL FreeImpl(struct OrtAllocator* this_, void* p) {
    auto& impl = *static_cast<NvTrtRtxOrtAllocator*>(this_);
    impl.allocator_->Free(p);
  }

  static const struct OrtMemoryInfo* ORT_API_CALL InfoImpl(const struct OrtAllocator* this_) {
    const NvTrtRtxOrtAllocator& impl = *static_cast<const NvTrtRtxOrtAllocator*>(this_);
    return impl.memory_info_;
  }

 private:
  const OrtMemoryInfo* memory_info_;
  std::unique_ptr<IAllocator> allocator_;
};

struct NvTrtRtxDataTransferImpl : OrtDataTransferImpl {
  NvTrtRtxDataTransferImpl(const OrtApi& ort_api_in)
      : ort_api{ort_api_in}, ep_api{*ort_api_in.GetEpApi()} {
    ort_version_supported = ORT_API_VERSION;

Review comment (suggestion):
OrtSyncStream is directly inherited from Stream so no cast should be necessary.

The definition is in core/framework/plugin_ep_stream.h

I think it would worked if you included the header.

Code after:
  NvTrtRtxOrtAllocator(const OrtMemoryInfo* mem_info, const OrtApi& api) : memory_info_{mem_info} {
    version = ORT_API_VERSION;
    Alloc = AllocImpl;
    AllocOnStream = AllocOnStreamImpl;
    Free = FreeImpl;
    Info = InfoImpl;
    Reserve = AllocImpl;  // no special behavior for Reserve so use AllocImpl
    GetStats = nullptr;   // GetStatsImpl. The CUDA allocators don't have stats currently so we can skip.

    const OrtEpApi& ep_api = *api.GetEpApi();
    const OrtMemoryDevice* mem_device = ep_api.MemoryInfo_GetMemoryDevice(mem_info);
    uint32_t device_id = ep_api.MemoryDevice_GetDeviceId(mem_device);
    const char* name = nullptr;
    auto* status = api.MemoryInfoGetName(mem_info, &name);
    static_cast<void>(status);  // GetName never fails

    if (ep_api.MemoryDevice_GetMemoryType(mem_device) == OrtDeviceMemoryType_HOST_ACCESSIBLE) {
      allocator_ = std::make_unique<CUDAPinnedAllocator>(device_id, name);
    } else {
      allocator_ = std::make_unique<CUDAAllocator>(device_id, name);
    }
  }

  static void* ORT_API_CALL AllocImpl(struct OrtAllocator* this_, size_t size) {
    auto& impl = *static_cast<NvTrtRtxOrtAllocator*>(this_);
    return impl.allocator_->Alloc(size);
  }

  static void* ORT_API_CALL AllocOnStreamImpl(struct OrtAllocator* this_, size_t size, OrtSyncStream* stream) {
    auto& impl = *static_cast<NvTrtRtxOrtAllocator*>(this_);
    return impl.allocator_->AllocOnStream(size, stream);
  }

  static void ORT_API_CALL FreeImpl(struct OrtAllocator* this_, void* p) {
    auto& impl = *static_cast<NvTrtRtxOrtAllocator*>(this_);
    impl.allocator_->Free(p);
  }

  static const struct OrtMemoryInfo* ORT_API_CALL InfoImpl(const struct OrtAllocator* this_) {
    const NvTrtRtxOrtAllocator& impl = *static_cast<const NvTrtRtxOrtAllocator*>(this_);
    return impl.memory_info_;
  }

 private:
  const OrtMemoryInfo* memory_info_;
  std::unique_ptr<IAllocator> allocator_;
};

struct NvTrtRtxDataTransferImpl : OrtDataTransferImpl {
  NvTrtRtxDataTransferImpl(const OrtApi& ort_api_in)
      : ort_api{ort_api_in}, ep_api{*ort_api_in.GetEpApi()} {