From b541f175394396e62fd27ff9b1bdaa22e48eef2c Mon Sep 17 00:00:00 2001
From: Junpei Kawamoto <kawamoto.junpei@gmail.com>
Date: Wed, 29 May 2024 00:58:45 -0600
Subject: [PATCH] test

---
 .github/workflows/build.yaml | 41 +++++++++++++++++++++++++++---------
 Cargo.toml                   |  1 +
 build.rs                     | 10 ++++++++-
 3 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
index c78565b..7a82d36 100644
--- a/.github/workflows/build.yaml
+++ b/.github/workflows/build.yaml
@@ -23,17 +23,30 @@ jobs:
       fail-fast: false
       matrix:
         os: [macos-13, macos-14]
-        feature: [ruy, accelerate]
+        backends: [ruy, openblas, accelerate]
+        openmp: ["", openmp]
+        exclude:
+          - os: macos-13
+            backends: openblas
+          - os: macos-14
+            backends: openblas
+            openmp: ""
     runs-on: ${{ matrix.os }}
 
     steps:
       - name: Install OpenBLAS
-        if: ${{ matrix.feature == 'openblas' }}
+        if: ${{ matrix.backends == 'openblas' }}
         run: |
           brew install openblas
           PREFIX=$(brew --prefix openblas)
           echo "LIBRARY_PATH=$PREFIX/lib:$LIBRARY_PATH" >> $GITHUB_ENV
-          echo "LD_LIBRARY_PATH=$PREFIX/lib:$LD_LIBRARY_PATH" >> $GITHUB_ENV
+          echo "CMAKE_PREFIX_PATH=$PREFIX:$CMAKE_PREFIX_PATH" >> $GITHUB_ENV
+      - name: Install OpenMP
+        if: ${{ matrix.openmp == 'openmp' }}
+        run: |
+          brew install libomp
+          PREFIX=$(brew --prefix libomp)
+          echo "LIBRARY_PATH=$PREFIX/lib:$LIBRARY_PATH" >> $GITHUB_ENV
           echo "CMAKE_PREFIX_PATH=$PREFIX:$CMAKE_PREFIX_PATH" >> $GITHUB_ENV
       - uses: actions/checkout@v4
         with:
@@ -41,22 +54,23 @@ jobs:
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
       - name: Build
-        run: cargo build -vv --no-default-features -F "${{ matrix.feature }}"
+        run: cargo build -vv --no-default-features -F "${{ matrix.backends }} ${{ matrix.openmp }}"
       - name: Run tests
-        run: cargo test -vv --no-default-features -F "${{ matrix.feature }}"
+        run: cargo test -vv --no-default-features -F "${{ matrix.backends }} ${{ matrix.openmp }}"
 
   build-linux:
     strategy:
       fail-fast: false
       matrix:
-        feature: [ruy, mkl, openblas]
+        backends: [ruy, mkl, openblas]
+        openmp: ["", openmp]
     runs-on: ubuntu-latest
     env:
       LIBRARY_PATH: /usr/lib/x86_64-linux-gnu
 
     steps:
       - name: Install MKL
-        if: ${{ matrix.feature == 'mkl' }}
+        if: ${{ matrix.backends == 'mkl' }}
         run: |
           wget https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS-2023.PUB
           sudo apt-key add GPG-PUB-KEY-INTEL-SW-PRODUCTS-2023.PUB
@@ -64,20 +78,27 @@ jobs:
           sudo apt-get update
           sudo apt-get install -y intel-mkl-64bit-2020.4-912
           echo "LD_LIBRARY_PATH=/opt/intel/lib/intel64:$LD_LIBRARY_PATH" >> $GITHUB_ENV
+          echo "LIBRARY_PATH=/opt/intel/lib/intel64:LIBRARY_PATH" >> $GITHUB_ENV
       - name: Install OpenBLAS
-        if: ${{ matrix.feature == 'openblas' }}
+        if: ${{ matrix.backends == 'openblas' }}
         run: |
           sudo apt-get update
           sudo apt-get install -y cmake libclang-dev libopenblas-dev
+      - name: Install OpenMP
+        if: ${{ matrix.openmp == 'openmp' }}
+        run: |
+          sudo apt-get install -y libomp-dev
+          echo "LD_LIBRARY_PATH=/usr/lib/llvm-14/lib:$LD_LIBRARY_PATH" >> $GITHUB_ENV
+          echo "LIBRARY_PATH=/usr/lib/llvm-14/lib:$LIBRARY_PATH" >> $GITHUB_ENV
       - uses: actions/checkout@v4
         with:
           submodules: recursive
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
       - name: Build
-        run: cargo build -vv --no-default-features -F "${{ matrix.feature }}"
+        run: cargo build -vv --no-default-features -F "${{ matrix.backends }} ${{ matrix.openmp }}"
       - name: Run tests
-        run: cargo test -vv --no-default-features -F "${{ matrix.feature }}"
+        run: cargo test -vv --no-default-features -F "${{ matrix.backends }} ${{ matrix.openmp }}"
 
   build-windows:
     strategy:
diff --git a/Cargo.toml b/Cargo.toml
index 12706f5..72855f0 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -53,6 +53,7 @@ ruy = []
 accelerate = []
 cuda = []
 cudnn = ["cuda"]
+openmp = []
 
 [[example]]
 name = "bart"
diff --git a/build.rs b/build.rs
index adbb6d6..53122f3 100644
--- a/build.rs
+++ b/build.rs
@@ -41,7 +41,15 @@ fn main() {
         .define("BUILD_CLI", "OFF")
         .define("BUILD_SHARED_LIBS", "OFF")
         .define("WITH_MKL", "OFF")
-        .define("OPENMP_RUNTIME", "NONE");
+        .define(
+            "OPENMP_RUNTIME",
+            if cfg!(feature = "openmp") {
+                println!("cargo:rustc-link-lib=omp");
+                "COMP"
+            } else {
+                "NONE"
+            },
+        );
     if cfg!(target_os = "windows") {
         println!("cargo::rustc-link-arg=/FORCE:MULTIPLE");
         cmake.profile("Release");