build.rs

fn main() {
    // Check if static-link feature is enabled
    let static_link = std::env::var("CARGO_FEATURE_STATIC_LINK").is_ok();
    let target_os = std::env::var("CARGO_CFG_TARGET_OS").unwrap_or_default();
    let target_family = std::env::var("CARGO_CFG_TARGET_FAMILY").unwrap_or_default();
    let is_windows = target_os == "windows";
    let is_unix = target_family == "unix";

    // Library directory - can be overridden via environment variable
    let default_lib_dir = if static_link {
        "3rdparty/tilelang/build-static/lib"
    } else {
        "3rdparty/tilelang/build/lib"
    };
    let lib_dir =
        std::env::var("FLASHTILE_TILELANG_LIB_DIR").unwrap_or_else(|_| default_lib_dir.into());

    println!("cargo:rerun-if-env-changed=FLASHTILE_TILELANG_LIB_DIR");
    println!("cargo:rerun-if-env-changed=CARGO_FEATURE_STATIC_LINK");
    println!("cargo:rerun-if-env-changed=FLASHTILE_LIBBACKTRACE_DIR");

    if static_link {
        // Static linking mode for binary distribution
        println!("cargo:rustc-link-search=native={}", lib_dir);

        // Link order matters for static initialization!
        // tvm_ffi_static must be linked first as it contains the FFI dispatch table
        // tvm_static contains the TVM runtime and compiler
        // tilelang_module_static contains TileLang extensions
        //
        // We use whole-archive to ensure all global constructors are included
        // (TVM uses static registration pattern)
        println!("cargo:rustc-link-lib=static:+whole-archive=tvm_ffi_static");
        println!("cargo:rustc-link-lib=static:+whole-archive=tvm_static");
        println!("cargo:rustc-link-lib=static:+whole-archive=tilelang_module_static");

        // Stack traces:
        // - Unix: tvm-ffi can build libbacktrace (via autoconf/configure).
        // - Windows: tvm-ffi uses DbgHelp, and libbacktrace isn't available.
        if is_windows {
            println!("cargo:rustc-link-lib=dylib=Dbghelp");
        } else {
            // libbacktrace is built by tvm-ffi and required for stack traces
            // Can be overridden via FLASHTILE_LIBBACKTRACE_DIR
            let libbacktrace_dir =
                std::env::var("FLASHTILE_LIBBACKTRACE_DIR").unwrap_or_else(|_| {
                    // Default: relative to lib_dir
                    let lib_path = std::path::Path::new(&lib_dir);
                    if let Some(parent) = lib_path.parent() {
                        parent
                            .join("tvm/3rdparty/tvm-ffi/libbacktrace/lib")
                            .to_string_lossy()
                            .to_string()
                    } else {
                        "3rdparty/tilelang/build-static/tvm/3rdparty/tvm-ffi/libbacktrace/lib"
                            .to_string()
                    }
                });
            println!("cargo:rustc-link-search=native={}", libbacktrace_dir);
            println!("cargo:rustc-link-lib=static=backtrace");
        }

        // System libraries required by TVM/TileLang
        if target_os == "linux" {
            println!("cargo:rustc-link-lib=dylib=pthread");
            println!("cargo:rustc-link-lib=dylib=dl");
            println!("cargo:rustc-link-lib=dylib=stdc++");
        } else if target_os == "macos" {
            println!("cargo:rustc-link-lib=dylib=c++");
        }

        // CUDA libraries (dynamic - CUDA remains dynamically linked)
        // Find CUDA path from environment or use default
        let cuda_path = std::env::var("CUDA_PATH")
            .or_else(|_| std::env::var("CUDA_HOME"))
            .unwrap_or_else(|_| "/usr/local/cuda".into());
        if is_windows {
            println!("cargo:rustc-link-search=native={}\\lib\\x64", cuda_path);
        } else {
            println!("cargo:rustc-link-search=native={}/lib64", cuda_path);
            println!("cargo:rustc-link-search=native={}/lib64/stubs", cuda_path);
        }

        // Link CUDA runtime and driver libraries
        println!("cargo:rustc-link-lib=dylib=cudart");
        println!("cargo:rustc-link-lib=dylib=cuda");
        println!("cargo:rustc-link-lib=dylib=nvrtc");

        // Tell tvm-ffi-sys to skip dynamic linking (we handle it here)
        println!("cargo:rustc-env=TVM_FFI_STATIC=1");
    } else {
        // Dynamic linking mode (default)
        //
        // Unix: set rpath so runtime can find libraries.
        // Windows: no rpath concept; rely on DLL search order / PATH.
        if is_unix {
            // NOTE: `tileiras` is invoked by cutile-python from arbitrary working
            // directories, so a relative RUNPATH like `3rdparty/...` will not work.
            // Use both:
            // - an absolute path (best effort), and
            // - an $ORIGIN-relative path so the binary is relocatable under the repo.
            let manifest_dir = std::path::PathBuf::from(
                std::env::var("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR is set by Cargo"),
            );
            let lib_dir_path = std::path::PathBuf::from(&lib_dir);
            let lib_dir_abs = if lib_dir_path.is_absolute() {
                lib_dir_path.clone()
            } else {
                manifest_dir.join(&lib_dir_path)
            };
            println!(
                "cargo:rustc-link-arg=-Wl,-rpath,{}",
                lib_dir_abs.to_string_lossy()
            );
            if !lib_dir_path.is_absolute() {
                // Binary is located at `target/<profile>/flashtile` so go back two
                // directories to repo root, then into `3rdparty/...`.
                let origin_rpath = format!("$ORIGIN/../../{}", lib_dir_path.to_string_lossy());
                println!("cargo:rustc-link-arg=-Wl,-rpath,{}", origin_rpath);
            }
        }
    }
}