[OrcJIT] Add LLVM ORC JIT v2 dynamic object loader addon

.github/workflows/tvm-ffi-orcjit.yml

@@ -0,0 +1,118 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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.
+
+name: OrcJIT Build
+
+on:
+  push:
+    paths:
+      - "addons/tvm-ffi-orcjit/**"
+      - ".github/workflows/tvm-ffi-orcjit.yml"
+  pull_request:
+    paths:
+      - "addons/tvm-ffi-orcjit/**"
+      - ".github/workflows/tvm-ffi-orcjit.yml"
+  workflow_dispatch:
+
+jobs:
+  build_wheels:
+    name: ${{ matrix.os }} (${{ matrix.arch }})
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - {os: ubuntu-latest, arch: x86_64, build: "cp312-manylinux_x86_64"}
+          - {os: ubuntu-24.04-arm, arch: aarch64, build: "cp312-manylinux_aarch64"}
+          - {os: macos-14, arch: arm64, build: "cp312-macosx_arm64"}
+          - {os: windows-latest, arch: AMD64, build: "cp312-win_amd64"}
+
+    steps:
+      - uses: actions/checkout@v5
+        with:
+          submodules: recursive
+
+      - uses: ./.github/actions/detect-env-vars
+        id: env_vars
+
+      # ---- Cache LLVM prefix ----
+      - name: Cache LLVM
+        uses: actions/cache@v4
+        id: llvm-cache
+        with:
+          path: ${{ runner.os == 'Windows' && 'C:/opt/llvm' || '/opt/llvm' }}
+          key: llvm-22.1.0-${{ runner.os }}-${{ matrix.arch }}-v3
+
+      # ---- Install LLVM via conda (cache miss only) ----
+      - name: Setup conda
+        if: steps.llvm-cache.outputs.cache-hit != 'true'
+        uses: conda-incubator/setup-miniconda@fc2d68f6413eb2d87b895e92f8584b5b94a10167  # v3.3.0
+        continue-on-error: true
+        id: conda1
+        with:
+          miniforge-version: latest
+
+      - name: Setup conda (retry with tar.bz2)
+        if: steps.llvm-cache.outputs.cache-hit != 'true' && steps.conda1.outcome == 'failure'
+        uses: conda-incubator/setup-miniconda@fc2d68f6413eb2d87b895e92f8584b5b94a10167  # v3.3.0
+        with:
+          miniforge-version: latest
+          use-only-tar-bz2: true
+
+      - name: Create /opt/llvm (macOS)
+        if: steps.llvm-cache.outputs.cache-hit != 'true' && runner.os == 'macOS'
+        run: sudo mkdir -p /opt/llvm && sudo chown -R $(whoami) /opt/llvm
+
+      - name: Install LLVM (Unix)
+        if: steps.llvm-cache.outputs.cache-hit != 'true' && runner.os != 'Windows'
+        shell: bash -l {0}
+        run: |
+          conda create -q -p /opt/llvm -c conda-forge \
+            llvmdev=22.1.0 clangdev=22.1.0 compiler-rt=22.1.0 zlib zstd-static \
+            -y
+
+      - name: Install LLVM (Windows)
+        if: steps.llvm-cache.outputs.cache-hit != 'true' && runner.os == 'Windows'
+        shell: cmd /C call {0}
+        run: |
+          conda create -q -p C:\opt\llvm -c conda-forge llvmdev=22.1.0 zlib zstd-static -y
+
+      # ---- Build and test wheels ----
+      - name: Build and test wheels
+        uses: pypa/cibuildwheel@298ed2fb2c105540f5ed055e8a6ad78d82dd3a7e  # v3.3.1
+        with:
+          package-dir: addons/tvm-ffi-orcjit
+          output-dir: wheelhouse
+        env:
+          CIBW_BUILD: ${{ matrix.build }}
+          CIBW_ARCHS_LINUX: ${{ matrix.arch }}
+          CIBW_ARCHS_MACOS: ${{ matrix.arch }}
+          CIBW_ARCHS_WINDOWS: ${{ matrix.arch }}
+          CIBW_BUILD_VERBOSITY: 1
+          CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }}
+          CIBW_ENVIRONMENT: LLVM_PREFIX=/opt/llvm
+          CIBW_ENVIRONMENT_WINDOWS: LLVM_PREFIX="C:/opt/llvm"
+          CIBW_CONTAINER_ENGINE: "docker; create_args: --volume /opt/llvm:/opt/llvm"
+          CIBW_TEST_REQUIRES: pytest ninja
+          CIBW_TEST_COMMAND: >-
+            pip install --force-reinstall {project} &&
+            python {project}/addons/tvm-ffi-orcjit/tests/run_all_tests.py
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: cibw-wheels-${{ matrix.os }}-${{ matrix.arch }}-${{ strategy.job-index }}
+          path: wheelhouse/*.whl

addons/tvm-ffi-orcjit/CMakeLists.txt

@@ -0,0 +1,178 @@
+cmake_minimum_required(VERSION 3.20)
+project(
+  tvm_ffi_orcjit
+  VERSION 0.1.0
+  LANGUAGES C CXX
+)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+
+# ---- LLVM_PREFIX handling ----
+# scikit-build-core overrides CMAKE_PREFIX_PATH via its init cache, so we read LLVM_PREFIX from the
+# environment instead.
+if (DEFINED ENV{LLVM_PREFIX})
+  list(APPEND CMAKE_PREFIX_PATH "$ENV{LLVM_PREFIX}")
+  if (EXISTS "$ENV{LLVM_PREFIX}/Library")
+    list(APPEND CMAKE_PREFIX_PATH "$ENV{LLVM_PREFIX}/Library")
+  endif ()
+endif ()
+
+# ---- Find packages ----
+find_package(LLVM REQUIRED CONFIG)
+message(STATUS "Found LLVM ${LLVM_PACKAGE_VERSION} in ${LLVM_DIR}")
+
+find_package(
+  Python
+  COMPONENTS Interpreter
+  REQUIRED
+)
+execute_process(
+  COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+  OUTPUT_VARIABLE tvm_ffi_ROOT
+)
+find_package(tvm_ffi CONFIG REQUIRED)
+
+# ---- Build shared library ----
+add_library(tvm_ffi_orcjit SHARED src/ffi/orcjit_session.cc src/ffi/orcjit_dylib.cc)
+set_target_properties(
+  tvm_ffi_orcjit PROPERTIES CXX_VISIBILITY_PRESET hidden VISIBILITY_INLINES_HIDDEN ON
+)
+target_include_directories(
+  tvm_ffi_orcjit PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src ${LLVM_INCLUDE_DIRS}
+)
+separate_arguments(LLVM_DEFINITIONS_LIST NATIVE_COMMAND ${LLVM_DEFINITIONS})
+target_compile_definitions(tvm_ffi_orcjit PRIVATE ${LLVM_DEFINITIONS_LIST})
+
+# ---- Static LLVM linking via llvm-config ----
+set(_llvm_config_hints ${LLVM_TOOLS_BINARY_DIR} "${LLVM_DIR}/../../../bin")
+if (DEFINED ENV{LLVM_PREFIX})
+  list(APPEND _llvm_config_hints "$ENV{LLVM_PREFIX}/bin" "$ENV{LLVM_PREFIX}/Library/bin")
+endif ()
+find_program(
+  LLVM_CONFIG_EXE
+  NAMES llvm-config-${LLVM_VERSION_MAJOR} llvm-config
+  HINTS ${_llvm_config_hints}
+)
+if (NOT LLVM_CONFIG_EXE)
+  message(FATAL_ERROR "llvm-config not found")
+endif ()
+
+execute_process(
+  COMMAND ${LLVM_CONFIG_EXE} --link-static --libs Core OrcJIT Support native
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+  OUTPUT_VARIABLE _llvm_libs
+)
+execute_process(
+  COMMAND ${LLVM_CONFIG_EXE} --link-static --ldflags
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+  OUTPUT_VARIABLE _llvm_ldflags
+)
+separate_arguments(_llvm_libs_list NATIVE_COMMAND "${_llvm_libs}")
+separate_arguments(_llvm_ldflags_list NATIVE_COMMAND "${_llvm_ldflags}")
+
+# ---- zlib (static, from LLVM prefix) ----
+cmake_path(GET LLVM_DIR PARENT_PATH _llvm_prefix)
+cmake_path(GET _llvm_prefix PARENT_PATH _llvm_prefix)
+cmake_path(GET _llvm_prefix PARENT_PATH _llvm_prefix)
+set(_lib_hints "${_llvm_prefix}/lib" "${_llvm_prefix}/lib64" "${_llvm_prefix}/Library/lib")
+
+if (WIN32)
+  find_library(
+    ZLIB_STATIC
+    NAMES zlibstatic z
+    HINTS ${_lib_hints} REQUIRED
+  )
+else ()
+  find_library(
+    ZLIB_STATIC
+    NAMES libz.a
+    HINTS ${_lib_hints} REQUIRED
+  )
+endif ()
+
+# ---- zstd (static, from LLVM prefix) ----
+if (WIN32)
+  find_library(
+    ZSTD_STATIC
+    NAMES zstd_static zstd
+    HINTS ${_lib_hints} REQUIRED
+  )
+else ()
+  find_library(
+    ZSTD_STATIC
+    NAMES libzstd.a
+    HINTS ${_lib_hints} REQUIRED
+  )
+endif ()
+
+# ---- Link everything ----
+target_link_libraries(
+  tvm_ffi_orcjit PRIVATE tvm_ffi::header tvm_ffi::shared ${_llvm_ldflags_list} ${_llvm_libs_list}
+                         ${ZLIB_STATIC} ${ZSTD_STATIC}
+)
+
+# LLVM system libs (ntdll/psapi on Windows, rt/dl on Linux), minus zlib/zstd (linked above).
+execute_process(
+  COMMAND ${LLVM_CONFIG_EXE} --system-libs
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+  OUTPUT_VARIABLE _sys_libs
+  RESULT_VARIABLE _sys_libs_rc
+)
+if (_sys_libs_rc EQUAL 0 AND _sys_libs)
+  separate_arguments(_sys_libs_list NATIVE_COMMAND "${_sys_libs}")
+  list(FILTER _sys_libs_list EXCLUDE REGEX "(zstd|^-lz$|^z\\.lib$|xml2)")
+  target_link_libraries(tvm_ffi_orcjit PRIVATE ${_sys_libs_list})
+endif ()
+
+# ---- Platform-specific fixups ----
+if (APPLE)
+  add_custom_command(
+    TARGET tvm_ffi_orcjit
+    POST_BUILD
+    COMMAND install_name_tool -change @rpath/libc++.1.dylib /usr/lib/libc++.1.dylib
+            $<TARGET_FILE:tvm_ffi_orcjit>
+    COMMENT "Fixing libc++ rpath to use system library"
+  )
+elseif (UNIX AND NOT WIN32)
+  target_link_options(tvm_ffi_orcjit PRIVATE -static-libstdc++ -static-libgcc)
+endif ()
+
+# ---- Find and bundle liborc_rt ----
+set(ORC_RT_PATH "")
+if (DEFINED ENV{ORC_RT_PATH})
+  set(ORC_RT_PATH "$ENV{ORC_RT_PATH}")
+else ()
+  if (WIN32)
+    file(GLOB_RECURSE _orc_rt_candidates "${_llvm_prefix}/lib/clang/*/lib/orc_rt*.lib")
+    if (NOT _orc_rt_candidates)
+      file(GLOB_RECURSE _orc_rt_candidates "${_llvm_prefix}/Library/lib/clang/*/lib/orc_rt*.lib")
+    endif ()
+  else ()
+    file(GLOB_RECURSE _orc_rt_candidates "${_llvm_prefix}/lib/clang/*/lib/liborc_rt*.a")
+  endif ()
+  if (_orc_rt_candidates)
+    list(GET _orc_rt_candidates 0 ORC_RT_PATH)
+  endif ()
+endif ()
+
+if (NOT ORC_RT_PATH)
+  message(WARNING "Could not find liborc_rt. ORC runtime features will be disabled.")
+endif ()
+
+if (ORC_RT_PATH)
+  message(STATUS "Found liborc_rt: ${ORC_RT_PATH}")
+  file(COPY "${ORC_RT_PATH}" DESTINATION "${CMAKE_BINARY_DIR}")
+  install(FILES "${ORC_RT_PATH}" DESTINATION lib)
+endif ()
+
+# ---- Install ----
+install(
+  TARGETS tvm_ffi_orcjit
+  LIBRARY DESTINATION lib
+  ARCHIVE DESTINATION lib
+  RUNTIME DESTINATION lib
+)