Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
166 changes: 165 additions & 1 deletion build_utils/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
//! Build utilities shared across monarch *-sys crates
//!
//! This module provides common functionality for Python environment discovery
//! and CUDA installation detection used by various build scripts.
//! and CUDA/HIP installation detection used by various build scripts.

use std::env;
use std::path::Path;
Expand Down Expand Up @@ -48,6 +48,18 @@ for library_path in cpp_extension.library_paths():
print('LIBTORCH_CXX11:', torch._C._GLIBCXX_USE_CXX11_ABI)
";

/// Python script to extract PyTorch details including HIP info
pub const PYTHON_PRINT_HIP_DETAILS: &str = r"
import torch
from torch.utils import cpp_extension
print('HIP_HOME:', cpp_extension.HIP_HOME)
for include_path in cpp_extension.include_paths():
print('LIBTORCH_INCLUDE:', include_path)
for library_path in cpp_extension.library_paths():
print('LIBTORCH_LIB:', library_path)
print('LIBTORCH_CXX11:', torch._C._GLIBCXX_USE_CXX11_ABI)
";

/// Python script to extract Python include paths
pub const PYTHON_PRINT_INCLUDE_PATH: &str = r"
import sysconfig
Expand All @@ -64,6 +76,14 @@ pub struct CudaConfig {
pub lib_dirs: Vec<PathBuf>,
}

/// Configuration structure for HIP environment
#[derive(Debug, Clone, Default)]
pub struct HipConfig {
pub hip_home: Option<PathBuf>,
pub include_dirs: Vec<PathBuf>,
pub lib_dirs: Vec<PathBuf>,
}

