Patch the build system to better support conda-build

install/cupy_builder/_environment.py

@@ -33,5 +33,7 @@ def get_nvtx_path() -> Optional[str]:
         nvtx = candidates[-1]
         print(f'Using NVTX at: {nvtx}')
         return nvtx
+    if os.environ.get('CONDA_BUILD', '0') == '1':
+        return os.environ['PREFIX']
     print('NVTX could not be found')
     return None

install/cupy_builder/_features.py

@@ -147,7 +147,6 @@ def _from_dict(d: Dict[str, Any], ctx: Context) -> Feature:
 
 # Libraries required for cudart_static
 _cudart_static_libs = (
-    ['cudart_static'] +
     (['pthread', 'rt', 'dl'] if sys.platform == 'linux' else [])
 )
 
@@ -273,6 +272,7 @@ def get_features(ctx: Context) -> Dict[str, Feature]:
             'cub/util_namespace.cuh',  # dummy
         ],
         'libraries': list(_cudart_static_libs),
+        'static_libraries': ['cudart_static'],
         'check_method': build.check_cub_version,
         'version_method': build.get_cub_version,
     }
@@ -292,6 +292,7 @@ def get_features(ctx: Context) -> Dict[str, Feature]:
             'cuda',
             'nvrtc',
         ] + _cudart_static_libs,
+        'static_libraries': ['cudart_static'],
         'check_method': build.check_jitify_version,
         'version_method': build.get_jitify_version,
     }
@@ -308,6 +309,7 @@ def get_features(ctx: Context) -> Dict[str, Feature]:
         'libraries': [
             'curand',
         ] + _cudart_static_libs,
+        'static_libraries': ['cudart_static'],
     }
     HIP_random = {
         'name': 'random',
@@ -394,6 +396,7 @@ def get_features(ctx: Context) -> Dict[str, Feature]:
             'thrust/version.h',
         ],
         'libraries': list(_cudart_static_libs),
+        'static_libraries': ['cudart_static'],
         'check_method': build.check_thrust_version,
         'version_method': build.get_thrust_version,
     }
@@ -460,6 +463,7 @@ def __init__(self, ctx: Context):
             # CUDA Toolkit
             ['cublas', 'cufft', 'curand', 'cusparse']
         )
+        self.static_libraries = ['cudart_static']
         self._version = self._UNDETERMINED
 
     def configure(self, compiler: Any, settings: Any) -> bool:

install/cupy_builder/cupy_setup_build.py

@@ -3,6 +3,7 @@
 import copy
 from distutils import ccompiler
 from distutils import sysconfig
+import logging
 import os
 import shutil
 import sys
@@ -260,20 +261,32 @@ def _rpath_base():
 def _find_static_library(name: str) -> str:
     if PLATFORM_LINUX:
         filename = f'lib{name}.a'
-        libdir = 'lib64'
+        if (int(os.environ.get('CONDA_BUILD_CROSS_COMPILATION', 0)) == 1 and
+                os.environ.get('CONDA_OVERRIDE_CUDA', '0').startswith('11')):
+            # CUDA 11 on conda-forge has an ad hoc layout to support cross
+            # compiling
+            libdirs = ['lib']
+            cuda_path = (f'{build.get_cuda_path()}/targets/'
+                         f'{build.conda_get_target_name()}/')
+        else:
+            libdirs = ['lib64', 'lib']
+            cuda_path = build.get_cuda_path()
     elif PLATFORM_WIN32:
         filename = f'{name}.lib'
-        libdir = 'lib\\x64'
+        libdirs = ['lib\\x64', 'lib']
+        cuda_path = build.get_cuda_path()
     else:
         raise Exception('not supported on this platform')
 
-    cuda_path = build.get_cuda_path()
+    logging.debug(f"{cuda_path=}")
     if cuda_path is None:
         raise Exception(f'Could not find {filename}: CUDA path unavailable')
