From 41ef0ac96c3268ac922d9b22e8d22bf9b72501a0 Mon Sep 17 00:00:00 2001 From: Leo Fang Date: Tue, 29 Oct 2024 03:25:47 +0000 Subject: [PATCH 1/2] switch to use new bindings --- CMakeLists.txt | 16 +- pynvjitlink/_nvjitlinklib.cpp | 368 -------------------------- pynvjitlink/_nvjitlinklib.py | 59 +++++ pynvjitlink/tests/test_pynvjitlink.py | 18 +- 4 files changed, 69 insertions(+), 392 deletions(-) delete mode 100644 pynvjitlink/_nvjitlinklib.cpp create mode 100644 pynvjitlink/_nvjitlinklib.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 87aaed4..7e89214 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -5,21 +5,7 @@ cmake_minimum_required(VERSION 3.26.4 FATAL_ERROR) project( pynvjitlink VERSION ${SKBUILD_PROJECT_VERSION} - LANGUAGES CXX CUDA + LANGUAGES CXX ) find_package(Python COMPONENTS Interpreter Development REQUIRED) - -Python_add_library(_nvjitlinklib MODULE pynvjitlink/_nvjitlinklib.cpp WITH_SOABI) - -find_package( - # Require CUDA 12.2 Update 2 to avoid nvjitlink bugs - CUDAToolkit 12.2.140 REQUIRED -) -target_link_libraries(_nvjitlinklib PRIVATE CUDA::nvJitLink_static CUDA::nvptxcompiler_static) - -target_compile_options(_nvjitlinklib PRIVATE -Werror -Wall) - -target_compile_features(_nvjitlinklib PRIVATE cxx_std_11) - -install(TARGETS _nvjitlinklib LIBRARY DESTINATION pynvjitlink) diff --git a/pynvjitlink/_nvjitlinklib.cpp b/pynvjitlink/_nvjitlinklib.cpp deleted file mode 100644 index 0bc9eb5..0000000 --- a/pynvjitlink/_nvjitlinklib.cpp +++ /dev/null @@ -1,368 +0,0 @@ -/* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#define PY_SSIZE_T_CLEAN -#include "nvJitLink.h" -#include -#include - -static const char *nvJitLinkGetErrorEnum(nvJitLinkResult error) { - switch (error) { - case NVJITLINK_SUCCESS: - return "NVJITLINK_SUCCESS"; - - case NVJITLINK_ERROR_UNRECOGNIZED_OPTION: - return "NVJITLINK_ERROR_UNRECOGNIZED_OPTION"; - - case NVJITLINK_ERROR_MISSING_ARCH: - return "NVJITLINK_ERROR_MISSING_ARCH"; - - case NVJITLINK_ERROR_INVALID_INPUT: - return "NVJITLINK_ERROR_INVALID_INPUT"; - - case NVJITLINK_ERROR_PTX_COMPILE: - return "NVJITLINK_ERROR_PTX_COMPILE"; - - case NVJITLINK_ERROR_NVVM_COMPILE: - return "NVJITLINK_ERROR_NVVM_COMPILE"; - - case NVJITLINK_ERROR_INTERNAL: - return "NVJITLINK_ERROR_INTERNAL"; - - default: - return ""; - } -} - -static void set_exception(PyObject *exception_type, const char *message_format, - nvJitLinkResult error) { - char exception_message[256]; - sprintf(exception_message, message_format, nvJitLinkGetErrorEnum(error)); - - PyErr_SetString(exception_type, exception_message); -} - -static PyObject *nvjitlink_version(PyObject *self, PyObject *Py_UNUSED(args)) { - unsigned int major; - unsigned int minor; - - nvJitLinkResult res = nvJitLinkVersion(&major, &minor); - - if (res != NVJITLINK_SUCCESS) { - set_exception(PyExc_RuntimeError, "%s error when calling nvJitLinkVersion", - res); - return nullptr; - } - - PyObject *py_version = PyTuple_New(2); - PyObject *py_major = PyLong_FromUnsignedLong(major); - PyObject *py_minor = PyLong_FromUnsignedLong(minor); - if (!py_version || !py_major || !py_minor) { - PyErr_SetString(PyExc_RuntimeError, "Failed to create version tuple"); - if (py_major) { - Py_DecRef(py_major); - } - if (py_minor) { - Py_DecRef(py_minor); - } - if (py_version) { - Py_DecRef(py_version); - } - return nullptr; - } - - PyTuple_SetItem(py_version, 0, py_major); - PyTuple_SetItem(py_version, 1, py_minor); - return py_version; -} - -static PyObject *create(PyObject *self, PyObject *args) { - PyObject *ret = nullptr; - const char **jitlink_options; - nvJitLinkHandle *jitlink; - - Py_ssize_t n_args = PyTuple_Size(args); - - try { - jitlink_options = new const char *[n_args]; - } catch (const std::bad_alloc &) { - PyErr_NoMemory(); - return nullptr; - } - - for (Py_ssize_t i = 0; i < n_args; ++i) { - PyObject *py_option = PyTuple_GetItem(args, i); - if (!PyUnicode_Check(py_option)) { - PyErr_SetString(PyExc_TypeError, - "Expecting only strings for jitlink args"); - delete[] jitlink_options; - return nullptr; - } - - jitlink_options[i] = PyUnicode_AsUTF8AndSize(py_option, nullptr); - } - - try { - jitlink = new nvJitLinkHandle; - } catch (const std::bad_alloc &) { - PyErr_NoMemory(); - delete[] jitlink_options; - return nullptr; - } - - nvJitLinkResult res = nvJitLinkCreate(jitlink, n_args, jitlink_options); - if (res != NVJITLINK_SUCCESS) { - set_exception(PyExc_RuntimeError, "%s error when calling nvJitLinkCreate", - res); - goto error; - } - - if ((ret = PyLong_FromUnsignedLongLong((unsigned long long)jitlink)) == - nullptr) { - // Attempt to destroy the linker - since we're already in an error - // condition, there's no point in checking the return code and taking any - // further action based on it though. - nvJitLinkDestroy(jitlink); - goto error; - } - - delete[] jitlink_options; - return ret; - -error: - delete jitlink; - delete[] jitlink_options; - return nullptr; -} - -static PyObject *destroy(PyObject *self, PyObject *args) { - nvJitLinkHandle *jitlink; - if (!PyArg_ParseTuple(args, "K", &jitlink)) - return nullptr; - - nvJitLinkResult res = nvJitLinkDestroy(jitlink); - - if (res != NVJITLINK_SUCCESS) { - set_exception(PyExc_RuntimeError, "%s error when calling nvJitLinkDestroy", - res); - return nullptr; - } - - delete jitlink; - - Py_RETURN_NONE; -} - -static PyObject *add_data(PyObject *self, PyObject *args) { - nvJitLinkHandle *jitlink; - nvJitLinkInputType input_type; - Py_buffer buf; - const char *name; - - if (!PyArg_ParseTuple(args, "Kiy*s", &jitlink, &input_type, &buf, &name)) { - return nullptr; - } - - const void *data = buf.buf; - size_t size = buf.len; - nvJitLinkResult res = - nvJitLinkAddData(*jitlink, input_type, data, size, name); - - if (res != NVJITLINK_SUCCESS) { - set_exception(PyExc_RuntimeError, "%s error when calling nvJitLinkAddData", - res); - return nullptr; - } - - Py_RETURN_NONE; -} - -static PyObject *add_file(PyObject *self, PyObject *args) { - set_exception(PyExc_NotImplementedError, "Unimplemented", NVJITLINK_SUCCESS); - - return nullptr; -} -static PyObject *complete(PyObject *self, PyObject *args) { - nvJitLinkHandle *jitlink; - if (!PyArg_ParseTuple(args, "K", &jitlink)) - return nullptr; - - nvJitLinkResult res = nvJitLinkComplete(*jitlink); - - if (res != NVJITLINK_SUCCESS) { - set_exception(PyExc_RuntimeError, "%s error when calling nvJitLinkComplete", - res); - return nullptr; - } - - Py_RETURN_NONE; - - set_exception(PyExc_NotImplementedError, "Unimplemented", NVJITLINK_SUCCESS); - - return nullptr; -} -static PyObject *get_error_log(PyObject *self, PyObject *args) { - nvJitLinkHandle *jitlink; - if (!PyArg_ParseTuple(args, "K", &jitlink)) - return nullptr; - - size_t error_log_size; - nvJitLinkResult res = nvJitLinkGetErrorLogSize(*jitlink, &error_log_size); - if (res != NVJITLINK_SUCCESS) { - set_exception(PyExc_RuntimeError, - "%s error when calling nvJitLinkGetErrorLogSize", res); - return nullptr; - } - - // The size returned doesn't include a trailing null byte - char *error_log = new char[error_log_size + 1]; - res = nvJitLinkGetErrorLog(*jitlink, error_log); - if (res != NVJITLINK_SUCCESS) { - set_exception(PyExc_RuntimeError, - "%s error when calling nvJitLinkGetErrorLog", res); - return nullptr; - } - - PyObject *py_log = PyUnicode_FromStringAndSize(error_log, error_log_size); - // Once we've copied the log to a Python object we can delete it - we don't - // need to check whether creation of the Unicode object succeeded, because we - // delete the log either way. - delete[] error_log; - - return py_log; -} - -static PyObject *get_info_log(PyObject *self, PyObject *args) { - nvJitLinkHandle *jitlink; - if (!PyArg_ParseTuple(args, "K", &jitlink)) - return nullptr; - - size_t info_log_size; - nvJitLinkResult res = nvJitLinkGetInfoLogSize(*jitlink, &info_log_size); - if (res != NVJITLINK_SUCCESS) { - set_exception(PyExc_RuntimeError, - "%s error when calling nvJitLinkGetInfoLogSize", res); - return nullptr; - } - - // The size returned doesn't include a trailing null byte - char *info_log = new char[info_log_size + 1]; - res = nvJitLinkGetInfoLog(*jitlink, info_log); - if (res != NVJITLINK_SUCCESS) { - set_exception(PyExc_RuntimeError, - "%s error when calling nvJitLinkGetInfoLog", res); - return nullptr; - } - - PyObject *py_log = PyUnicode_FromStringAndSize(info_log, info_log_size); - // Once we've copied the log to a Python object we can delete it - we don't - // need to check whether creation of the Unicode object succeeded, because we - // delete the log either way. - delete[] info_log; - - return py_log; -} - -static PyObject *get_linked_ptx(PyObject *self, PyObject *args) { - nvJitLinkHandle *jitlink; - if (!PyArg_ParseTuple(args, "K", &jitlink)) - return nullptr; - - size_t linked_ptx_size; - nvJitLinkResult res = nvJitLinkGetLinkedPtxSize(*jitlink, &linked_ptx_size); - if (res != NVJITLINK_SUCCESS) { - set_exception(PyExc_RuntimeError, - "%s error when calling nvJitLinkGetLinkedPtxSize", res); - return nullptr; - } - - char *linked_ptx = new char[linked_ptx_size]; - res = nvJitLinkGetLinkedPtx(*jitlink, linked_ptx); - if (res != NVJITLINK_SUCCESS) { - set_exception(PyExc_RuntimeError, - "%s error when calling nvJitLinkGetLinkedPtx", res); - } - - PyObject *py_ptx = PyBytes_FromStringAndSize(linked_ptx, linked_ptx_size); - // Once we've copied the compiled program to a Python object we can delete it - // - we don't need to check whether creation of the Unicode object succeeded, - // because we delete the compiled program either way. - delete[] linked_ptx; - - return py_ptx; -} -static PyObject *get_linked_cubin(PyObject *self, PyObject *args) { - nvJitLinkHandle *jitlink; - if (!PyArg_ParseTuple(args, "K", &jitlink)) - return nullptr; - - size_t linked_cubin_size; - nvJitLinkResult res = - nvJitLinkGetLinkedCubinSize(*jitlink, &linked_cubin_size); - if (res != NVJITLINK_SUCCESS) { - set_exception(PyExc_RuntimeError, - "%s error when calling nvJitLinkGetLinkedCubinSize", res); - return nullptr; - } - - char *linked_cubin = new char[linked_cubin_size]; - res = nvJitLinkGetLinkedCubin(*jitlink, linked_cubin); - if (res != NVJITLINK_SUCCESS) { - set_exception(PyExc_RuntimeError, - "%s error when calling nvJitLinkGetLinkedCubin", res); - } - - PyObject *py_cubin = - PyBytes_FromStringAndSize(linked_cubin, linked_cubin_size); - // Once we've copied the compiled program to a Python object we can delete it - // - we don't need to check whether creation of the Unicode object succeeded, - // because we delete the compiled program either way. - delete[] linked_cubin; - - return py_cubin; -} - -static PyMethodDef ext_methods[] = { - {"nvjitlink_version", (PyCFunction)nvjitlink_version, METH_NOARGS, - "Returns the nvJitLink version"}, - {"create", (PyCFunction)create, METH_VARARGS, - "Returns a handle to a new nvJitLink object"}, - {"destroy", (PyCFunction)destroy, METH_VARARGS, - "Given a handle, destroy an nvJitLink object"}, - {"add_data", (PyCFunction)add_data, METH_VARARGS, - "Add data to the link for the given handle"}, - {"add_file", (PyCFunction)add_file, METH_VARARGS, - "Add a file to the link for the given handle"}, - {"complete", (PyCFunction)complete, METH_VARARGS, - "Given a handle, complete the link"}, - {"get_error_log", (PyCFunction)get_error_log, METH_VARARGS, - "Given a handle, return the error log"}, - {"get_info_log", (PyCFunction)get_info_log, METH_VARARGS, - "Given a handle, return the info log"}, - {"get_linked_ptx", (PyCFunction)get_linked_ptx, METH_VARARGS, - "Given a handle, provide the linked PTX"}, - {"get_linked_cubin", (PyCFunction)get_linked_cubin, METH_VARARGS, - "Given a handle, provide the linked cubin"}, - {nullptr}}; - -static struct PyModuleDef moduledef = { - PyModuleDef_HEAD_INIT, "pynvjitlink", - "Provides access to nvJitLink API methods", -1, ext_methods}; - -PyMODINIT_FUNC PyInit__nvjitlinklib(void) { - PyObject *m = PyModule_Create(&moduledef); - return m; -} diff --git a/pynvjitlink/_nvjitlinklib.py b/pynvjitlink/_nvjitlinklib.py new file mode 100644 index 0000000..dab12a5 --- /dev/null +++ b/pynvjitlink/_nvjitlinklib.py @@ -0,0 +1,59 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. + +# A shim layer to use new bindings from CUDA Python +from cuda.bindings import nvjitlink +from cuda.bindings.nvjitlink import nvJitLinkError + + +def nvjitlink_version(): + return nvjitlink.version() + + +def create(*options): + return nvjitlink.create(len(options), options) + + +def destroy(handle): + return nvjitlink.destroy(handle) + + +def add_data(handle, input_type, data, filename): + nvjitlink.add_data(handle, input_type, data, len(data), filename) + + +def add_file(*args, **kwargs): + raise NotImplementedError("seems unused by pynvjitlink") + + +def complete(handle): + nvjitlink.complete(handle) + + +def get_error_log(handle): + log_size = nvjitlink.get_error_log_size(handle) + log_size += 1 + print(f"{log_size=}") + log = bytearray(log_size) + nvjitlink.get_error_log(handle, log) + return log.decode() + + +def get_info_log(handle): + log_size = nvjitlink.get_info_log_size(handle) + log = bytearray(log_size) + nvjitlink.get_info_log(handle, log) + return log.decode() + + +def get_linked_ptx(handle): + ptx_size = nvjitlink.get_linked_ptx_size(handle) + ptx = bytearray(ptx_size) + nvjitlink.get_linked_ptx(handle, ptx) + return ptx.decode() + + +def get_linked_cubin(handle): + cubin_size = nvjitlink.get_linked_cubin_size(handle) + cubin = bytearray(cubin_size) + nvjitlink.get_linked_cubin(handle, cubin) + return cubin diff --git a/pynvjitlink/tests/test_pynvjitlink.py b/pynvjitlink/tests/test_pynvjitlink.py index ad0c5b7..af047c8 100644 --- a/pynvjitlink/tests/test_pynvjitlink.py +++ b/pynvjitlink/tests/test_pynvjitlink.py @@ -7,24 +7,23 @@ def test_create_no_arch_error(): - # nvjitlink expects at least the architecture to be specified. - with pytest.raises(RuntimeError, match="NVJITLINK_ERROR_MISSING_ARCH error"): + with pytest.raises(_nvjitlinklib.nvJitLinkError, match="ERROR_INVALID_INPUT"): _nvjitlinklib.create() def test_invalid_arch_error(): # sm_XX is not a valid architecture - with pytest.raises(RuntimeError, match="NVJITLINK_ERROR_UNRECOGNIZED_OPTION error"): + with pytest.raises(_nvjitlinklib.nvJitLinkError, match="ERROR_UNRECOGNIZED_OPTION"): _nvjitlinklib.create("-arch=sm_XX") def test_unrecognized_option_error(): - with pytest.raises(RuntimeError, match="NVJITLINK_ERROR_UNRECOGNIZED_OPTION error"): + with pytest.raises(_nvjitlinklib.nvJitLinkError, match="ERROR_UNRECOGNIZED_OPTION"): _nvjitlinklib.create("-fictitious_option") def test_invalid_option_type_error(): - with pytest.raises(TypeError, match="Expecting only strings"): + with pytest.raises(TypeError, match="an integer is required"): _nvjitlinklib.create("-arch", 53) @@ -74,14 +73,15 @@ def test_get_error_log(undefined_extern_cubin, gpu_arch_flag): filename, data = undefined_extern_cubin input_type = InputType.CUBIN.value _nvjitlinklib.add_data(handle, input_type, data, filename) - with pytest.raises(RuntimeError): + with pytest.raises(_nvjitlinklib.nvJitLinkError): + # FIXME: For some reason this API would leak the error log to stderr _nvjitlinklib.complete(handle) error_log = _nvjitlinklib.get_error_log(handle) _nvjitlinklib.destroy(handle) assert ( "Undefined reference to '_Z5undefff' " "in 'undefined_extern.cubin'" in error_log - ) + ), f"{error_log=}" def test_get_info_log(device_functions_cubin, gpu_arch_flag): @@ -116,7 +116,7 @@ def test_get_linked_cubin_link_not_complete_error( filename, data = device_functions_cubin input_type = InputType.CUBIN.value _nvjitlinklib.add_data(handle, input_type, data, filename) - with pytest.raises(RuntimeError, match="NVJITLINK_ERROR_INTERNAL error"): + with pytest.raises(_nvjitlinklib.nvJitLinkError, match="ERROR_INTERNAL"): _nvjitlinklib.get_linked_cubin(handle) _nvjitlinklib.destroy(handle) @@ -159,7 +159,7 @@ def test_get_linked_ptx_link_not_complete_error( filename, data = device_functions_ltoir_object input_type = InputType.OBJECT.value _nvjitlinklib.add_data(handle, input_type, data, filename) - with pytest.raises(RuntimeError, match="NVJITLINK_ERROR_INTERNAL error"): + with pytest.raises(_nvjitlinklib.nvJitLinkError, match="ERROR_INTERNAL"): _nvjitlinklib.get_linked_ptx(handle) _nvjitlinklib.destroy(handle) From 2c362d0dd6ca8b8b45e537b99cda792c9cd27d50 Mon Sep 17 00:00:00 2001 From: Leo Fang Date: Tue, 29 Oct 2024 03:51:55 +0000 Subject: [PATCH 2/2] update test_patch.py --- pynvjitlink/tests/test_patch.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pynvjitlink/tests/test_patch.py b/pynvjitlink/tests/test_patch.py index f6db86e..69efb23 100644 --- a/pynvjitlink/tests/test_patch.py +++ b/pynvjitlink/tests/test_patch.py @@ -5,7 +5,8 @@ import pytest from numba import cuda -from pynvjitlink import NvJitLinkError, patch +from pynvjitlink import patch +from pynvjitlink._nvjitlinklib import nvJitLinkError from pynvjitlink.patch import ( PatchedLinker, _numba_version_ok, @@ -48,7 +49,7 @@ def test_create_no_cc_error(): def test_invalid_arch_error(): # CC 0.0 is not a valid compute capability with pytest.raises( - NvJitLinkError, match="NVJITLINK_ERROR_UNRECOGNIZED_OPTION error" + nvJitLinkError, match="ERROR_UNRECOGNIZED_OPTION" ): PatchedLinker(cc=(0, 0))