/// Result of Python environment discovery
#[derive(Debug, Clone)]
pub struct PythonConfig {
Expand All @@ -75,6 +95,7 @@ pub struct PythonConfig {
#[derive(Debug)]
pub enum BuildError {
CudaNotFound,
HipNotFound,
PythonNotFound,
CommandFailed(String),
PathNotFound(String),
Expand All @@ -84,6 +105,7 @@ impl std::fmt::Display for BuildError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
BuildError::CudaNotFound => write!(f, "CUDA installation not found"),
BuildError::HipNotFound => write!(f, "HIP installation not found"),
BuildError::PythonNotFound => write!(f, "Python interpreter not found"),
BuildError::CommandFailed(cmd) => write!(f, "Command failed: {}", cmd),
BuildError::PathNotFound(path) => write!(f, "Path not found: {}", path),
Expand Down Expand Up @@ -140,6 +162,47 @@ pub fn find_cuda_home() -> Option<String> {
cuda_home
}

/// Find HIP home directory using various heuristics
///
/// This function attempts to locate HIP installation through:
/// 1. HIP_HOME environment variable
/// 2. HIP_PATH environment variable
/// 3. Finding hipcc in PATH and deriving hip home
/// 4. Platform-specific default locations
pub fn find_hip_home() -> Option<String> {
// Guess #1: Environment variables
let mut hip_home = env::var("HIP_HOME")
.ok()
.or_else(|| env::var("HIP_PATH").ok());

if hip_home.is_none() {
// Guess #2: Find hipcc in PATH
if let Ok(hipcc_path) = which("hipcc") {
// Get parent directory twice (hipcc is in HIP_HOME/bin)
if let Some(hip_dir) = hipcc_path.parent().and_then(|p| p.parent()) {
hip_home = Some(hip_dir.to_string_lossy().into_owned());
}
} else {
// Guess #3: Platform-specific defaults
if cfg!(windows) {
let pattern = r"C:\Program Files\NVIDIA GPU Computing Toolkit\HIP\v*.*";
let hip_homes: Vec<_> = glob(pattern).unwrap().filter_map(Result::ok).collect();
if !hip_homes.is_empty() {
hip_home = Some(hip_homes[0].to_string_lossy().into_owned());
}
} else {
// Unix-like systems
let hip_candidate = "/opt/rocm/hip";
if Path::new(hip_candidate).exists() {
hip_home = Some(hip_candidate.to_string());
}
}
}
}

hip_home
}

/// Discover CUDA configuration including home, include dirs, and lib dirs
pub fn discover_cuda_config() -> Result<CudaConfig, BuildError> {
let cuda_home = find_cuda_home().ok_or(BuildError::CudaNotFound)?;
Expand Down Expand Up @@ -169,6 +232,35 @@ pub fn discover_cuda_config() -> Result<CudaConfig, BuildError> {
Ok(config)
}

/// Discover HIP configuration including home, include dirs, and lib dirs
pub fn discover_hip_config() -> Result<HipConfig, BuildError> {
let hip_home = find_hip_home().ok_or(BuildError::HipNotFound)?;
let hip_home_path = PathBuf::from(&hip_home);

let mut config = HipConfig {
hip_home: Some(hip_home_path.clone()),
include_dirs: Vec::new(),
lib_dirs: Vec::new(),
};

// Add standard include directory
let include_dir = hip_home_path.join("include");
if include_dir.exists() {
config.include_dirs.push(include_dir);
}

// Add standard library directories
for lib_subdir in &["lib64", "lib", "lib/x64"] {
let lib_dir = hip_home_path.join(lib_subdir);
if lib_dir.exists() {
config.lib_dirs.push(lib_dir);
break; // Use first found
}
}

Ok(config)
}

/// Validate CUDA installation exists and is complete
pub fn validate_cuda_installation() -> Result<String, BuildError> {
let cuda_config = discover_cuda_config()?;
Expand All @@ -187,6 +279,24 @@ pub fn validate_cuda_installation() -> Result<String, BuildError> {
Ok(cuda_home_str)
}

/// Validate HIP installation exists and is complete
pub fn validate_hip_installation() -> Result<String, BuildError> {
let hip_config = discover_hip_config()?;
let hip_home = hip_config.hip_home.ok_or(BuildError::HipNotFound)?;
let hip_home_str = hip_home.to_string_lossy().to_string();

// Verify HIP include directory exists
let hip_include_path = hip_home.join("include");
if !hip_include_path.exists() {
return Err(BuildError::PathNotFound(format!(
"HIP include directory at {}",
hip_include_path.display()
)));
}

Ok(hip_home_str)
}

/// Get CUDA library directory
pub fn get_cuda_lib_dir() -> Result<String, BuildError> {
// Check if user explicitly set CUDA_LIB_DIR
Expand All @@ -212,6 +322,31 @@ pub fn get_cuda_lib_dir() -> Result<String, BuildError> {
))
}

/// Get HIP library directory
pub fn get_hip_lib_dir() -> Result<String, BuildError> {
// Check if user explicitly set HIP_LIB_DIR
if let Ok(hip_lib_dir) = env::var("HIP_LIB_DIR") {
return Ok(hip_lib_dir);
}

// Try to deduce from HIP configuration
let hip_config = discover_hip_config()?;
if let Some(hip_home) = hip_config.hip_home {
let lib64_path = hip_home.join("lib64");
if lib64_path.exists() {
return Ok(lib64_path.to_string_lossy().to_string());
}
let lib_path = hip_home.join("lib");
if lib_path.exists() {
return Ok(lib_path.to_string_lossy().to_string());
}
}

Err(BuildError::PathNotFound(
"HIP library directory".to_string(),
))
}

/// Discover Python environment directories using sysconfig
///
/// Returns tuple of (include_dir, lib_dir) as optional strings
Expand Down Expand Up @@ -264,6 +399,18 @@ pub fn print_cuda_error_help() {
eprintln!("Example: export CUDA_HOME=/usr/local/cuda-12.0");
}

/// Print helpful error message for HIP not found
pub fn print_hip_error_help() {
eprintln!("Error: HIP installation not found!");
eprintln!("Please ensure HIP is installed and one of the following is true:");
eprintln!(" 1. Set HIP_HOME environment variable to your HIP installation directory");
eprintln!(" 2. Set HIP_PATH environment variable to your HIP installation directory");
eprintln!(" 3. Ensure 'hipcc' is in your PATH");
eprintln!(" 4. Install HIP to the default location (/opt/rocm on Linux)");
eprintln!();
eprintln!("Example: export HIP_HOME=/opt/rocm/");
}

/// Print helpful error message for CUDA lib dir not found
pub fn print_cuda_lib_error_help() {
eprintln!("Error: CUDA library directory not found!");
Expand All @@ -273,6 +420,15 @@ pub fn print_cuda_lib_error_help() {
eprintln!("Or: export CUDA_LIB_DIR=/usr/lib64");
}

/// Print helpful error message for HIP lib dir not found
pub fn print_hip_lib_error_help() {
eprintln!("Error: HIP library directory not found!");
eprintln!("Please set HIP_LIB_DIR environment variable to your HIP library directory.");
eprintln!();
eprintln!("Example: export HIP_LIB_DIR=/opt/rocm/lib");
eprintln!("Or: export HIP_LIB_DIR=/usr/lib64");
}

#[cfg(test)]
mod tests {
use super::*;
Expand All @@ -285,6 +441,14 @@ mod tests {
assert_eq!(result, Some("/test/cuda".to_string()));
}

#[test]
fn test_find_hip_home_env_var() {
env::set_var("HIP_HOME", "/test/hip");
let result = find_hip_home();
env::remove_var("HIP_HOME");
assert_eq!(result, Some("/test/hip".to_string()));
}

#[test]
fn test_python_scripts_constants() {
assert!(PYTHON_PRINT_DIRS.contains("sysconfig"));
Expand Down
39 changes: 18 additions & 21 deletions hip-sys/build.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,15 +15,14 @@ fn main() {}

#[cfg(not(target_os = "macos"))]
fn main() {
// Validate CUDA installation and get CUDA home path
// let cuda_home = match build_utils::validate_cuda_installation() {
// Ok(home) => home,
// Err(_) => {
// build_utils::print_cuda_error_help();
// std::process::exit(1);
// }
// };
let hip_home = "/opt/rocm";
// Validate HIP installation and get HIP home path
let hip_home = match build_utils::validate_hip_installation() {
Ok(home) => home,
Err(_) => {
build_utils::print_hip_error_help();
std::process::exit(1);
}
};

// Start building the bindgen configuration
let mut builder = bindgen::Builder::default()
Expand Down Expand Up @@ -70,20 +69,18 @@ fn main() {
println!("cargo::metadata=LIB_PATH={}", lib_dir);
}

// Get CUDA library directory and emit link directives
// let cuda_lib_dir = match build_utils::get_cuda_lib_dir() {
// Ok(dir) => dir,
// Err(_) => {
// build_utils::print_cuda_lib_error_help();
// std::process::exit(1);
// }
// };
let hip_lib_dir ="/opt/rocm/lib";

// Get HIP library directory and emit link directives
let hip_lib_dir = match build_utils::get_hip_lib_dir() {
Ok(dir) => dir,
Err(_) => {
build_utils::print_hip_lib_error_help();
std::process::exit(1);
}
};
println!("cargo:rustc-link-search=native={}", hip_lib_dir);
// println!("cargo:rustc-link-lib=cuda");
// println!("cargo:rustc-link-lib=cudart");
println!("cargo:rustc-link-lib=hip");
println!("cargo:rustc-link-lib=hiprtc");
println!("{}", hip_lib_dir);

// Write the bindings to the $OUT_DIR/bindings.rs file
match env::var("OUT_DIR") {
Expand Down
36 changes: 16 additions & 20 deletions rdmaxcel-sys-hip/build.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,16 +25,14 @@ fn main() {
println!("cargo:rerun-if-changed=src/rdmaxcel.h");
println!("cargo:rerun-if-changed=src/rdmaxcel.c");

// TODO: hardcoding hip_home for now
// Validate CUDA installation and get CUDA home path
// let cuda_home = match build_utils::validate_cuda_installation() {
// Ok(home) => home,
// Err(_) => {
// build_utils::print_cuda_error_help();
// std::process::exit(1);
// }
// };
let hip_home = "/opt/rocm";
// Validate HIP installation and get HIP home path
let hip_home = match build_utils::validate_hip_installation() {
Ok(home) => home,
Err(_) => {
build_utils::print_hip_error_help();
std::process::exit(1);
}
};

// Get the directory of the current crate
let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap_or_else(|_| {
Expand Down Expand Up @@ -142,16 +140,14 @@ fn main() {
println!("cargo:metadata=LIB_PATH={}", lib_dir);
}

// TODO: hardcoding hip_lib_dir for now
// Get CUDA library directory and emit link directives
// let cuda_lib_dir = match build_utils::get_cuda_lib_dir() {
// Ok(dir) => dir,
// Err(_) => {
// build_utils::print_cuda_lib_error_help();
// std::process::exit(1);
// }
// };
let hip_lib_dir ="/opt/rocm/lib";
// Get HIP library directory and emit link directives
let hip_lib_dir = match build_utils::get_hip_lib_dir() {
Ok(dir) => dir,
Err(_) => {
build_utils::print_hip_lib_error_help();
std::process::exit(1);
}
};
println!("cargo:rustc-link-search=native={}", hip_lib_dir);

// TODO: confirm hip is enough to replace cuda/cudart
Expand Down
Loading