build.rs

use std::path::Path;
use std::process::Command;
use std::{env, path::PathBuf};

fn main() {
    // let _cpphelpers = cmake::Config::new("cpphelpers").build();
    let zfp_dst = cmake::Config::new("rust/external/zfp")
        .define("BUILD_SHARED_LIBS", "OFF")
        .build();
    let eigen_dst = cmake::Config::new("rust/external/eigen")
        .define("CMAKE_POLICY_VERSION_MINIMUM", "3.5")
        .build();
    let octree_src = Path::new("rust/external/tucker-octree");
    let cpp_file = octree_src.join("toctree.cpp");
    let mut octree_compiled = false;

    let patch_1 = Command::new("sed")
        .args(&["-i", "s/ColMajor/RowMajor/g", cpp_file.to_str().unwrap()])
        .status();

    let patch_2 = Command::new("sed")
        .args(&[
            "-i",
            "s/Eigen::Vector<T, Eigen::Dynamic>/Eigen::Matrix<T, Eigen::Dynamic, 1>/g",
            cpp_file.to_str().unwrap(),
        ])
        .status();

    if patch_1.map_or(false, |s| s.success()) && patch_2.map_or(false, |s| s.success()) {
        let octree_dst_result = cmake::Config::new(octree_src)
            .define("TOCTREE_L2ERROR", "true")
            .define("CMAKE_BUILD_TYPE", "Release")
            .define("BUILD_SHARED_LIBS", "OFF")
            .define("zfp_DIR", format!("{}/lib/cmake/zfp", zfp_dst.display()))
            .define(
                "Eigen3_DIR",
                format!("{}/share/eigen3/cmake", eigen_dst.display()),
            )
            .build();

        let lib_path = octree_dst_result.join("lib/libtoctree_compressor.a");
        if lib_path.exists() {
            octree_compiled = true;
            setup_octree_linking(&octree_dst_result, &zfp_dst);
        }
    }

    if !octree_compiled {
        println!("cargo:warning=Octree build failed. Setting no_octree.");
        println!("cargo:rustc-cfg=no_octree");
    }

    let mut ml_compiled = false;
    let asterix_dir = Path::new("rust/external/asterix");
    let mlp_src = asterix_dir.join("src/vdf_compressor_nn.cu");
    let include_path = asterix_dir.join("include");

    let root_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());
    let ml_lib_name = "vlasiator_vdf_compressor_nn";
    let ml_lib_path = root_dir.join(format!("lib{}.so", ml_lib_name));

    let compiler = if is_program_in_path("nvcc") {
        Some("nvcc")
    } else if is_program_in_path("hipcc") {
        Some("hipcc")
    } else {
        None
    };

    if let Some(cc) = compiler {
        println!("cargo:rerun-if-changed={}", mlp_src.display());
        let mut gpu_cmd = Command::new(cc);
        if cc == "nvcc" {
            gpu_cmd.args(&[
                mlp_src.to_str().unwrap(),
                "--std=c++20",
                "-DTINYAI_MEMORY_GB=4",
                "-DNOPROFILE",
                &format!("-I{}", include_path.display()),
                "--shared",
                "-o",
                ml_lib_path.to_str().unwrap(),
                "-Xcompiler=-fPIC",
                "-lcublas",
                "-lblas",
            ]);
        } else if cc == "hipcc" {
            gpu_cmd.args(&[
                mlp_src.to_str().unwrap(),
                "--std=c++20",
                "-I/opt/rocm/include/hipblas",
                "-DTINYAI_MEMORY_GB=4",
                "-DNOPROFILE",
                &format!("-I{}", include_path.display()),
                "--shared",
                "-o",
                ml_lib_path.to_str().unwrap(),
                "-fPIC",
                "-lhipblas",
                "-DSKIP_HOSTBLAS",
            ]);
        }

        if gpu_cmd.status().map(|s| s.success()).unwrap_or(false) {
            ml_compiled = true;
            println!("cargo:rustc-link-search=native={}", root_dir.display());
            println!("cargo:rustc-link-lib=dylib={}", ml_lib_name);

            if env::var("CARGO_CFG_TARGET_OS").as_deref() == Ok("linux") {
                println!("cargo:rustc-link-arg=-Wl,-rpath,{}", root_dir.display());
                println!(
                    "cargo:warning=MLP Library created at: {}",
                    ml_lib_path.display()
                );
            }
            if cc == "nvcc" {
                println!("cargo:rustc-link-search=native=/usr/local/cuda/lib64");
                println!("cargo:rustc-link-lib=cudart");
                println!("cargo:rustc-link-lib=cublas");
            } else if cc == "hipcc" {
                println!("cargo:rustc-link-lib=hipblas");
            }
        } else {
            println!("cargo:warning=MLP build failed. Setting no_nn.");
            println!("cargo:rustc-cfg=no_nn");
        }
    } else {
        println!("cargo:rustc-cfg=no_nn");
    }

    for cfg in ["no_nn", "no_octree"] {
        println!("cargo:rustc-check-cfg=cfg({cfg})");
    }

    println!("cargo:rerun-if-env-changed=MLP_COMPRESSION_DIR");
    println!("cargo:rerun-if-changed=rust/external/tucker-octree/toctree.cpp");

    let ml_msg = if ml_compiled { "OK" } else { "FAILED (no_nn)" };
    let octree_msg = if octree_compiled {
        "OK"
    } else {
        "FAILED (no_octree)"
    };
    let zfp_lib_dir = zfp_dst.join("lib");
    let zfp_lib64_dir = zfp_dst.join("lib64");

    if zfp_lib64_dir.exists() {
        println!("cargo:rustc-link-search=native={}", zfp_lib64_dir.display());
    } else {
        println!("cargo:rustc-link-search=native={}", zfp_lib_dir.display());
    }
    println!("cargo:warning=Octree: {} | MLP: {}", octree_msg, ml_msg);
}

fn setup_octree_linking(octree_dst: &PathBuf, zfp_dst: &PathBuf) {
    let octree_lib_dir = octree_dst.join("lib");
    let zfp_lib_dir = zfp_dst.join("lib");
    let zfp_lib64_dir = zfp_dst.join("lib64");

    if zfp_lib64_dir.exists() {
        println!("cargo:rustc-link-search=native={}", zfp_lib64_dir.display());
    } else {
        println!("cargo:rustc-link-search=native={}", zfp_lib_dir.display());
    }

    println!(
        "cargo:rustc-link-search=native={}",
        octree_lib_dir.display()
    );
    println!("cargo:rustc-link-lib=static=toctree_compressor");
    println!("cargo:rustc-link-lib=static=zfp");
    println!("cargo:rustc-link-lib=stdc++");
    if env::var("CARGO_CFG_TARGET_OS").as_deref() == Ok("linux") {
        println!("cargo:rustc-link-arg=-lstdc++");
    }
}

fn is_program_in_path(program: &str) -> bool {
    Command::new("which")
        .arg(program)
        .status()
        .map(|s| s.success())
        .unwrap_or(false)
}