ri-kernels 0.1.0__tar.gz

ri_kernels-0.1.0/.github/workflows/wheels.yml

@@ -0,0 +1,221 @@
+name: Wheels
+
+# Builds the three distributions published from this repository:
+#
+#   ri_kernels         Python code + libri_kernels.so (CPU)
+#   ri_kernels_cuda12  libri_kernels_cuda.so built against CUDA 12
+#   ri_kernels_cuda13  libri_kernels_cuda.so built against CUDA 13
+#
+# The CUDA wheels are add-ons: they contain no Python code and depend on
+# ri_kernels of the exact same version. ci/make_variant.py derives their build
+# metadata from pyproject.toml. Wheels are built on every push and pull request;
+# they are only uploaded to PyPI when a GitHub Release is published.
+
+on:
+  push:
+    branches: [main]
+    tags: ["v*"]
+  pull_request:
+  workflow_dispatch:
+  release:
+    types: [published]
+
+permissions:
+  contents: read
+
+concurrency:
+  group: wheels-${{ github.ref }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
+jobs:
+  sdist:
+    name: sdist (ri_kernels)
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - run: pipx run build --sdist
+      - uses: actions/upload-artifact@v4
+        with:
+          name: sdist-ri_kernels
+          path: dist/*.tar.gz
+
+  build_cpu:
+    name: ri_kernels (${{ matrix.arch }}, ${{ matrix.os }})
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - { os: ubuntu-latest, arch: x86_64 }
+          - { os: ubuntu-24.04-arm, arch: aarch64 }
+          - { os: macos-14, arch: arm64 }
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: pypa/cibuildwheel@v4.2.0
+        env:
+          CIBW_BUILD_VERBOSITY: 1
+          CIBW_TEST_REQUIRES: pytest numpy
+          CIBW_TEST_SOURCES: tests
+          # The test module skips itself at module level when no library is
+          # found, so assert the library is present before running it -
+          # otherwise a wheel missing libri_kernels.so would pass vacuously.
+          CIBW_TEST_COMMAND: >
+            python -c "from ri_kernels.jax_api import rfi_vis_op as m; assert m._TAB_LIB is not None, 'libri_kernels.so missing from wheel'"
+            && pytest ./tests -q
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: wheels-ri_kernels-${{ matrix.os }}-${{ matrix.arch }}
+          path: wheelhouse/*.whl
+
+  build_cuda:
+    name: ri_kernels_cuda${{ matrix.cuda.major }} (${{ matrix.platform.arch }})
+    runs-on: ${{ matrix.platform.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        # `archs` overrides the CMAKE_CUDA_ARCHITECTURES default in
+        # CMakeLists.txt, which is kept to what any CUDA 12.x can build. The
+        # published wheels target one cubin per compute-capability *major*
+        # (plus 7.5, the CUDA 13 floor): cubins are binary compatible upward
+        # within a major, so 8.0 covers 8.6/8.7/8.9, 10.0 covers 10.3 (B300)
+        # and 12.0 covers 12.1 (GB10). 10.x and 12.x are separate families and
+        # need separate cubins. This is the set the jax CUDA plugin itself
+        # ships, minus 7.0, which only CUDA 12 can still build. The trailing
+        # `-virtual` embeds PTX once, as the JIT fallback for future GPUs.
+        cuda:
+          - major: "12"
+            pkg: "12-9"
+            home: "/usr/local/cuda-12.9"
+            archs: "60-real;70-real;75-real;80-real;90-real;100-real;120-real;120-virtual"
+          - major: "13"
+            pkg: "13-0"
+            home: "/usr/local/cuda-13.0"
+            # CUDA 13 dropped Maxwell/Pascal/Volta; Turing is the floor.
+            archs: "75-real;80-real;90-real;100-real;120-real;120-virtual"
+        platform:
+          - { os: ubuntu-latest, arch: x86_64, repo_arch: x86_64 }
+          - { os: ubuntu-24.04-arm, arch: aarch64, repo_arch: sbsa }
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Rewrite metadata for the CUDA ${{ matrix.cuda.major }} variant
+        run: |
+          python -m pip install --disable-pip-version-check tomlkit
+          python ci/make_variant.py --cuda ${{ matrix.cuda.major }}
+
+      - uses: pypa/cibuildwheel@v4.2.0
+        env:
+          CIBW_BUILD_VERBOSITY: 1
+          # nvcc, the CUDA runtime headers/static lib, and CUB, installed into
+          # the manylinux container. Use the cuda-toolkit-${{ matrix.cuda.pkg }}
+          # metapackage instead if these split packages ever stop resolving.
+          CIBW_BEFORE_ALL_LINUX: >
+            dnf -y install dnf-plugins-core &&
+            dnf config-manager --add-repo
+            https://developer.download.nvidia.com/compute/cuda/repos/rhel8/${{ matrix.platform.repo_arch }}/cuda-rhel8.repo &&
+            dnf -y install
+            cuda-nvcc-${{ matrix.cuda.pkg }}
+            cuda-cudart-devel-${{ matrix.cuda.pkg }}
+            cuda-cccl-${{ matrix.cuda.pkg }}
+          # CMAKE_ARGS is forwarded to the configure step by scikit-build-core.
+          # It has to carry the architectures rather than SKBUILD_CMAKE_DEFINE,
+          # which splits its own value on `;`. The quotes keep cibuildwheel's
+          # shell-style parser from reading the separators as command
+          # terminators.
+          CIBW_ENVIRONMENT_LINUX: >
+            CUDACXX=${{ matrix.cuda.home }}/bin/nvcc
+            PATH=${{ matrix.cuda.home }}/bin:$PATH
+            CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=${{ matrix.cuda.archs }}"
+          # No test here: installing the wheel would pull ri_kernels==<version>
+          # from the index, which does not exist yet for an unreleased version. The
+          # library is smoke-tested below instead.
+          CIBW_TEST_COMMAND: ""
+
+      - name: Smoke-test the built library
+        run: python ci/check_library.py wheelhouse ri_kernels_cuda${{ matrix.cuda.major }}
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: wheels-ri_kernels_cuda${{ matrix.cuda.major }}-${{ matrix.platform.arch }}
+          path: wheelhouse/*.whl
+
+  check_release:
+    name: Check release tag
+    if: github.event_name == 'release'
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Tag must match the version in pyproject.toml
+        env:
+          TAG: ${{ github.event.release.tag_name }}
+        run: |
+          version=$(python -c "import tomllib; print(tomllib.load(open('pyproject.toml','rb'))['project']['version'])")
+          if [ "$TAG" != "v$version" ] && [ "$TAG" != "$version" ]; then
+            echo "::error::release tag '$TAG' does not match pyproject.toml version '$version'"
+            exit 1
+          fi
+          echo "Releasing version $version"
+
+  # One publish job per PyPI project: trusted publishing binds a project to
+  # a single workflow + environment, so they cannot share one.
+  publish_ri_kernels:
+    name: Publish ri_kernels
+    if: github.event_name == 'release'
+    needs: [check_release, sdist, build_cpu, build_cuda]
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi-ri_kernels
+      url: https://pypi.org/p/ri_kernels
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          pattern: wheels-ri_kernels-*
+          path: dist
+          merge-multiple: true
+      - uses: actions/download-artifact@v4
+        with:
+          name: sdist-ri_kernels
+          path: dist
+      - uses: pypa/gh-action-pypi-publish@release/v1
+
+  # The CUDA packages pin ri_kernels==<version>, so the base package has to be
+  # on the index first.
+  publish_ri_kernels_cuda12:
+    name: Publish ri_kernels_cuda12
+    if: github.event_name == 'release'
+    needs: [publish_ri_kernels]
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi-ri_kernels_cuda12
+      url: https://pypi.org/p/ri_kernels_cuda12
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          pattern: wheels-ri_kernels_cuda12-*
+          path: dist
+          merge-multiple: true
+      - uses: pypa/gh-action-pypi-publish@release/v1
+
+  publish_ri_kernels_cuda13:
+    name: Publish ri_kernels_cuda13
+    if: github.event_name == 'release'
+    needs: [publish_ri_kernels]
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi-ri_kernels_cuda13
+      url: https://pypi.org/p/ri_kernels_cuda13
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          pattern: wheels-ri_kernels_cuda13-*
+          path: dist
+          merge-multiple: true
+      - uses: pypa/gh-action-pypi-publish@release/v1

ri_kernels-0.1.0/CMakeLists.txt

@@ -0,0 +1,203 @@
+cmake_minimum_required(VERSION 3.20 FATAL_ERROR)
+
+# Parse version from pyproject.toml
+file(READ "${CMAKE_CURRENT_SOURCE_DIR}/pyproject.toml" PYPROJECT_TOML)
+string(REGEX MATCH "\nversion = \"([^\"]+)\"" _ "${PYPROJECT_TOML}")
+set(RI_KERNELS_VERSION "${CMAKE_MATCH_1}")
+
+project(ri_kernels LANGUAGES CXX VERSION "${RI_KERNELS_VERSION}")
+set(RI_KERNELS_SO_VERSION ${CMAKE_PROJECT_VERSION_MAJOR})
+
+
+# set default build type to RELEASE
+if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
+    set(CMAKE_BUILD_TYPE "Release" CACHE STRING "Build type" FORCE)
+	set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS
+		"Debug" "Release" "MinSizeRel" "RelWithDebInfo"
+	)
+endif()
+
+# set language and standard
+set(CMAKE_CXX_STANDARD 20)
+set(CUDA_STANDARD 20)
+set(CUDA_STANDARD_REQUIRED ON)
+
+include(CMakeDependentOption)
+include(FetchContent)
+include(CheckLinkerFlag)
+
+# --exclude-libs is a GNU ld / lld feature; Apple's linker rejects it.
+check_linker_flag(CXX "-Wl,--exclude-libs,ALL" RI_KERNELS_HAVE_EXCLUDE_LIBS)
+
+# Export only the FFI entry points annotated with RI_KERNELS_API (see
+# src/visibility.h). Everything else, including symbols pulled in from static
+# archives such as Highway or a statically linked CUDA runtime, stays out of
+# the dynamic symbol table. Without this the CUDA runtime symbols would be
+# exported and could be interposed by the runtime XLA already loaded into the
+# process, since the dynamic loader searches the global scope before a dlopen'd
+# object's own definitions.
+function(ri_kernels_restrict_exports target)
+	set_target_properties(${target} PROPERTIES
+		C_VISIBILITY_PRESET hidden
+		CXX_VISIBILITY_PRESET hidden
+		CUDA_VISIBILITY_PRESET hidden
+		HIP_VISIBILITY_PRESET hidden
+		VISIBILITY_INLINES_HIDDEN ON
+	)
+	if(RI_KERNELS_HAVE_EXCLUDE_LIBS)
+		target_link_options(${target} PRIVATE "-Wl,--exclude-libs,ALL")
+	endif()
+endfunction()
+
+if (CMAKE_VERSION VERSION_GREATER_EQUAL "3.24.0")
+	cmake_policy(SET CMP0135 NEW)
+endif()
+
+# Options
+option(RI_KERNELS_CPU "Build CPU extension" ON)
+option(RI_KERNELS_CUDA "Build CUDA extension" OFF)
+option(RI_KERNELS_ROCM "Build ROCm extension" OFF)
+
+option(RI_KERNELS_BUNDLED_LIBS "Use all bundled libraries" ON)
+cmake_dependent_option(RI_KERNELS_BUNDLED_HIGHWAY "Use bundled highway lib" ON "RI_KERNELS_BUNDLED_LIBS" OFF)
+option(RI_KERNELS_MULTI_ARCH "Build kernels for multiple CPU architectues with dynamic dispatch. When disabled, arch flags should be set through CMAKE_CXX_FLAGS." ON)
+
+# Directory inside the wheel the GPU libraries are installed into. The CPU
+# library always ships in the ri_kernels package, but the GPU libraries are
+# published as separate add-on distributions (ri_kernels_cuda12,
+# ri_kernels_cuda13) that own a top-level package directory of their own.
+set(RI_KERNELS_GPU_INSTALL_DIR "ri_kernels" CACHE STRING
+	"Directory inside the wheel that the GPU libraries are installed into")
+
+set(CMAKE_CUDA_ARCHITECTURES "75-real;80-real;90-real;90-virtual" CACHE STRING "CUDA Architectures")
+set(CMAKE_HIP_ARCHITECTURES "gfx90a;gfx942;gfx1030;gfx1100" CACHE STRING "HIP Architectures")
+
+set(Python_EXECUTABLE "python3" CACHE STRING "The python interpreter")
+
+# Detect jaxlib include directory (override with -DJAXLIB_HOME=...)
+if(NOT DEFINED JAXLIB_HOME)
+	execute_process(
+		COMMAND ${Python_EXECUTABLE} -c
+			"import jaxlib, os; print(os.path.dirname(jaxlib.__file__))"
+		OUTPUT_VARIABLE JAXLIB_HOME
+		OUTPUT_STRIP_TRAILING_WHITESPACE
+		RESULT_VARIABLE _jaxlib_result
+	)
+	if(NOT _jaxlib_result EQUAL 0)
+		message(FATAL_ERROR "Failed to locate jaxlib via python3. "
+			"Install jaxlib or pass -DJAXLIB_HOME=<path>.")
+	endif()
+endif()
+message(STATUS "Using JAXLIB_HOME: ${JAXLIB_HOME}")
+
+
+set(RI_KERNELS_CPU_SOURCES
+  ./src/rfi_kernel.cpp
+  ./src/rfi_jvp_kernel.cpp
+  ./src/rfi_transpose_kernel.cpp
+)
+
+set(RI_KERNELS_GPU_SOURCES
+  ./src/rfi_kernel_gpu.cu
+  ./src/rfi_jvp_kernel_gpu.cu
+  ./src/rfi_transpose_kernel_gpu.cu
+  ./src/util_gpu.cu
+)
+
+if(RI_KERNELS_CPU)
+  add_library(ri_kernels SHARED ${RI_KERNELS_CPU_SOURCES})
+  target_include_directories(ri_kernels PRIVATE ${JAXLIB_HOME}/include ./src)
+
+  target_compile_options(ri_kernels PRIVATE
+	$<$<COMPILE_LANG_AND_ID:CXX,GNU>:-Wno-return-type>
+	$<$<COMPILE_LANG_AND_ID:CXX,GNU>:-Wno-attributes>
+  )
+
+  if(RI_KERNELS_BUNDLED_HIGHWAY)
+	# add google highway
+	set(HWY_ENABLE_CONTRIB ON CACHE BOOL "")
+	set(HWY_ENABLE_EXAMPLES OFF CACHE BOOL "")
+	set(HWY_ENABLE_INSTALL OFF CACHE BOOL "")
+	set(HWY_ENABLE_TESTS OFF CACHE BOOL "")
+	set(HWY_FORCE_STATIC_LIBS ON CACHE BOOL "")
+	set(HWY_ENABLE_CONTRIB ON CACHE BOOL "")
+	FetchContent_Declare(
+	  hwy
+	  URL https://github.com/google/highway/archive/refs/tags/1.4.0.tar.gz
+	  URL_MD5 9d335797777e17f827c7980b8313a34b
+	)
+	FetchContent_MakeAvailable(hwy)
+	if(NOT TARGET hwy::hwy)
+	  add_library(hwy::hwy ALIAS hwy)
+	endif()
+  else()
+	find_package(hwy CONFIG REQUIRED)
+  endif()
+  target_link_libraries(ri_kernels PRIVATE hwy::hwy)
+
+  if(RI_KERNELS_MULTI_ARCH)
+	target_compile_definitions(ri_kernels PUBLIC -DRI_KERNELS_MULTI_ARCH)
+  endif()
+endif()
+
+if(RI_KERNELS_CUDA)
+  enable_language(CUDA)
+  # find toolkit after language is enabled to ensure version matching
+  find_package(CUDAToolkit REQUIRED)
+
+  set_source_files_properties(${RI_KERNELS_GPU_SOURCES} PROPERTIES LANGUAGE CUDA)
+
+  add_library(ri_kernels_cuda SHARED ${RI_KERNELS_GPU_SOURCES})
+  # set_target_properties(ri_kernels_cuda PROPERTIES OUTPUT_NAME "ri_kernels_cuda_${CUDAToolkit_VERSION_MAJOR}")
+  target_include_directories(ri_kernels_cuda PRIVATE ${JAXLIB_HOME}/include ./src)
+  target_link_libraries(ri_kernels_cuda PRIVATE CUDA::cudart_static)
+  target_compile_options(ri_kernels_cuda PRIVATE
+	$<$<COMPILE_LANGUAGE:CUDA>:--diag-suppress=940>
+	$<$<COMPILE_LANGUAGE:CUDA>:--diag-suppress=2473>
+  )
+endif()
+
+if(RI_KERNELS_ROCM)
+  enable_language(HIP)
+  find_package(hip CONFIG REQUIRED)
+
+  set_source_files_properties(${RI_KERNELS_GPU_SOURCES} PROPERTIES LANGUAGE HIP)
+
+  add_library(ri_kernels_hip SHARED ${RI_KERNELS_GPU_SOURCES})
+  target_include_directories(ri_kernels_hip PRIVATE ${JAXLIB_HOME}/include ./src)
+  target_link_libraries(ri_kernels_hip PRIVATE hip::host)
+endif()
+
+# Normalize shared-library suffix across platforms so Python loader finds
+# libri_kernels*.so on macOS as well as Linux, and restrict the exported
+# symbols to the annotated FFI entry points.
+if(TARGET ri_kernels)
+  set_target_properties(ri_kernels PROPERTIES SUFFIX ".so")
+  ri_kernels_restrict_exports(ri_kernels)
+endif()
+if(TARGET ri_kernels_cuda)
+  set_target_properties(ri_kernels_cuda PROPERTIES SUFFIX ".so")
+  ri_kernels_restrict_exports(ri_kernels_cuda)
+endif()
+if(TARGET ri_kernels_hip)
+  set_target_properties(ri_kernels_hip PROPERTIES SUFFIX ".so")
+  ri_kernels_restrict_exports(ri_kernels_hip)
+endif()
+
+# Install rules so scikit-build-core stages libraries into the wheel at the
+# location where ri_kernels/rfi_vis_op.py loads them from.
+if(TARGET ri_kernels)
+  install(TARGETS ri_kernels
+	LIBRARY DESTINATION ri_kernels
+	RUNTIME DESTINATION ri_kernels)
+endif()
+if(TARGET ri_kernels_cuda)
+  install(TARGETS ri_kernels_cuda
+	LIBRARY DESTINATION ${RI_KERNELS_GPU_INSTALL_DIR}
+	RUNTIME DESTINATION ${RI_KERNELS_GPU_INSTALL_DIR})
+endif()
+if(TARGET ri_kernels_hip)
+  install(TARGETS ri_kernels_hip
+	LIBRARY DESTINATION ${RI_KERNELS_GPU_INSTALL_DIR}
+	RUNTIME DESTINATION ${RI_KERNELS_GPU_INSTALL_DIR})
+endif()