-    path = os.path.join(cuda_path, libdir, filename)
-    if not os.path.exists(path):
+    for libdir in libdirs:
+        path = os.path.join(cuda_path, libdir, filename)
+        if os.path.exists(path):
+            return path
+    else:
         raise Exception(f'Could not find {filename}: {path} does not exist')
-    return path
 
 
 def make_extensions(ctx: Context, compiler, use_cython):
@@ -317,16 +330,42 @@ def make_extensions(ctx: Context, compiler, use_cython):
         settings['define_macros'].append(('__HIP_PLATFORM_HCC__', '1'))
     settings['define_macros'].append(('CUPY_CACHE_KEY', ctx.cupy_cache_key))
 
-    available_modules = []
-    if no_cuda:
-        available_modules = [m['name'] for m in MODULES]
-    else:
-        available_modules, settings = preconfigure_modules(
-            ctx, MODULES, compiler, settings)
-        required_modules = get_required_modules(MODULES)
-        if not (set(required_modules) <= set(available_modules)):
-            raise Exception('Your CUDA environment is invalid. '
-                            'Please check above error log.')
+    try:
+        host_compiler = compiler
+        if int(os.environ.get('CONDA_BUILD_CROSS_COMPILATION', 0)) == 1:
+            os.symlink(f'{os.environ["BUILD_PREFIX"]}/x86_64-conda-linux-gnu/'
+                       'bin/x86_64-conda-linux-gnu-ld',
+                       f'{os.environ["BUILD_PREFIX"]}/bin/ld')
+        if (int(os.environ.get('CONDA_BUILD_CROSS_COMPILATION', 0)) == 1 or
+                os.environ.get('CONDA_OVERRIDE_CUDA', '0').startswith('12')):
+            # If cross-compiling, we need build_and_run() & build_shlib() to
+            # use the compiler on the build platform to generate stub files
+            # that are executable in the build environment, not the target
+            # environment.
+            compiler = ccompiler.new_compiler()
+            cc = os.environ['CC_FOR_BUILD' if PLATFORM_LINUX else 'CC']
+            cxx = os.environ['CXX_FOR_BUILD' if PLATFORM_LINUX else 'CXX']
+            compiler.compiler = [cc,]
+            compiler.compiler_cxx = [cxx,]
+            compiler.compiler_so = [cc,]
+            compiler.linker_exe = [cc, f'-B{os.environ["BUILD_PREFIX"]}/bin']
+            compiler.linker_so = [cc, f'-B{os.environ["BUILD_PREFIX"]}/bin',
+                                  '-shared']
+
+        available_modules = []
+        if no_cuda:
+            available_modules = [m['name'] for m in MODULES]
+        else:
+            available_modules, settings = preconfigure_modules(
+                ctx, MODULES, compiler, settings)
+            required_modules = get_required_modules(MODULES)
+            if not (set(required_modules) <= set(available_modules)):
+                raise Exception('Your CUDA environment is invalid. '
+                                'Please check above error log.')
+    finally:
+        compiler = host_compiler
+        if int(os.environ.get('CONDA_BUILD_CROSS_COMPILATION', 0)) == 1:
+            os.remove(f'{os.environ["BUILD_PREFIX"]}/bin/ld')
 
     ret = []
     for module in MODULES:

install/cupy_builder/install_build.py

@@ -1,7 +1,9 @@
 # mypy: ignore-errors
 
 import contextlib
+import logging
 import os
+import platform
 import re
 import shlex
 import shutil
@@ -16,6 +18,9 @@
 from cupy_builder._context import Context
 
 
+if os.environ.get('CONDA_BUILD', '0') == '1':
+    logging.basicConfig(level=logging.DEBUG)
+
 PLATFORM_LINUX = sys.platform.startswith('linux')
 PLATFORM_WIN32 = sys.platform.startswith('win32')
 
@@ -692,9 +697,75 @@ def get_cusparselt_version(formatted=False):
     return _cusparselt_version
 
 
