Add GEM Blockset

.github/workflows/ci.yml

@@ -16,12 +16,20 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
+        with:
+          submodules: recursive
 
       - name: Install Rust
         uses: dtolnay/rust-toolchain@stable
         with:
           components: rustfmt, clippy
 
+      - name: Set up CUDA
+        uses: Jimver/cuda-toolkit@v0.2.23
+        with:
+          cuda: '12.8.0'
+          log-file-suffix: '-${{ github.job }}'
+
       - name: Cache cargo
         uses: actions/cache@v4
         with:
@@ -31,7 +39,7 @@ jobs:
             ~/.cargo/registry/cache/
             ~/.cargo/git/db/
             target/
-          key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
+          key: ${{ runner.os }}-cargo-check-${{ hashFiles('**/Cargo.lock') }}
 
       - name: Check formatting
         run: cargo fmt --all -- --check
@@ -44,10 +52,18 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
+        with:
+          submodules: recursive
 
       - name: Install Rust
         uses: dtolnay/rust-toolchain@stable
 
+      - name: Set up CUDA
+        uses: Jimver/cuda-toolkit@v0.2.23
+        with:
+          cuda: '12.8.0'
+          log-file-suffix: '-${{ github.job }}'
+
       - name: Cache cargo
         uses: actions/cache@v4
         with:
@@ -57,7 +73,7 @@ jobs:
             ~/.cargo/registry/cache/
             ~/.cargo/git/db/
             target/
-          key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
+          key: ${{ runner.os }}-cargo-test-${{ hashFiles('**/Cargo.lock') }}
 
       - name: Run tests
         run: cargo test --all-features
@@ -68,10 +84,18 @@ jobs:
     needs: [check, test]
     steps:
       - uses: actions/checkout@v4
+        with:
+          submodules: recursive
 
       - name: Install Rust
         uses: dtolnay/rust-toolchain@stable
 
+      - name: Set up CUDA
+        uses: Jimver/cuda-toolkit@v0.2.23
+        with:
+          cuda: '12.8.0'
+          log-file-suffix: '-${{ github.job }}'
+
       - name: Cache cargo
         uses: actions/cache@v4
         with:
@@ -81,13 +105,13 @@ jobs:
             ~/.cargo/registry/cache/
             ~/.cargo/git/db/
             target/
-          key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
+          key: ${{ runner.os }}-cargo-build-${{ hashFiles('**/Cargo.lock') }}
 
       - name: Build release
         run: cargo build --release
 
       - name: Upload binary
         uses: actions/upload-artifact@v4
         with:
-          name: cubert
+          name: cubert-binary
           path: target/release/cubert

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "dependencies/gem-blockset"]
+	path = dependencies/gem-blockset
+	url = https://github.com/Rosnaky/gem-blockset.git

Cargo.toml

@@ -19,3 +19,7 @@ uuid = { version = "1.21.0", features = ["v4"] }
 [dev-dependencies]
 tempfile = "3"
 tokio-test = "0.4"
+
+[build-dependencies]
+bindgen = "0.72.1"
+cmake = "0.1.57"

build.rs

@@ -1,11 +1,39 @@
+use std::env;
+use std::path::PathBuf;
+
 fn main() {
-    let libpath = format!("{}/model", env!("CARGO_MANIFEST_DIR"));
+    let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());
+    let gem_dir = manifest_dir.join("dependencies/gem-blockset");
+
+    let cuda_path = env::var("CUDA_PATH")
+        .or_else(|_| env::var("CUDA_HOME"))
+        .unwrap_or_else(|_| "/opt/cuda".to_string());
+
+    let dst = cmake::Config::new(&gem_dir)
+        .define("CMAKE_CUDA_ARCHITECTURES", "75")
+        .define("CMAKE_CUDA_COMPILER", format!("{}/bin/nvcc", cuda_path))
+        .define("GEM_BUILD_FFI", "ON")
+        .build_target("gem_blockset")
+        .build();
+
+    println!("cargo:rustc-link-search=native={}/build", dst.display());
+    println!("cargo:rustc-link-lib=static=gem_blockset");
+
+    println!("cargo:rustc-link-search=native={}/lib64", cuda_path);
+    println!("cargo:rustc-link-search=native={}/lib", cuda_path);
+    println!("cargo:rustc-link-lib=dylib=cudart");
 
-    println!("cargo:rustc-link-search=native={}", libpath);
+    println!("cargo:rustc-link-lib=dylib=stdc++");
 
-    println!("cargo:rustc-link-lib=cudastats");
+    let bindings = bindgen::Builder::default()
+        .header(gem_dir.join("models/ffi.h").to_str().unwrap())
+        .generate()
+        .expect("failed to generate bindings");
 
-    println!("cargo:rustc-link-arg=-Wl,-rpath,{}", libpath);
+    let out_path = PathBuf::from(env::var("OUT_DIR").unwrap());
+    bindings
+        .write_to_file(out_path.join("gem_bindings.rs"))
+        .expect("failed to write bindings");
 
-    println!("cargo:rerun-if-changed=model/libcudastats.so");
+    println!("cargo:rerun-if-changed=dependencies/gem-blockset/models");
 }