Create Custom Node in CUDA

[cocotdf]

Describe the issue

I am interested in creating custom nodes in ONNX. To do this, I based my work on the example code provided in the build, which allows creating the nodes CustomOpOne and CustomOpTwo. I noticed that these nodes were generated in the DLL "custom_op_library.dll".

When I tried to run the model by adding the DLL using the function "RegisterCustomOpsLibrary", the execution worked with the CPU but not with CUDA.

I would like to understand if I missed something to make it work, or if this behavior is expected.

To reproduce

For my tests, I used the model available here: https://github.com/microsoft/onnxruntime/blob/rel-1.16.0/onnxruntime/test/testdata/custom_op_library/custom_op_test.onnx.

#include <onnxruntime_cxx_api.h>
#include <onnxruntime_lite_custom_op.h>
#include <iostream>

int main() {
    try {
        // Initialisation des options de session
        Ort::SessionOptions session_options;

        // Chemin de la bibliothèque de nœuds personnalisés
        const wchar_t* custom_library_path = L"custom_op_library.dll";

        // Configuration des options de nœuds personnalisés
        Ort::CustomOpConfigs custom_op_configs;

        // Configuration du fournisseur d'exécution CUDA
        OrtCUDAProviderOptions cuda_options;
        cuda_options.device_id = 0; // Utiliser le premier GPU
        session_options.AppendExecutionProvider_CUDA(cuda_options);

        // Enregistrer la bibliothèque de nœuds personnalisés
        session_options.RegisterCustomOpsLibrary(custom_library_path, custom_op_configs);

        // Initialisation de l'environnement ONNX Runtime
        Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "test");

        // Chemin du modèle ONNX
        const wchar_t* model_path = L"custom_op_test.onnx";

        // Création de la session avec les options spécifiées
        Ort::Session session(env, model_path, session_options);

        // Vecteur pour stocker les tenseurs d'entrée
        std::vector<Ort::Value> input_tensors;

        // Données d'entrée
        std::vector<float> input_data_1(3 * 5);
        std::vector<float> input_data_2(3 * 5);

        // Forme des données d'entrée
        std::vector<int64_t> input_shape = { 3, 5 };

        // Initialiser les données d'entrée
        float value = 2.0f;
        for (auto& element : input_data_1) {
            element = value++;
        }
        value = 2.5f;
        for (auto& element : input_data_2)  {
            element = value++;
        }

        // Créer les tenseurs d'entrée
        Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
        Ort::Value input_tensor_1 = Ort::Value::CreateTensor<float>(memory_info, input_data_1.data(), input_data_1.size(),
            input_shape.data(), input_shape.size());
        Ort::Value input_tensor_2 = Ort::Value::CreateTensor<float>(memory_info, input_data_2.data(), input_data_2.size(),
            input_shape.data(), input_shape.size());

        // Ajouter les tenseurs d'entrée au vecteur
        input_tensors.push_back(std::move(input_tensor_1));
        input_tensors.push_back(std::move(input_tensor_2));

        // Noms des entrées et des sorties du modèle
        std::vector<const char*> input_names = { "input_1", "input_2" };
        std::vector<const char*> output_names = { "output" };

        // Exécuter l'inférence
        auto output_tensors = session.Run(Ort::RunOptions{ nullptr }, input_names.data(), input_tensors.data(), 2, output_names.data(), 1);

        // Récupérer et afficher les résultats
        float* output_data = output_tensors.front().GetTensorMutableData<float>();
        std::cout << "Inference results :" << std::endl << '\t';
        for (size_t i = 0; i < (3 * 5); ++i) {
            std::cout << output_data[i] << " ";
            if ((i + 1) % 5 == 0) std::cout << std::endl << '\t';
        }

        return 0;
    }
    catch (const std::exception& e) {
        std::cout << e.what() << std::endl;
        return -1;
    }
}

Urgency

No response

Platform

Windows

OS Version

10

ONNX Runtime Installation

Built from Source

ONNX Runtime Version or Commit ID

1.18.0

ONNX Runtime API

C++

Architecture

X64

Execution Provider

CUDA

Execution Provider Library Version

12.3

[skottmckay]

#21417 (comment)