+def conda_get_target_name():
+    out = None
+    if PLATFORM_LINUX:
+        plat = platform.processor()
+        if plat == "aarch64":
+            out = "sbsa-linux"
+        else:
+            out = f"{plat}-linux"
+    elif PLATFORM_WIN32:
+        out = 'x64'
+    else:
+        assert False
+    logging.debug(f"{out=}")
+    return out
+
+
+def conda_update_dirs(include_dirs, library_dirs):
+    # Note: These hacks are needed for the dependency detection stage to
+    # function, because we create a fresh compiler instance that does not
+    # honor CFLAGS etc set in the conda-build environment.
+    include_dirs = list(include_dirs)
+    library_dirs = list(library_dirs)
+
+    if (int(os.environ.get('CONDA_BUILD_CROSS_COMPILATION', 0)) == 1):
+        # If we're cross compiling, we need to generate stub files that are
+        # executable in the build environment, not the target environment.
+        # This assumes, however, that the build/host environments see the same
+        # CUDA Toolkit.
+        if os.environ.get('CONDA_OVERRIDE_CUDA', '0').startswith('12'):
+            include_dirs.insert(
+                0,
+                f'{os.environ["BUILD_PREFIX"]}/targets/x86_64-linux/include')
+            library_dirs.insert(
+                0, f'{os.environ["BUILD_PREFIX"]}/targets/x86_64-linux/lib')
+            library_dirs.insert(0, f'{os.environ["BUILD_PREFIX"]}/lib/stubs')
+        elif os.environ.get('CONDA_OVERRIDE_CUDA', '0').startswith('11'):
+            include_dirs.append('/usr/local/cuda/include')
+            library_dirs.append('/usr/local/cuda/lib64/stubs')
+
+        # for optional dependencies
+        include_dirs.append(f'{os.environ["BUILD_PREFIX"]}/include')
+        library_dirs.append(f'{os.environ["BUILD_PREFIX"]}/lib')
+
+    if os.environ.get('CONDA_OVERRIDE_CUDA', '0').startswith('12'):
+        if PLATFORM_LINUX:
+            include_dirs.append(
+                f'{os.environ["BUILD_PREFIX"]}/targets/'
+                f'{conda_get_target_name()}/include')  # for crt headers
+            library_dirs.append(f'{os.environ["PREFIX"]}/lib/stubs')
+            # for optional dependencies
+            include_dirs.append(f'{os.environ["PREFIX"]}/include')
+            library_dirs.append(f'{os.environ["PREFIX"]}/lib')
+        else:
+            # there seems to be no stubs for windows
+            # for optional dependencies
+            include_dirs.append(
+                f'{os.environ["LIBRARY_INC"]}')  # $PREFIX/Library/include
+            library_dirs.append(
+                f'{os.environ["LIBRARY_LIB"]}')  # $PREFIX/Library/lib
+
+    return include_dirs, library_dirs
+
+
 def build_shlib(compiler, source, libraries=(),
                 include_dirs=(), library_dirs=(), define_macros=None,
                 extra_compile_args=()):
+    include_dirs, library_dirs = conda_update_dirs(include_dirs, library_dirs)
+    logging.debug(include_dirs)
+
     with _tempdir() as temp_dir:
         fname = os.path.join(temp_dir, 'a.cpp')
         with open(fname, 'w') as f:
@@ -720,6 +791,9 @@ def build_shlib(compiler, source, libraries=(),
 def build_and_run(compiler, source, libraries=(),
                   include_dirs=(), library_dirs=(), define_macros=None,
                   extra_compile_args=()):
+    include_dirs, library_dirs = conda_update_dirs(include_dirs, library_dirs)
+    logging.debug(include_dirs)
+
     with _tempdir() as temp_dir:
         fname = os.path.join(temp_dir, 'a.cpp')
         with open(fname, 'w') as f: