yum/gpu_fft

git clone https://git.yummers.dev/yum/gpu_fft

yumAdd gpu DFT code69b45432026-09-12

use anyhow::Context;
use std::fs;
use std::fs::File;
use std::fs::OpenOptions;
use std::io::Write;
use std::path::Path;
use std::process::{Command, Output};

fn say(message: &str) {
    let file = OpenOptions::new().write(true).open("/dev/tty");
    if file.is_ok() {
        let _ = writeln!(file.unwrap(), "gpu_fft: {message}");
    }
}

fn run_cmd(cmd: &str, dir: &str) -> anyhow::Result<Output> {
    println!("cargo::warning=Running {cmd}...");
    let output = Command::new("sh")
        .arg("-c")
        .arg(cmd)
        .current_dir(dir)
        .output()
        .with_context(|| format!("failed to execute `{cmd}`"))?;

    if !output.status.success() {
        anyhow::bail!(
            "Command failed `{cmd}`:\nstdout:{}\nstderr:\n{}",
            String::from_utf8_lossy(&output.stdout),
            String::from_utf8_lossy(&output.stderr)
        );
    }

    Ok(output)
}

fn ensure_slang() -> Result<(), anyhow::Error> {
    let sentinel = "./slang/.did_build";
    if Path::new(sentinel).exists() {
        // Already built - short-circuit.
        return Ok(());
    }

    say(concat!(
        "The slang compiler is not present, or is missing the\n",
        "sentinel file. We need the slang compiler in order to translate\n",
        ".slang shaders into .wgsl, which wgpu (webGPU) uses. We're going to\n",
        "acquire the slang compiler, then build it. This will take a few\n",
        "minutes and use a lot of CPU. This only needs to run once."
    ));

    // Assume we're in some weird intermediate state and start fresh.
    if Path::new("./slang").exists() {
        say("Removing old slang compiler...");
        fs::remove_dir_all("./slang")?;
    }

    say("cargo::warning=Acquiring slang compiler...");
    run_cmd(
        "git clone --recursive https://github.com/shader-slang/slang",
        ".",
    )?;
    run_cmd("mkdir ./slang/build", ".")?;

    say("Building slang compiler...");
    run_cmd("cmake -B build --preset default", "./slang")?;
    run_cmd("cmake --build --preset releaseWithDebugInfo", "./slang")?;
    File::create(sentinel)?;
    Ok(())
}

// Convert slang shaders in shaders/ to wgsl. (I refuse to use the term "transpile" - it is
// redundant and ignores the fact that all compilation is merely a translation from one language to
// another.)
fn compile_shaders(slang_dir: &str, wgsl_dir: &str) -> Result<(), anyhow::Error> {
    let compiler_path = "./slang/build/RelWithDebInfo/bin/slangc";

    if Path::new(wgsl_dir).exists() {
        fs::remove_dir_all(wgsl_dir)?;
    }
    fs::create_dir_all(wgsl_dir)
        .with_context(|| format!("Failed to create output directory {wgsl_dir}"))?;

    for entry in fs::read_dir(slang_dir)? {
        let entry = entry?.path();
        let Some(ext) = entry.extension() else {
            // No extension.
            continue;
        };
        if ext != "slang" {
            continue;
        }
        let stem = entry
            .file_stem()
            .with_context(|| format!("Failed to get file stem for {}", entry.display()))?;
        let new_path = Path::new(wgsl_dir).join(stem).with_extension("wgsl");
        let slang_dir = entry.parent().ok_or(anyhow::anyhow!(
            "Failed to get parent directory for {}",
            entry.display()
        ))?;
        let cmd = format!(
            "{compiler_path} {} -o {} -no-mangle -I {}",
            entry.display(),
            new_path.display(),
            slang_dir.display(),
        );
        run_cmd(&cmd, ".")?;
    }

    Ok(())
}

fn main() -> anyhow::Result<()> {
    println!("cargo:rerun-if-changed=slang");
    ensure_slang()?;

    println!("cargo:rerun-if-changed=shaders");
    compile_shaders("./shaders", "./src/generated/shaders")?;

    Ok(())
}