diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index fed97c3c5de..49e6e01b2f8 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -669,6 +669,25 @@ target_link_libraries( $ ) +add_library(cudfcuinit_intercept SHARED src/utilities/cuinit_intercept.cpp) +set_target_properties( + cudfcuinit_intercept + PROPERTIES BUILD_RPATH "\$ORIGIN" + INSTALL_RPATH "\$ORIGIN" + # set target compile options + CXX_STANDARD 17 + CXX_STANDARD_REQUIRED ON + CUDA_STANDARD 17 + CUDA_STANDARD_REQUIRED ON + POSITION_INDEPENDENT_CODE ON + INTERFACE_POSITION_INDEPENDENT_CODE ON +) + +if(TARGET conda_env) + target_link_libraries(cudfcuinit_intercept PRIVATE conda_env) +endif() +target_link_libraries(cudfcuinit_intercept PUBLIC CUDA::cudart cuda dl) + # Add Conda library, and include paths if specified if(TARGET conda_env) target_link_libraries(cudf PRIVATE conda_env) diff --git a/cpp/src/utilities/cuinit_intercept.cpp b/cpp/src/utilities/cuinit_intercept.cpp new file mode 100644 index 00000000000..97119efd3c2 --- /dev/null +++ b/cpp/src/utilities/cuinit_intercept.cpp @@ -0,0 +1,118 @@ +/* + * Copyright (c) 2022, 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. + */ + +#ifndef _GNU_SOURCE +#define _GNU_SOURCE +#endif +#include +#include +#include + +#if defined(__GLIBC__) && __GLIBC__ >= 2 && defined(__GLIBC_MINOR__) && __GLIBC_MINOR__ >= 1 +namespace { +static int cuInitCount{0}; +using init_t = CUresult (*)(unsigned int); +using proc_t = CUresult (*)(const char*, + void**, + int, + cuuint64_t +#if CUDA_VERSION >= 12000 + , + CUdriverProcAddressQueryResult* +#endif +); +using dlsym_t = void* (*)(void*, const char*); +static init_t original_cuInit{nullptr}; +static proc_t original_cuGetProcAddress{nullptr}; +static dlsym_t original_dlsym{nullptr}; + +static __attribute__((constructor)) void init_cuInit_hack() +{ + // Hack hack hack, only for glibc, this magic number can be found in + // glibc's sysdeps/unix/sysv/linux/x86_64/64/libc.abilist (glibc >= + // 2.34) (or libdl.abilist (glibc < 2.34). + original_dlsym = (dlsym_t)dlvsym(RTLD_NEXT, "dlsym", "GLIBC_2.2.5"); + if (original_dlsym) { + original_cuGetProcAddress = (proc_t)original_dlsym(RTLD_NEXT, "cuGetProcAddress"); + } +} + +extern "C" { +CUresult cuInit(unsigned int flags) +{ + if (!original_cuInit) { + void* ptr{nullptr}; + CUresult err = original_cuGetProcAddress("cuInit", + &ptr, + CUDA_VERSION, + CU_GET_PROC_ADDRESS_DEFAULT +#if CUDA_VERSION >= 12000 + , + nullptr +#endif + ); + if (err != CUDA_SUCCESS) { return err; } + if (ptr) { original_cuInit = (init_t)(ptr); } + } + std::cerr << "cuInit has been called " << ++cuInitCount << " times" << std::endl; + if (original_cuInit) { + return original_cuInit(flags); + } else { + return CUDA_ERROR_NOT_INITIALIZED; + } +} + +CUresult cuGetProcAddress(const char* symbol, + void** pfn, + int cudaVersion, + cuuint64_t flags +#if CUDA_VERSION >= 12000 + , + CUdriverProcAddressQueryResult* symbolStatus +#endif +) +{ + if (!original_cuGetProcAddress) { return CUDA_ERROR_NOT_SUPPORTED; } + CUresult err = original_cuGetProcAddress(symbol, + pfn, + cudaVersion, + flags +#if CUDA_VERSION >= 12000 + , + symbolStatus +#endif + ); + if (std::string{symbol} == "cuInit") { + original_cuInit = (init_t)(*pfn); + *pfn = (void*)cuInit; + } + return err; +} + +void* dlsym(void* handle, const char* name_) +{ + std::string name{name_}; + if (name == "cuInit") { + return (void*)cuInit; + } else if (name == "cuGetProcAddress") { + return (void*)cuGetProcAddress; + } else { + return original_dlsym(handle, name_); + } +} +} +} // namespace +#endif diff --git a/python/cudf/cudf/tests/test_nocuinit.py b/python/cudf/cudf/tests/test_nocuinit.py new file mode 100644 index 00000000000..e9fb5d319fb --- /dev/null +++ b/python/cudf/cudf/tests/test_nocuinit.py @@ -0,0 +1,28 @@ +# Copyright (c) 2022, NVIDIA CORPORATION. + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +location = Path(__file__) +cpp_build_dir = location / ".." / ".." / ".." / ".." / ".." / "cpp" / "build" +libintercept = (cpp_build_dir / "libcudfcuinit_intercept.so").resolve() + + +@pytest.mark.skipif( + not libintercept.exists(), + reason="libcudfcuinit_intercept.so not built, can't check for cuInit", +) +def test_import_no_cuinit(): + env = os.environ.copy() + env["RAPIDS_NO_INITIALIZE"] = "1" + env["LD_PRELOAD"] = str(libintercept) + output = subprocess.check_output( + [sys.executable, "-c", "import cudf"], + env=env, + stderr=subprocess.STDOUT, + ) + assert "cuInit has been called" not in output.decode()