From 24ffd33beb61968e28dc694a664c799238694e50 Mon Sep 17 00:00:00 2001
From: Hyunsu Cho <chohyu01@cs.washington.edu>
Date: Tue, 1 Feb 2022 17:18:26 -0800
Subject: [PATCH 01/18] Add back OpenMP

---
 CMakeLists.txt                     |   1 +
 cmake/TreeliteConfig.cmake.in      |   5 +-
 src/CMakeLists.txt                 |  17 ++-
 src/annotator.cc                   |   8 +-
 src/threading_utils/parallel_for.h | 159 ++++++++++++++++++++---------
 tests/cpp/test_threading_utils.cc  |  58 ++---------
 6 files changed, 144 insertions(+), 104 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 635aebd9..dec67470 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -12,6 +12,7 @@ if(MSVC)
 endif()
 
 option(TEST_COVERAGE "C++ test coverage" OFF)
+option(USE_OPENMP "Use OpenMP" ON)
 option(BUILD_CPP_TEST "Build C++ tests" OFF)
 option(BUILD_STATIC_LIBS "Build static libs, in addition to dynamic libs" OFF)
 option(DETECT_CONDA_ENV "Enable detection of conda environment for dependencies" ON)
diff --git a/cmake/TreeliteConfig.cmake.in b/cmake/TreeliteConfig.cmake.in
index 18a875c4..587ebd71 100644
--- a/cmake/TreeliteConfig.cmake.in
+++ b/cmake/TreeliteConfig.cmake.in
@@ -2,7 +2,10 @@
 
 include(CMakeFindDependencyMacro)
 
-find_dependency(Threads)
+set(USE_OPENMP @USE_OPENMP@)
+if(USE_OPENMP)
+  find_dependency(OpenMP)
+endif()
 
 if(NOT TARGET treelite::treelite)
   include(${CMAKE_CURRENT_LIST_DIR}/TreeliteTargets.cmake)
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index 7fb97183..b75ad944 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -11,7 +11,16 @@ endif(UNIX)
 
 add_library(objtreelite_common OBJECT)  # Component shared by both main package and runtime
 
-find_package(Threads REQUIRED)
+if(USE_OPENMP)
+  if (APPLE)
+    # Require CMake 3.16+ on Mac OSX, as previous versions of CMake had trouble locating
+    # OpenMP on Mac. See https://github.com/dmlc/xgboost/pull/5146#issuecomment-568312706
+    cmake_minimum_required(VERSION 3.16)
+  endif (APPLE)
+  find_package(OpenMP REQUIRED)
+else()
+  message(STATUS "Disabling OpenMP")
+endif()
 
 if(ENABLE_ALL_WARNINGS)
   foreach(target objtreelite objtreelite_runtime objtreelite_runtime)
@@ -24,7 +33,6 @@ foreach(lib objtreelite objtreelite_runtime objtreelite_common)
     $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}/include>
     $<BUILD_INTERFACE:${PROJECT_BINARY_DIR}/include>
     $<INSTALL_INTERFACE:$<INSTALL_PREFIX>/include>)
-  target_link_libraries(${lib} PUBLIC Threads::Threads)
   if(MSVC)
     target_compile_options(${lib} PRIVATE /MP)
     target_compile_definitions(${lib} PRIVATE -DNOMINMAX)
@@ -32,6 +40,10 @@ foreach(lib objtreelite objtreelite_runtime objtreelite_common)
   else()
     target_compile_options(${lib} PRIVATE -funroll-loops)
   endif()
+  if(USE_OPENMP)
+    target_link_libraries(${lib} PUBLIC OpenMP::OpenMP_CXX)
+    target_compile_definitions(${lib} PRIVATE -DTREELITE_OPENMP_SUPPORT)
+  endif()
   if(TEST_COVERAGE)
     if(MSVC)
       message(FATAL_ERROR "Test coverage not available on Windows")
@@ -101,6 +113,7 @@ target_sources(objtreelite
     ${PROJECT_SOURCE_DIR}/include/treelite/frontend.h
     ${PROJECT_SOURCE_DIR}/include/treelite/frontend_impl.h
     ${PROJECT_SOURCE_DIR}/include/treelite/gtil.h
+    ${PROJECT_SOURCE_DIR}/include/treelite/omp.h
     ${PROJECT_SOURCE_DIR}/include/treelite/optional.h
     ${PROJECT_SOURCE_DIR}/include/treelite/thread_local.h
     ${PROJECT_SOURCE_DIR}/include/treelite/tree.h
diff --git a/src/annotator.cc b/src/annotator.cc
index 33eb3b21..7de197e6 100644
--- a/src/annotator.cc
+++ b/src/annotator.cc
@@ -77,7 +77,8 @@ inline void ComputeBranchLoopImpl(
   size_t num_col = dmat->num_col;
   ElementType missing_value = dmat->missing_value;
   bool nan_missing = treelite::math::CheckNAN(missing_value);
-  treelite::threading_utils::ParallelFor(rbegin, rend, nthread,
+  auto sched = treelite::threading_utils::ParallelSchedule::Static();
+  treelite::threading_utils::ParallelFor(rbegin, rend, nthread, sched,
                                          [&](std::size_t rid, std::size_t thread_id) {
     const ElementType* row = &dmat->data[rid * num_col];
     const size_t off = dmat->num_col * thread_id;
@@ -107,7 +108,8 @@ inline void ComputeBranchLoopImpl(
   std::vector<Entry<ElementType>> inst(nthread * dmat->num_col, {-1});
   size_t ntree = model.trees.size();
   TREELITE_CHECK_LE(rbegin, rend);
-  treelite::threading_utils::ParallelFor(rbegin, rend, nthread,
+  auto sched = treelite::threading_utils::ParallelSchedule::Static();
+  treelite::threading_utils::ParallelFor(rbegin, rend, nthread, sched,
                                          [&](std::size_t rid, std::size_t thread_id) {
     const size_t off = dmat->num_col * thread_id;
     const size_t off2 = count_row_ptr[ntree] * thread_id;
@@ -192,7 +194,7 @@ AnnotateImpl(
 
   count_row_ptr = {0};
   const size_t ntree = model.trees.size();
-  const int max_thread = static_cast<int>(std::thread::hardware_concurrency());
+  const int max_thread = static_cast<int>(threading_utils::MaxNumThread());
   nthread = (nthread == 0) ? max_thread : std::min(nthread, max_thread);
   for (const treelite::Tree<ThresholdType, LeafOutputType>& tree : model.trees) {
     count_row_ptr.push_back(count_row_ptr.back() + tree.num_nodes);
diff --git a/src/threading_utils/parallel_for.h b/src/threading_utils/parallel_for.h
index c966b779..33f7a7bf 100644
--- a/src/threading_utils/parallel_for.h
+++ b/src/threading_utils/parallel_for.h
@@ -7,75 +7,134 @@
 #ifndef TREELITE_THREADING_UTILS_PARALLEL_FOR_H_
 #define TREELITE_THREADING_UTILS_PARALLEL_FOR_H_
 
+#include <treelite/omp.h>
 #include <treelite/logging.h>
-#include <future>
-#include <thread>
-#include <algorithm>
-#include <vector>
+#include <type_traits>
+#include <exception>
+#include <mutex>
 #include <cstddef>
 
 namespace treelite {
 namespace threading_utils {
 
-template <typename IndexType>
-std::vector<IndexType> ComputeWorkRange(IndexType begin, IndexType end, std::size_t nthread);
+/*!
+ * \brief OMP Exception class catches, saves and rethrows exception from OMP blocks
+ */
+class OMPException {
+ private:
+  // exception_ptr member to store the exception
+  std::exception_ptr omp_exception_;
+  // mutex to be acquired during catch to set the exception_ptr
+  std::mutex mutex_;
+
+ public:
+  /*!
+   * \brief Parallel OMP blocks should be placed within Run to save exception
+   */
+  template <typename Function, typename... Parameters>
+  void Run(Function f, Parameters... params) {
+    try {
+      f(params...);
+    } catch (std::exception& ex) {
+      std::lock_guard<std::mutex> lock(mutex_);
+      if (!omp_exception_) {
+        omp_exception_ = std::current_exception();
+      }
+    }
+  }
+
+  /*!
+   * \brief should be called from the main thread to rethrow the exception
+   */
+  void Rethrow() {
+    if (this->omp_exception_) {
+      std::rethrow_exception(this->omp_exception_);
+    }
+  }
+};
+
+inline int MaxNumThread() {
+  return omp_get_max_threads();
+}
+
+// OpenMP schedule
+struct ParallelSchedule {
+  enum {
+    kAuto,
+    kDynamic,
+    kStatic,
+    kGuided,
+  } sched;
+  std::size_t chunk{0};
+
+  ParallelSchedule static Auto() { return ParallelSchedule{kAuto}; }
+  ParallelSchedule static Dynamic(std::size_t n = 0) { return ParallelSchedule{kDynamic, n}; }
+  ParallelSchedule static Static(std::size_t n = 0) { return ParallelSchedule{kStatic, n}; }
+  ParallelSchedule static Guided() { return ParallelSchedule{kGuided}; }
+};
 
 template <typename IndexType, typename FuncType>
-void ParallelFor(IndexType begin, IndexType end, std::size_t nthread, FuncType func) {
+inline void ParallelFor(IndexType begin, IndexType end, int nthread, ParallelSchedule sched,
+                        FuncType func) {
   TREELITE_CHECK_GT(nthread, 0) << "nthread must be positive";
-  TREELITE_CHECK_LE(nthread, std::thread::hardware_concurrency())
-    << "nthread cannot exceed " << std::thread::hardware_concurrency();
+  TREELITE_CHECK_LE(nthread, MaxNumThread()) << "nthread cannot exceed " << MaxNumThread();
   if (begin == end) {
     return;
   }
-  /* Divide the range [begin, end) equally among the threads.
-   * The i-th thread gets the range [work_range[i], work_range[i+1]). */
-  std::vector<IndexType> work_range = ComputeWorkRange(begin, end, nthread);
 
-  // Launch (nthread - 1) threads, as the main thread should also perform work.
-  std::vector<std::future<void>> async_tasks;
-  for (std::size_t thread_id = 1; thread_id < nthread; ++thread_id) {
-    async_tasks.push_back(std::async(std::launch::async, [&work_range, &func, thread_id]() {
-      const IndexType begin_ = work_range[thread_id];
-      const IndexType end_ = work_range[thread_id + 1];
-      for (IndexType i = begin_; i < end_; ++i) {
-        func(i, thread_id);
+#if defined(_MSC_VER)
+  // msvc doesn't support unsigned integer as openmp index.
+  using OmpInd = std::conditional_t<std::is_signed<IndexType>::value, IndexType, std::int64_t>;
+#else
+  using OmpInd = IndexType;
+#endif
+
+  OMPException exc;
+  switch (sched.sched) {
+  case ParallelSchedule::kAuto: {
+#pragma omp parallel for num_threads(nthread)
+    for (OmpInd i = begin; i < end; ++i) {
+      exc.Run(func, i, omp_get_thread_num());
+    }
+    break;
+  }
+  case ParallelSchedule::kDynamic: {
+    if (sched.chunk == 0) {
+#pragma omp parallel for num_threads(nthread) schedule(dynamic)
+      for (OmpInd i = begin; i < end; ++i) {
+        exc.Run(func, i, omp_get_thread_num());
+      }
+    } else {
+#pragma omp parallel for num_threads(nthread) schedule(dynamic, sched.chunk)
+      for (OmpInd i = begin; i < end; ++i) {
+        exc.Run(func, i, omp_get_thread_num());
       }
-    }));
+    }
+    break;
   }
-  {
-    const IndexType begin_ = work_range[0];
-    const IndexType end_ = work_range[1];
-    for (IndexType i = begin_; i < end_; ++i) {
-      func(i, 0);
+  case ParallelSchedule::kStatic: {
+    if (sched.chunk == 0) {
+#pragma omp parallel for num_threads(nthread) schedule(static)
+      for (OmpInd i = begin; i < end; ++i) {
+        exc.Run(func, i, omp_get_thread_num());
+      }
+    } else {
+#pragma omp parallel for num_threads(nthread) schedule(static, sched.chunk)
+      for (OmpInd i = begin; i < end; ++i) {
+        exc.Run(func, i, omp_get_thread_num());
+      }
     }
+    break;
   }
-  // Join threads
-  for (auto& task : async_tasks) {
-    task.get();
+  case ParallelSchedule::kGuided: {
+#pragma omp parallel for num_threads(nthread) schedule(guided)
+    for (OmpInd i = begin; i < end; ++i) {
+      exc.Run(func, i, omp_get_thread_num());
+    }
+    break;
   }
-}
-
-template <typename IndexType>
-std::vector<IndexType> ComputeWorkRange(IndexType begin, IndexType end, std::size_t nthread) {
-  TREELITE_CHECK_GE(end, 0) << "end must be 0 or greater";
-  TREELITE_CHECK_GE(begin, 0) << "begin must be 0 or greater";
-  TREELITE_CHECK_GE(end, begin) << "end cannot be less than begin";
-  TREELITE_CHECK_GT(nthread, 0) << "nthread must be positive";
-  IndexType num_elem = end - begin;
-  const IndexType portion = num_elem / nthread + !!(num_elem % nthread);
-  // integer division, rounded-up
-
-  std::vector<IndexType> work_range(nthread + 1);
-  work_range[0] = begin;
-  IndexType acc = begin;
-  for (std::size_t i = 0; i < nthread; ++i) {
-    acc += portion;
-    work_range[i + 1] = std::min(acc, end);
   }
-  TREELITE_CHECK_EQ(work_range[nthread], end);
-
-  return work_range;
+  exc.Rethrow();
 }
 
 }  // namespace threading_utils
diff --git a/tests/cpp/test_threading_utils.cc b/tests/cpp/test_threading_utils.cc
index 5adf4b3b..84b8f371 100644
--- a/tests/cpp/test_threading_utils.cc
+++ b/tests/cpp/test_threading_utils.cc
@@ -8,7 +8,6 @@
 #include <gmock/gmock.h>
 #include <vector>
 #include <algorithm>
-#include <thread>
 #include <random>
 #include <cstddef>
 #include <cstdint>
@@ -53,56 +52,19 @@ class RandomGenerator {
 
 namespace treelite {
 namespace threading_utils {
-
-TEST(ThreadingUtils, ComputeWorkRange) {
-  /* Test error handling */
-  EXPECT_THROW(ComputeWorkRange(0, 100, 0), treelite::Error);
-  EXPECT_THROW(ComputeWorkRange(-100, 100, 3), treelite::Error);
-  EXPECT_THROW(ComputeWorkRange(-200, -100, 3), treelite::Error);
-  EXPECT_THROW(ComputeWorkRange(200, 100, 3), treelite::Error);
-
-  /* Property-based testing with randomly generated parameters */
-  RandomGenerator rng;
-
-  constexpr int kNumTrial = 200;
-  for (int i = 0; i < kNumTrial; ++i) {
-    int64_t begin = rng.DrawInteger(0, 10000);
-    std::size_t nthread = static_cast<std::size_t>(rng.DrawInteger(1, 100));
-    int64_t end = rng.DrawInteger(begin, 10000);
-    auto range = ComputeWorkRange(begin, end, nthread);
-    EXPECT_EQ(range.size(), nthread + 1);
-    EXPECT_EQ(range[0], begin);
-    EXPECT_EQ(range[nthread], end);
-    for (std::size_t i = 0; i < nthread; ++i) {
-      EXPECT_GE(range[i + 1], range[i]);
-    }
-  }
-  // Test the case with begin == end
-  for (int i = 0; i < 10; ++i) {
-    int64_t begin = rng.DrawInteger(0, 10000);
-    int64_t end = begin;
-    std::size_t nthread = static_cast<std::size_t>(rng.DrawInteger(1, 100));
-    auto range = ComputeWorkRange(begin, end, nthread);
-    EXPECT_EQ(range.size(), nthread + 1);
-    EXPECT_EQ(range[0], begin);
-    EXPECT_EQ(range[nthread], begin);
-    for (std::size_t i = 0; i < nthread; ++i) {
-      EXPECT_EQ(range[i + 1], range[i]);
-    }
-  }
-}
-
 TEST(ThreadingUtils, ParallelFor) {
   /* Test error handling */
-  const int max_thread = std::thread::hardware_concurrency();
+  const int max_thread = treelite::threading_utils::MaxNumThread();
+
+  auto sched = treelite::threading_utils::ParallelSchedule::Guided();
 
   auto dummy_func = [](int, std::size_t) {};
-  EXPECT_THROW(ParallelFor(0, 100, 0, dummy_func), treelite::Error);
-  EXPECT_THROW(ParallelFor(200, 100, 3, dummy_func), treelite::Error);
-  EXPECT_THROW(ParallelFor(-100, 100, 3, dummy_func), treelite::Error);
-  EXPECT_THROW(ParallelFor(-200, -100, 3, dummy_func), treelite::Error);
-  EXPECT_THROW(ParallelFor(200, 100, 3, dummy_func), treelite::Error);
-  EXPECT_THROW(ParallelFor(10, 20, 3 * max_thread, dummy_func), treelite::Error);
+  EXPECT_THROW(ParallelFor(0, 100, 0, sched, dummy_func), treelite::Error);
+  EXPECT_THROW(ParallelFor(200, 100, 3, sched, dummy_func), treelite::Error);
+  EXPECT_THROW(ParallelFor(-100, 100, 3, sched, dummy_func), treelite::Error);
+  EXPECT_THROW(ParallelFor(-200, -100, 3, sched, dummy_func), treelite::Error);
+  EXPECT_THROW(ParallelFor(200, 100, 3, sched, dummy_func), treelite::Error);
+  EXPECT_THROW(ParallelFor(10, 20, 3 * max_thread, sched, dummy_func), treelite::Error);
 
   /* Property-based testing with randomly generated parameters */
   constexpr int kVectorLength = 10000;
@@ -123,7 +85,7 @@ TEST(ThreadingUtils, ParallelFor) {
     std::size_t nthread = static_cast<std::size_t>(rng.DrawInteger(1, max_thread + 1));
     int64_t end = rng.DrawInteger(begin, kVectorLength);
 
-    ParallelFor(begin, end, nthread, [&a, &b, &c](int64_t i, std::size_t) {
+    ParallelFor(begin, end, nthread, sched, [&a, &b, &c](int64_t i, std::size_t) {
       c[i] = a[i] + b[i];
     });
 

From 8c32348b1bf0fad0fa3bd4c334230af2d0926626 Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 17:35:18 -0800
Subject: [PATCH 02/18] Explicitly specify llvm-openmp dependency

---
 tests/travis/run_test.sh | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/travis/run_test.sh b/tests/travis/run_test.sh
index 4143b25d..d4b108cf 100755
--- a/tests/travis/run_test.sh
+++ b/tests/travis/run_test.sh
@@ -9,7 +9,7 @@ then
   conda activate python3
   conda --version
   python --version
-  conda install -c conda-forge numpy scipy pandas pytest pytest-cov scikit-learn coverage ninja lcov cmake
+  conda install -c conda-forge numpy scipy pandas pytest pytest-cov scikit-learn coverage ninja lcov cmake llvm-openmp
 
   # Run coverage test
   set -x
@@ -38,7 +38,7 @@ then
   conda activate python3
   conda --version
   python --version
-  conda install -c conda-forge ninja cmake rapidjson fmt
+  conda install -c conda-forge ninja cmake rapidjson fmt llvm-openmp
 
   # Install Treelite C++ library into the Conda env
   set -x
@@ -64,7 +64,7 @@ then
   conda activate python3
   conda --version
   python --version
-  conda install -c conda-forge numpy scipy pandas pytest scikit-learn coverage ninja cmake
+  conda install -c conda-forge numpy scipy pandas pytest scikit-learn coverage ninja cmake llvm-openmp
 
   # Build binary wheel
   set -x

From 69cead73b01a3c455ee8aa87c4d4fbe3a37f4a3c Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 18:58:59 -0800
Subject: [PATCH 03/18] [CI] Set up pipeline to build MacOS wheels with libomp
 bundled

---
 .travis.yml                           |  8 +++
 tests/ci_build/build_python_wheels.sh | 72 +++++++++++++++++++++++++++
 tests/travis/run_test.sh              | 29 ++++++-----
 3 files changed, 97 insertions(+), 12 deletions(-)
 create mode 100755 tests/ci_build/build_python_wheels.sh

diff --git a/.travis.yml b/.travis.yml
index 60c8e729..fd370060 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -24,6 +24,14 @@ jobs:
     - os: osx
       env: TASK=python_coverage_test
       osx_image: xcode10.2
+    # Build Python wheels for MacOS Intel and Apple Silicon
+    # pypa/cibuildwheel is used, hence CIBW prefix
+    - os: osx
+      env: TASK=python_wheels CIBW_PLATFORM_ID=macosx_x86_64
+      osx_image: xcode12.5
+    - os: osx
+      env: TASK=python_wheels CIBW_PLATFORM_ID=macosx_arm64
+      osx_image: xcode12.5
 
 # dependent brew packages
 addons:
diff --git a/tests/ci_build/build_python_wheels.sh b/tests/ci_build/build_python_wheels.sh
new file mode 100755
index 00000000..2b53ca20
--- /dev/null
+++ b/tests/ci_build/build_python_wheels.sh
@@ -0,0 +1,72 @@
+#!/bin/bash
+
+set -e
+set -x
+
+if [[ $# -ne 2 ]]; then
+  echo "Usage: $0 [platform_id] [commit ID]"
+  exit 1
+fi
+
+platform_id=$1
+commit_id=$2
+
+# Bundle libomp 11.1.0 when targeting MacOS.
+# This is a workaround in order to prevent segfaults when running inside a Conda environment.
+# See https://github.com/dmlc/xgboost/issues/7039#issuecomment-1025125003 for more context.
+# The workaround is also used by the scikit-learn and XGBoost projects.
+if [[ "$platform_id" == macosx_* ]]; then
+    # Make sure to use a libomp version binary compatible with the oldest
+    # supported version of the macos SDK as libomp will be vendored into the
+    # Treelite wheels for MacOS.
+
+    if [[ "$platform_id" == macosx_arm64 ]]; then
+        # MacOS, Apple Silicon
+        # arm64 builds must cross compile because CI is on x64
+        # cibuildwheel will take care of cross-compilation.
+        wheel_tag=macosx_12_0_arm64
+        cpython_ver=38
+        setup_env_var='CIBW_TARGET_OSX_ARM64=1'  # extra flag to be passed to setup.py
+        export PYTHON_CROSSENV=1
+        export MACOSX_DEPLOYMENT_TARGET=12.0
+        OPENMP_URL="https://anaconda.org/conda-forge/llvm-openmp/11.1.0/download/osx-arm64/llvm-openmp-11.1.0-hf3c4609_1.tar.bz2"
+    elif [[ "$platform_id" == macosx_x86_64 ]]; then
+        # MacOS, Intel
+        wheel_tag=macosx_10_15_x86_64.macosx_11_0_x86_64.macosx_12_0_x86_64
+        cpython_ver=37
+        export MACOSX_DEPLOYMENT_TARGET=10.13
+        OPENMP_URL="https://anaconda.org/conda-forge/llvm-openmp/11.1.0/download/osx-64/llvm-openmp-11.1.0-hda6cdc1_1.tar.bz2"
+    else
+        echo "Platform not supported: $platform_id"
+        exit 3
+    fi
+    # Set up environment variables to configure cibuildwheel
+    export CIBW_BUILD=cp${cpython_ver}-${platform_id}
+    export CIBW_ARCHS=all
+    export CIBW_ENVIRONMENT=${setup_env_var}
+    export CIBW_TEST_SKIP='*-macosx_arm64'
+    export CIBW_BUILD_VERBOSITY=3
+
+    conda create -n build $OPENMP_URL
+    PREFIX="/usr/local/miniconda/envs/build"
+
+    # Set up build flags for cibuildwheel
+    # This is needed to bundle libomp lib we downloaded earlier
+    export CC=/usr/bin/clang
+    export CXX=/usr/bin/clang++
+    export CPPFLAGS="$CPPFLAGS -Xpreprocessor -fopenmp"
+    export CFLAGS="$CFLAGS -I$PREFIX/include"
+    export CXXFLAGS="$CXXFLAGS -I$PREFIX/include"
+    export LDFLAGS="$LDFLAGS -Wl,-rpath,$PREFIX/lib -L$PREFIX/lib -lomp"
+else
+    echo "Platform not supported: $platform_id"
+    exit 2
+fi
+
+python -m pip install cibuildwheel
+python -m cibuildwheel python --output-dir wheelhouse
+python -m cibuildwheel runtime/python --output-dir wheelhouse-runtime
+python tests/ci_build/rename_whl.py wheelhouse/*.whl ${commit_id} ${wheel_tag}
+python tests/ci_build/rename_whl.py wheelhouse-runtime/*.whl ${commit_id} ${wheel_tag}
+mv -v wheelhouse/*.whl .
+mv -v wheelhouse-runtime/*.whl .
diff --git a/tests/travis/run_test.sh b/tests/travis/run_test.sh
index d4b108cf..f7c89978 100755
--- a/tests/travis/run_test.sh
+++ b/tests/travis/run_test.sh
@@ -97,18 +97,6 @@ then
   python -m pip install --pre xgboost --no-binary :all:
   python -m pip install lightgbm --no-binary :all:
   python -m pytest -v --fulltrace tests/python
-
-  # Deploy binary wheel to S3
-  python -m pip install awscli
-  if [ "${TRAVIS_BRANCH}" == "mainline" ]
-  then
-    S3_DEST="s3://treelite-wheels/"
-  elif [ -z "${TRAVIS_TAG}" ]
-  then
-    S3_DEST="s3://treelite-wheels/${TRAVIS_BRANCH}/"
-  fi
-  python -m awscli s3 cp python/dist/treelite-*.whl "${S3_DEST}" --acl public-read || true
-  python -m awscli s3 cp runtime/python/dist/treelite_runtime*.whl "${S3_DEST}" --acl public-read || true
 fi
 
 if [ ${TASK} == "python_sdist_test" ]; then
@@ -145,3 +133,20 @@ if [ ${TASK} == "python_sdist_test" ]; then
   python -m awscli s3 cp treelite-*.tar.gz "${S3_DEST}" --acl public-read || true
   python -m awscli s3 cp treelite_runtime-*.tar.gz "${S3_DEST}" --acl public-read || true
 fi
+
+if [ ${TASK} == "python_wheels" ]; then
+  tests/ci_build/build_python_wheels.sh ${CIBW_PLATFORM_ID} ${TRAVIS_COMMIT}
+  # Deploy binary wheels to S3
+  conda activate python3
+  python --version
+  python -m pip install awscli
+  if [ "${TRAVIS_BRANCH}" == "mainline" ]
+  then
+    S3_DEST="s3://treelite-wheels/"
+  elif [ -z "${TRAVIS_TAG}" ]
+  then
+    S3_DEST="s3://treelite-wheels/${TRAVIS_BRANCH}/"
+  fi
+  python -m awscli s3 cp treelite-*.whl "${S3_DEST}" --acl public-read || true
+  python -m awscli s3 cp treelite_runtime-*.whl "${S3_DEST}" --acl public-read || true
+fi

From 9a9b101ff7303f94a8f06017030b25cfc9db3234 Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 19:06:44 -0800
Subject: [PATCH 04/18] Use correct Python

---
 tests/ci_build/build_python_wheels.sh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/ci_build/build_python_wheels.sh b/tests/ci_build/build_python_wheels.sh
index 2b53ca20..8f54a1c7 100755
--- a/tests/ci_build/build_python_wheels.sh
+++ b/tests/ci_build/build_python_wheels.sh
@@ -63,6 +63,7 @@ else
     exit 2
 fi
 
+conda activate python3
 python -m pip install cibuildwheel
 python -m cibuildwheel python --output-dir wheelhouse
 python -m cibuildwheel runtime/python --output-dir wheelhouse-runtime

From 7aa3d381a88a842b716680f4fc6a8289ebb385ba Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 19:11:18 -0800
Subject: [PATCH 05/18] fix

---
 tests/ci_build/build_python_wheels.sh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/ci_build/build_python_wheels.sh b/tests/ci_build/build_python_wheels.sh
index 8f54a1c7..5d22d174 100755
--- a/tests/ci_build/build_python_wheels.sh
+++ b/tests/ci_build/build_python_wheels.sh
@@ -63,6 +63,7 @@ else
     exit 2
 fi
 
+source $HOME/miniconda/bin/activate
 conda activate python3
 python -m pip install cibuildwheel
 python -m cibuildwheel python --output-dir wheelhouse

From ef5de4675106b0ea8601c53dbaf8eaab28073a72 Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 19:19:11 -0800
Subject: [PATCH 06/18] Take command-line args out of context

---
 tests/ci_build/build_python_wheels.sh | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/ci_build/build_python_wheels.sh b/tests/ci_build/build_python_wheels.sh
index 5d22d174..23d0fc35 100755
--- a/tests/ci_build/build_python_wheels.sh
+++ b/tests/ci_build/build_python_wheels.sh
@@ -9,7 +9,9 @@ if [[ $# -ne 2 ]]; then
 fi
 
 platform_id=$1
-commit_id=$2
+shift
+commit_id=$1
+shift
 
 # Bundle libomp 11.1.0 when targeting MacOS.
 # This is a workaround in order to prevent segfaults when running inside a Conda environment.

From 309990192518096d8a7b08a9b70373df4ff95a6f Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 19:26:04 -0800
Subject: [PATCH 07/18] fix

---
 tests/ci_build/build_python_wheels.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/ci_build/build_python_wheels.sh b/tests/ci_build/build_python_wheels.sh
index 23d0fc35..412cffff 100755
--- a/tests/ci_build/build_python_wheels.sh
+++ b/tests/ci_build/build_python_wheels.sh
@@ -49,7 +49,7 @@ if [[ "$platform_id" == macosx_* ]]; then
     export CIBW_TEST_SKIP='*-macosx_arm64'
     export CIBW_BUILD_VERBOSITY=3
 
-    conda create -n build $OPENMP_URL
+    sudo conda create -n build $OPENMP_URL
     PREFIX="/usr/local/miniconda/envs/build"
 
     # Set up build flags for cibuildwheel

From db405dc628042786fcef13484606c68f31bee3d6 Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 19:32:40 -0800
Subject: [PATCH 08/18] Show env path

---
 tests/ci_build/build_python_wheels.sh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/ci_build/build_python_wheels.sh b/tests/ci_build/build_python_wheels.sh
index 412cffff..3bc8a6ec 100755
--- a/tests/ci_build/build_python_wheels.sh
+++ b/tests/ci_build/build_python_wheels.sh
@@ -50,6 +50,7 @@ if [[ "$platform_id" == macosx_* ]]; then
     export CIBW_BUILD_VERBOSITY=3
 
     sudo conda create -n build $OPENMP_URL
+    conda info -e
     PREFIX="/usr/local/miniconda/envs/build"
 
     # Set up build flags for cibuildwheel

From 1f8cee92adf6686c544834ddc755b9d7c797afc8 Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 19:37:46 -0800
Subject: [PATCH 09/18] Use correct PREFIX

---
 tests/ci_build/build_python_wheels.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/ci_build/build_python_wheels.sh b/tests/ci_build/build_python_wheels.sh
index 3bc8a6ec..b260a28d 100755
--- a/tests/ci_build/build_python_wheels.sh
+++ b/tests/ci_build/build_python_wheels.sh
@@ -51,7 +51,7 @@ if [[ "$platform_id" == macosx_* ]]; then
 
     sudo conda create -n build $OPENMP_URL
     conda info -e
-    PREFIX="/usr/local/miniconda/envs/build"
+    PREFIX="$HOME/miniconda/envs/build"
 
     # Set up build flags for cibuildwheel
     # This is needed to bundle libomp lib we downloaded earlier

From 16148210aef0c50fb20885057365e152bd7b405a Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 20:12:47 -0800
Subject: [PATCH 10/18] Update setup.py

---
 python/setup.py         | 7 +++++++
 runtime/python/setup.py | 7 +++++++
 2 files changed, 14 insertions(+)

diff --git a/python/setup.py b/python/setup.py
index d7c296ca..38d429fc 100644
--- a/python/setup.py
+++ b/python/setup.py
@@ -98,6 +98,13 @@ def build(self, src_dir, build_dir, generator, build_tool=None):
         """Build the core library with CMake."""
         cmake_cmd = ['cmake', src_dir, generator]
 
+        # Flag for cross-compiling for Apple Silicon
+        # We use environment variable because it's the only way to pass down custom flags
+        # through the cibuildwheel package, which otherwise calls `python setup.py bdist_wheel`
+        # command.
+        if 'CIBW_TARGET_OSX_ARM64' in os.environ:
+            cmake_cmd.append("-DCMAKE_OSX_ARCHITECTURES=arm64")
+
         self.logger.info('Run CMake command: %s', str(cmake_cmd))
         subprocess.check_call(cmake_cmd, cwd=build_dir)
 
diff --git a/runtime/python/setup.py b/runtime/python/setup.py
index 9323692f..e3106cd8 100644
--- a/runtime/python/setup.py
+++ b/runtime/python/setup.py
@@ -97,6 +97,13 @@ def build(self, src_dir, build_dir, generator, build_tool=None, use_omp=1):
         """Build the runtime with CMake."""
         cmake_cmd = ['cmake', src_dir, generator]
 
+        # Flag for cross-compiling for Apple Silicon
+        # We use environment variable because it's the only way to pass down custom flags
+        # through the cibuildwheel package, which otherwise calls `python setup.py bdist_wheel`
+        # command.
+        if 'CIBW_TARGET_OSX_ARM64' in os.environ:
+            cmake_cmd.append("-DCMAKE_OSX_ARCHITECTURES=arm64")
+
         self.logger.info('Run CMake command: %s', str(cmake_cmd))
         subprocess.check_call(cmake_cmd, cwd=build_dir)
 

From c75adcbc39aba3772cfac74ac7b9730cc4602443 Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 22:15:11 -0800
Subject: [PATCH 11/18] Use int for thread ID

---
 src/annotator.cc                  | 4 ++--
 tests/cpp/test_threading_utils.cc | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/annotator.cc b/src/annotator.cc
index 7de197e6..4ce65d8d 100644
--- a/src/annotator.cc
+++ b/src/annotator.cc
@@ -79,7 +79,7 @@ inline void ComputeBranchLoopImpl(
   bool nan_missing = treelite::math::CheckNAN(missing_value);
   auto sched = treelite::threading_utils::ParallelSchedule::Static();
   treelite::threading_utils::ParallelFor(rbegin, rend, nthread, sched,
-                                         [&](std::size_t rid, std::size_t thread_id) {
+                                         [&](std::size_t rid, int thread_id) {
     const ElementType* row = &dmat->data[rid * num_col];
     const size_t off = dmat->num_col * thread_id;
     const size_t off2 = count_row_ptr[ntree] * thread_id;
@@ -110,7 +110,7 @@ inline void ComputeBranchLoopImpl(
   TREELITE_CHECK_LE(rbegin, rend);
   auto sched = treelite::threading_utils::ParallelSchedule::Static();
   treelite::threading_utils::ParallelFor(rbegin, rend, nthread, sched,
-                                         [&](std::size_t rid, std::size_t thread_id) {
+                                         [&](std::size_t rid, int thread_id) {
     const size_t off = dmat->num_col * thread_id;
     const size_t off2 = count_row_ptr[ntree] * thread_id;
     const size_t ibegin = dmat->row_ptr[rid];
diff --git a/tests/cpp/test_threading_utils.cc b/tests/cpp/test_threading_utils.cc
index 84b8f371..69019d64 100644
--- a/tests/cpp/test_threading_utils.cc
+++ b/tests/cpp/test_threading_utils.cc
@@ -58,7 +58,7 @@ TEST(ThreadingUtils, ParallelFor) {
 
   auto sched = treelite::threading_utils::ParallelSchedule::Guided();
 
-  auto dummy_func = [](int, std::size_t) {};
+  auto dummy_func = [](int, int) {};
   EXPECT_THROW(ParallelFor(0, 100, 0, sched, dummy_func), treelite::Error);
   EXPECT_THROW(ParallelFor(200, 100, 3, sched, dummy_func), treelite::Error);
   EXPECT_THROW(ParallelFor(-100, 100, 3, sched, dummy_func), treelite::Error);
@@ -85,7 +85,7 @@ TEST(ThreadingUtils, ParallelFor) {
     std::size_t nthread = static_cast<std::size_t>(rng.DrawInteger(1, max_thread + 1));
     int64_t end = rng.DrawInteger(begin, kVectorLength);
 
-    ParallelFor(begin, end, nthread, sched, [&a, &b, &c](int64_t i, std::size_t) {
+    ParallelFor(begin, end, nthread, sched, [&a, &b, &c](int64_t i, int) {
       c[i] = a[i] + b[i];
     });
 

From fc6d25964e0b50880c9c9549571944798cffc94d Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 22:47:41 -0800
Subject: [PATCH 12/18] Don't bundle libomp in runtime package

---
 tests/ci_build/build_python_wheels.sh | 35 ++++++++++++++-------------
 tests/travis/run_test.sh              |  3 ++-
 2 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/tests/ci_build/build_python_wheels.sh b/tests/ci_build/build_python_wheels.sh
index b260a28d..441323e7 100755
--- a/tests/ci_build/build_python_wheels.sh
+++ b/tests/ci_build/build_python_wheels.sh
@@ -3,8 +3,8 @@
 set -e
 set -x
 
-if [[ $# -ne 2 ]]; then
-  echo "Usage: $0 [platform_id] [commit ID]"
+if [[ $# -ne 3 ]]; then
+  echo "Usage: $0 [platform_id] [commit ID] [python package location]"
   exit 1
 fi
 
@@ -12,6 +12,8 @@ platform_id=$1
 shift
 commit_id=$1
 shift
+python_pkg_root=$1
+shift
 
 # Bundle libomp 11.1.0 when targeting MacOS.
 # This is a workaround in order to prevent segfaults when running inside a Conda environment.
@@ -49,18 +51,20 @@ if [[ "$platform_id" == macosx_* ]]; then
     export CIBW_TEST_SKIP='*-macosx_arm64'
     export CIBW_BUILD_VERBOSITY=3
 
-    sudo conda create -n build $OPENMP_URL
-    conda info -e
-    PREFIX="$HOME/miniconda/envs/build"
+    if [[ "${NO_OPENMP}" != 1 ]]; then
+        sudo conda create -n build $OPENMP_URL
+        conda info -e
+        PREFIX="$HOME/miniconda/envs/build"
 
-    # Set up build flags for cibuildwheel
-    # This is needed to bundle libomp lib we downloaded earlier
-    export CC=/usr/bin/clang
-    export CXX=/usr/bin/clang++
-    export CPPFLAGS="$CPPFLAGS -Xpreprocessor -fopenmp"
-    export CFLAGS="$CFLAGS -I$PREFIX/include"
-    export CXXFLAGS="$CXXFLAGS -I$PREFIX/include"
-    export LDFLAGS="$LDFLAGS -Wl,-rpath,$PREFIX/lib -L$PREFIX/lib -lomp"
+        # Set up build flags for cibuildwheel
+        # This is needed to bundle libomp lib we downloaded earlier
+        export CC=/usr/bin/clang
+        export CXX=/usr/bin/clang++
+        export CPPFLAGS="$CPPFLAGS -Xpreprocessor -fopenmp"
+        export CFLAGS="$CFLAGS -I$PREFIX/include"
+        export CXXFLAGS="$CXXFLAGS -I$PREFIX/include"
+        export LDFLAGS="$LDFLAGS -Wl,-rpath,$PREFIX/lib -L$PREFIX/lib -lomp"
+    fi
 else
     echo "Platform not supported: $platform_id"
     exit 2
@@ -69,9 +73,6 @@ fi
 source $HOME/miniconda/bin/activate
 conda activate python3
 python -m pip install cibuildwheel
-python -m cibuildwheel python --output-dir wheelhouse
-python -m cibuildwheel runtime/python --output-dir wheelhouse-runtime
+python -m cibuildwheel ${python_pkg_root} --output-dir wheelhouse
 python tests/ci_build/rename_whl.py wheelhouse/*.whl ${commit_id} ${wheel_tag}
-python tests/ci_build/rename_whl.py wheelhouse-runtime/*.whl ${commit_id} ${wheel_tag}
 mv -v wheelhouse/*.whl .
-mv -v wheelhouse-runtime/*.whl .
diff --git a/tests/travis/run_test.sh b/tests/travis/run_test.sh
index f7c89978..d21af79a 100755
--- a/tests/travis/run_test.sh
+++ b/tests/travis/run_test.sh
@@ -135,7 +135,8 @@ if [ ${TASK} == "python_sdist_test" ]; then
 fi
 
 if [ ${TASK} == "python_wheels" ]; then
-  tests/ci_build/build_python_wheels.sh ${CIBW_PLATFORM_ID} ${TRAVIS_COMMIT}
+  tests/ci_build/build_python_wheels.sh ${CIBW_PLATFORM_ID} ${TRAVIS_COMMIT} python
+  NO_OPENMP=1 tests/ci_build/build_python_wheels.sh ${CIBW_PLATFORM_ID} ${TRAVIS_COMMIT} runtime/python
   # Deploy binary wheels to S3
   conda activate python3
   python --version

From 28105c95f16ea371b84d829405f8ab46d6711e6f Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 22:56:02 -0800
Subject: [PATCH 13/18] Revert "Don't bundle libomp in runtime package"

This reverts commit fc6d25964e0b50880c9c9549571944798cffc94d.
---
 tests/ci_build/build_python_wheels.sh | 35 +++++++++++++--------------
 tests/travis/run_test.sh              |  3 +--
 2 files changed, 18 insertions(+), 20 deletions(-)

diff --git a/tests/ci_build/build_python_wheels.sh b/tests/ci_build/build_python_wheels.sh
index 441323e7..b260a28d 100755
--- a/tests/ci_build/build_python_wheels.sh
+++ b/tests/ci_build/build_python_wheels.sh
@@ -3,8 +3,8 @@
 set -e
 set -x
 
-if [[ $# -ne 3 ]]; then
-  echo "Usage: $0 [platform_id] [commit ID] [python package location]"
+if [[ $# -ne 2 ]]; then
+  echo "Usage: $0 [platform_id] [commit ID]"
   exit 1
 fi
 
@@ -12,8 +12,6 @@ platform_id=$1
 shift
 commit_id=$1
 shift
-python_pkg_root=$1
-shift
 
 # Bundle libomp 11.1.0 when targeting MacOS.
 # This is a workaround in order to prevent segfaults when running inside a Conda environment.
@@ -51,20 +49,18 @@ if [[ "$platform_id" == macosx_* ]]; then
     export CIBW_TEST_SKIP='*-macosx_arm64'
     export CIBW_BUILD_VERBOSITY=3
 
-    if [[ "${NO_OPENMP}" != 1 ]]; then
-        sudo conda create -n build $OPENMP_URL
-        conda info -e
-        PREFIX="$HOME/miniconda/envs/build"
+    sudo conda create -n build $OPENMP_URL
+    conda info -e
+    PREFIX="$HOME/miniconda/envs/build"
 
-        # Set up build flags for cibuildwheel
-        # This is needed to bundle libomp lib we downloaded earlier
-        export CC=/usr/bin/clang
-        export CXX=/usr/bin/clang++
-        export CPPFLAGS="$CPPFLAGS -Xpreprocessor -fopenmp"
-        export CFLAGS="$CFLAGS -I$PREFIX/include"
-        export CXXFLAGS="$CXXFLAGS -I$PREFIX/include"
-        export LDFLAGS="$LDFLAGS -Wl,-rpath,$PREFIX/lib -L$PREFIX/lib -lomp"
-    fi
+    # Set up build flags for cibuildwheel
+    # This is needed to bundle libomp lib we downloaded earlier
+    export CC=/usr/bin/clang
+    export CXX=/usr/bin/clang++
+    export CPPFLAGS="$CPPFLAGS -Xpreprocessor -fopenmp"
+    export CFLAGS="$CFLAGS -I$PREFIX/include"
+    export CXXFLAGS="$CXXFLAGS -I$PREFIX/include"
+    export LDFLAGS="$LDFLAGS -Wl,-rpath,$PREFIX/lib -L$PREFIX/lib -lomp"
 else
     echo "Platform not supported: $platform_id"
     exit 2
@@ -73,6 +69,9 @@ fi
 source $HOME/miniconda/bin/activate
 conda activate python3
 python -m pip install cibuildwheel
-python -m cibuildwheel ${python_pkg_root} --output-dir wheelhouse
+python -m cibuildwheel python --output-dir wheelhouse
+python -m cibuildwheel runtime/python --output-dir wheelhouse-runtime
 python tests/ci_build/rename_whl.py wheelhouse/*.whl ${commit_id} ${wheel_tag}
+python tests/ci_build/rename_whl.py wheelhouse-runtime/*.whl ${commit_id} ${wheel_tag}
 mv -v wheelhouse/*.whl .
+mv -v wheelhouse-runtime/*.whl .
diff --git a/tests/travis/run_test.sh b/tests/travis/run_test.sh
index d21af79a..f7c89978 100755
--- a/tests/travis/run_test.sh
+++ b/tests/travis/run_test.sh
@@ -135,8 +135,7 @@ if [ ${TASK} == "python_sdist_test" ]; then
 fi
 
 if [ ${TASK} == "python_wheels" ]; then
-  tests/ci_build/build_python_wheels.sh ${CIBW_PLATFORM_ID} ${TRAVIS_COMMIT} python
-  NO_OPENMP=1 tests/ci_build/build_python_wheels.sh ${CIBW_PLATFORM_ID} ${TRAVIS_COMMIT} runtime/python
+  tests/ci_build/build_python_wheels.sh ${CIBW_PLATFORM_ID} ${TRAVIS_COMMIT}
   # Deploy binary wheels to S3
   conda activate python3
   python --version

From 466a2ddb5980f7bfb63d763d3354f4498f442c96 Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 23:07:16 -0800
Subject: [PATCH 14/18] Run auditwheel on Linux wheels

---
 azure-pipelines.yml                         | 11 ++++++++---
 tests/ci_build/Dockerfile.auditwheel_x86_64 | 15 +++++++++++++++
 2 files changed, 23 insertions(+), 3 deletions(-)
 create mode 100644 tests/ci_build/Dockerfile.auditwheel_x86_64

diff --git a/azure-pipelines.yml b/azure-pipelines.yml
index ba86b871..26d639cb 100644
--- a/azure-pipelines.yml
+++ b/azure-pipelines.yml
@@ -24,14 +24,22 @@ jobs:
   - script: tests/ci_build/ci_build.sh cpu tests/ci_build/build_via_cmake.sh
     displayName: 'Building Treelite...'
   - script: |
+      TAG=manylinux2014_x86_64
       tests/ci_build/ci_build.sh cpu bash -c "cd python/ && python setup.py bdist_wheel --universal"
+      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py python/dist/*.whl $(Build.SourceVersion) ${TAG}
+      tests/ci_build/ci_build.sh auditwheel_x86_64 auditwheel repair --plat ${TAG} python/dist/*.whl
+      mv -v wheelhouse/*.whl python/dist/
     displayName: 'Packaging Python wheel for Treelite...'
   - task: PublishPipelineArtifact@0
     inputs:
       artifactName: 'python_linux_whl'
       targetPath: 'python/dist/'
   - script: |
+      TAG=manylinux2014_x86_64
       tests/ci_build/ci_build.sh cpu bash -c "cd runtime/python/ && python setup.py bdist_wheel --universal"
+      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py runtime/python/dist/*.whl $(Build.SourceVersion) ${TAG}
+      tests/ci_build/ci_build.sh auditwheel_x86_64 auditwheel repair --plat ${TAG} runtime/python/dist/*.whl
+      mv -v wheelhouse/*.whl runtime/python/dist/
     displayName: 'Packaging Python wheel for Treelite runtime...'
   - task: PublishPipelineArtifact@1
     inputs:
@@ -210,9 +218,6 @@ jobs:
   - script: python -m pytest -v --fulltrace tests/python/test_basic.py
     displayName: 'Running Python tests...'
   - script: |
-      TAG=manylinux2014_x86_64
-      python tests/ci_build/rename_whl.py main/*.whl $(Build.SourceVersion) ${TAG}
-      python tests/ci_build/rename_whl.py runtime/*.whl $(Build.SourceVersion) ${TAG}
       python -m awscli s3 cp main/*.whl s3://treelite-wheels/ --acl public-read || true
       python -m awscli s3 cp runtime/*.whl s3://treelite-wheels/ --acl public-read || true
     displayName: 'Uploading Python wheels...'
diff --git a/tests/ci_build/Dockerfile.auditwheel_x86_64 b/tests/ci_build/Dockerfile.auditwheel_x86_64
new file mode 100644
index 00000000..6c215014
--- /dev/null
+++ b/tests/ci_build/Dockerfile.auditwheel_x86_64
@@ -0,0 +1,15 @@
+FROM quay.io/pypa/manylinux2014_x86_64
+
+# Install lightweight sudo (not bound to TTY)
+ENV GOSU_VERSION 1.10
+RUN set -ex; \
+    curl -o /usr/local/bin/gosu -L "https://github.com/tianon/gosu/releases/download/$GOSU_VERSION/gosu-amd64" && \
+    chmod +x /usr/local/bin/gosu && \
+    gosu nobody true
+
+# Default entry-point to use if running locally
+# It will preserve attributes of created files
+COPY entrypoint.sh /scripts/
+
+WORKDIR /workspace
+ENTRYPOINT ["/scripts/entrypoint.sh"]

From 3f9b8384ff1293de1a28bb5f43c5c85535eb5be5 Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 23:14:59 -0800
Subject: [PATCH 15/18] Rename after repair

---
 azure-pipelines.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/azure-pipelines.yml b/azure-pipelines.yml
index 26d639cb..803f6e48 100644
--- a/azure-pipelines.yml
+++ b/azure-pipelines.yml
@@ -26,9 +26,9 @@ jobs:
   - script: |
       TAG=manylinux2014_x86_64
       tests/ci_build/ci_build.sh cpu bash -c "cd python/ && python setup.py bdist_wheel --universal"
-      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py python/dist/*.whl $(Build.SourceVersion) ${TAG}
       tests/ci_build/ci_build.sh auditwheel_x86_64 auditwheel repair --plat ${TAG} python/dist/*.whl
       mv -v wheelhouse/*.whl python/dist/
+      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py python/dist/*.whl $(Build.SourceVersion) ${TAG}
     displayName: 'Packaging Python wheel for Treelite...'
   - task: PublishPipelineArtifact@0
     inputs:
@@ -37,9 +37,9 @@ jobs:
   - script: |
       TAG=manylinux2014_x86_64
       tests/ci_build/ci_build.sh cpu bash -c "cd runtime/python/ && python setup.py bdist_wheel --universal"
-      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py runtime/python/dist/*.whl $(Build.SourceVersion) ${TAG}
       tests/ci_build/ci_build.sh auditwheel_x86_64 auditwheel repair --plat ${TAG} runtime/python/dist/*.whl
       mv -v wheelhouse/*.whl runtime/python/dist/
+      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py runtime/python/dist/*.whl $(Build.SourceVersion) ${TAG}
     displayName: 'Packaging Python wheel for Treelite runtime...'
   - task: PublishPipelineArtifact@1
     inputs:

From fe20d9a3b05426eb891ae6be1df109a391cacaee Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 23:24:18 -0800
Subject: [PATCH 16/18] Make sure that repaired wheel overwrites original

---
 azure-pipelines.yml | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/azure-pipelines.yml b/azure-pipelines.yml
index 803f6e48..083e7084 100644
--- a/azure-pipelines.yml
+++ b/azure-pipelines.yml
@@ -26,9 +26,10 @@ jobs:
   - script: |
       TAG=manylinux2014_x86_64
       tests/ci_build/ci_build.sh cpu bash -c "cd python/ && python setup.py bdist_wheel --universal"
+      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py python/dist/*.whl $(Build.SourceVersion) ${TAG}
       tests/ci_build/ci_build.sh auditwheel_x86_64 auditwheel repair --plat ${TAG} python/dist/*.whl
+      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py wheelhouse/*.whl $(Build.SourceVersion) ${TAG}
       mv -v wheelhouse/*.whl python/dist/
-      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py python/dist/*.whl $(Build.SourceVersion) ${TAG}
     displayName: 'Packaging Python wheel for Treelite...'
   - task: PublishPipelineArtifact@0
     inputs:
@@ -37,9 +38,10 @@ jobs:
   - script: |
       TAG=manylinux2014_x86_64
       tests/ci_build/ci_build.sh cpu bash -c "cd runtime/python/ && python setup.py bdist_wheel --universal"
+      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py runtime/python/dist/*.whl $(Build.SourceVersion) ${TAG}
       tests/ci_build/ci_build.sh auditwheel_x86_64 auditwheel repair --plat ${TAG} runtime/python/dist/*.whl
+      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py wheelhouse/*.whl $(Build.SourceVersion) ${TAG}
       mv -v wheelhouse/*.whl runtime/python/dist/
-      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py runtime/python/dist/*.whl $(Build.SourceVersion) ${TAG}
     displayName: 'Packaging Python wheel for Treelite runtime...'
   - task: PublishPipelineArtifact@1
     inputs:

From aa595ebad850a1f66018c912731c9c458290807a Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 23:27:05 -0800
Subject: [PATCH 17/18] Rename only once

---
 azure-pipelines.yml | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/azure-pipelines.yml b/azure-pipelines.yml
index 083e7084..2a7efd3a 100644
--- a/azure-pipelines.yml
+++ b/azure-pipelines.yml
@@ -27,8 +27,7 @@ jobs:
       TAG=manylinux2014_x86_64
       tests/ci_build/ci_build.sh cpu bash -c "cd python/ && python setup.py bdist_wheel --universal"
       tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py python/dist/*.whl $(Build.SourceVersion) ${TAG}
-      tests/ci_build/ci_build.sh auditwheel_x86_64 auditwheel repair --plat ${TAG} python/dist/*.whl
-      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py wheelhouse/*.whl $(Build.SourceVersion) ${TAG}
+      tests/ci_build/ci_build.sh auditwheel_x86_64 auditwheel repair --only-plat --plat ${TAG} python/dist/*.whl
       mv -v wheelhouse/*.whl python/dist/
     displayName: 'Packaging Python wheel for Treelite...'
   - task: PublishPipelineArtifact@0
@@ -39,8 +38,7 @@ jobs:
       TAG=manylinux2014_x86_64
       tests/ci_build/ci_build.sh cpu bash -c "cd runtime/python/ && python setup.py bdist_wheel --universal"
       tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py runtime/python/dist/*.whl $(Build.SourceVersion) ${TAG}
-      tests/ci_build/ci_build.sh auditwheel_x86_64 auditwheel repair --plat ${TAG} runtime/python/dist/*.whl
-      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py wheelhouse/*.whl $(Build.SourceVersion) ${TAG}
+      tests/ci_build/ci_build.sh auditwheel_x86_64 auditwheel repair --only-plat --plat ${TAG} runtime/python/dist/*.whl
       mv -v wheelhouse/*.whl runtime/python/dist/
     displayName: 'Packaging Python wheel for Treelite runtime...'
   - task: PublishPipelineArtifact@1

From 5bb230bfcbd5eccaa769ad98d099e8471417f202 Mon Sep 17 00:00:00 2001
From: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
Date: Wed, 2 Feb 2022 23:35:16 -0800
Subject: [PATCH 18/18] One last try

---
 azure-pipelines.yml | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/azure-pipelines.yml b/azure-pipelines.yml
index 2a7efd3a..8acbd91a 100644
--- a/azure-pipelines.yml
+++ b/azure-pipelines.yml
@@ -26,9 +26,10 @@ jobs:
   - script: |
       TAG=manylinux2014_x86_64
       tests/ci_build/ci_build.sh cpu bash -c "cd python/ && python setup.py bdist_wheel --universal"
-      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py python/dist/*.whl $(Build.SourceVersion) ${TAG}
       tests/ci_build/ci_build.sh auditwheel_x86_64 auditwheel repair --only-plat --plat ${TAG} python/dist/*.whl
+      rm -v python/dist/*.whl
       mv -v wheelhouse/*.whl python/dist/
+      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py python/dist/*.whl $(Build.SourceVersion) ${TAG}
     displayName: 'Packaging Python wheel for Treelite...'
   - task: PublishPipelineArtifact@0
     inputs:
@@ -37,9 +38,10 @@ jobs:
   - script: |
       TAG=manylinux2014_x86_64
       tests/ci_build/ci_build.sh cpu bash -c "cd runtime/python/ && python setup.py bdist_wheel --universal"
-      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py runtime/python/dist/*.whl $(Build.SourceVersion) ${TAG}
       tests/ci_build/ci_build.sh auditwheel_x86_64 auditwheel repair --only-plat --plat ${TAG} runtime/python/dist/*.whl
+      rm -v runtime/python/dist/*.whl
       mv -v wheelhouse/*.whl runtime/python/dist/
+      tests/ci_build/ci_build.sh cpu python tests/ci_build/rename_whl.py runtime/python/dist/*.whl $(Build.SourceVersion) ${TAG}
     displayName: 'Packaging Python wheel for Treelite runtime...'
   - task: PublishPipelineArtifact@1
     inputs: