(4/N) Add Cargo support to cuda-sys and rdmacore-sys (#548)

Summary:
Pull Request resolved: #548

This diff adds Cargo support for `cuda-sys` and `rdmacore-sys`, which will be necessary for the open source build.

Note - I found that `arc autocargo` does not support `rust_bindgen_library` so these Cargo.tomls were created manually.

Reviewed By: dstaay-fb

Differential Revision: D78414608

fbshipit-source-id: 41c91978cec09ec2ce9e9a57e39281bea878b61c

Author: allenwang28

Cargo.toml

@@ -2,6 +2,7 @@
 resolver = "2"
 members = [
     "controller",
+    "cuda-sys",
     "hyper",
     "hyperactor",
     "hyperactor_macros",
@@ -12,5 +13,6 @@ members = [
     "monarch_extension",
     "monarch_tensor_worker",
     "nccl-sys",
+    "rdmacore-sys",
     "torch-sys",
 ]

cuda-sys/build.rs

@@ -0,0 +1,104 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use std::env;
+use std::path::Path;
+use std::path::PathBuf;
+
+use glob::glob;
+use which::which;
+
+// Translated from torch/utils/cpp_extension.py
+fn find_cuda_home() -> Option<String> {
+    // Guess #1
+    let mut cuda_home = env::var("CUDA_HOME")
+        .ok()
+        .or_else(|| env::var("CUDA_PATH").ok());
+
+    if cuda_home.is_none() {
+        // Guess #2
+        if let Ok(nvcc_path) = which("nvcc") {
+            // Get parent directory twice
+            if let Some(cuda_dir) = nvcc_path.parent().and_then(|p| p.parent()) {
+                cuda_home = Some(cuda_dir.to_string_lossy().into_owned());
+            }
+        } else {
+            // Guess #3
+            if cfg!(windows) {
+                // Windows code
+                let pattern = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v*.*";
+                let cuda_homes: Vec<_> = glob(pattern).unwrap().filter_map(Result::ok).collect();
+                if !cuda_homes.is_empty() {
+                    cuda_home = Some(cuda_homes[0].to_string_lossy().into_owned());
+                } else {
+                    cuda_home = None;
+                }
+            } else {
+                // Not Windows
+                let cuda_candidate = "/usr/local/cuda";
+                if Path::new(cuda_candidate).exists() {
+                    cuda_home = Some(cuda_candidate.to_string());
+                } else {
+                    cuda_home = None;
+                }
+            }
+        }
+    }
+    cuda_home
+}
+
+fn main() {
+    let cuda_home = find_cuda_home().expect("Could not find CUDA installation");
+
+    // Tell cargo to look for shared libraries in the CUDA directory
+    println!("cargo:rustc-link-search={}/lib64", cuda_home);
+    println!("cargo:rustc-link-search={}/lib", cuda_home);
+
+    // Link against the CUDA libraries
+    println!("cargo:rustc-link-lib=cuda");
+    println!("cargo:rustc-link-lib=cudart");
+
+    // Tell cargo to invalidate the built crate whenever the wrapper changes
+    println!("cargo:rerun-if-changed=src/wrapper.h");
+
+    // Add cargo metadata
+    println!("cargo:rustc-cfg=cargo");
+    println!("cargo:rustc-check-cfg=cfg(cargo)");
+
+    // The bindgen::Builder is the main entry point to bindgen
+    let bindings = bindgen::Builder::default()
+        // The input header we would like to generate bindings for
+        .header("src/wrapper.h")
+        // Add the CUDA include directory
+        .clang_arg(format!("-I{}/include", cuda_home))
+        // Parse as C++
+        .clang_arg("-x")
+        .clang_arg("c++")
+        .clang_arg("-std=gnu++20")
+        // Allow the specified functions and types
+        .allowlist_function("cu.*")
+        .allowlist_function("CU.*")
+        .allowlist_type("cu.*")
+        .allowlist_type("CU.*")
+        // Use newtype enum style
+        .default_enum_style(bindgen::EnumVariation::NewType {
+            is_bitfield: false,
+            is_global: false,
+        })
+        // Finish the builder and generate the bindings
+        .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
+        .generate()
+        // Unwrap the Result and panic on failure
+        .expect("Unable to generate bindings");
+
+    // Write the bindings to the $OUT_DIR/bindings.rs file
+    let out_path = PathBuf::from(env::var("OUT_DIR").unwrap());
+    bindings
+        .write_to_file(out_path.join("bindings.rs"))
+        .expect("Couldn't write bindings!");
+}

cuda-sys/src/lib.rs

@@ -5,3 +5,49 @@
  * This source code is licensed under the BSD-style license found in the
  * LICENSE file in the root directory of this source tree.
  */
+
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use cxx::ExternType;
+use cxx::type_id;
+
+/// SAFETY: bindings
+unsafe impl ExternType for CUstream_st {
+    type Id = type_id!("CUstream_st");
+    type Kind = cxx::kind::Opaque;
+}
+
+// When building with cargo, this is actually the lib.rs file for a crate.
+// Include the generated bindings.rs and suppress lints.
+#[allow(non_camel_case_types)]
+#[allow(non_upper_case_globals)]
+#[allow(non_snake_case)]
+mod inner {
+    #[cfg(cargo)]
+    include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
+}
+
+pub use inner::*;
+
+#[cfg(test)]
+mod tests {
+    use std::mem::MaybeUninit;
+
+    use super::*;
+
+    #[test]
+    fn sanity() {
+        // SAFETY: testing bindings
+        unsafe {
+            let mut version = MaybeUninit::<i32>::uninit();
+            let result = cuDriverGetVersion(version.as_mut_ptr());
+            assert_eq!(result, cudaError_enum(0));
+        }
+    }
+}

monarch_rdma/Cargo.toml

@@ -10,8 +10,10 @@ license = "BSD-3-Clause"
 [dependencies]
 anyhow = "1.0.98"
 async-trait = "0.1.86"
+cuda-sys = { path = "../cuda-sys" }
 hyperactor = { version = "0.0.0", path = "../hyperactor" }
 rand = { version = "0.8", features = ["small_rng"] }
+rdmacore-sys = { path = "../rdmacore-sys" }
 serde = { version = "1.0.185", features = ["derive", "rc"] }
 tracing = { version = "0.1.41", features = ["attributes", "valuable"] }
 

rdmacore-sys/build.rs

@@ -0,0 +1,69 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use std::env;
+use std::path::PathBuf;
+
+fn main() {
+    // Tell cargo to look for shared libraries in the specified directory
+    println!("cargo:rustc-link-search=/usr/lib");
+    println!("cargo:rustc-link-search=/usr/lib64");
+
+    // Link against the ibverbs library
+    println!("cargo:rustc-link-lib=ibverbs");
+
+    // Link against the mlx5 library
+    println!("cargo:rustc-link-lib=mlx5");
+
+    // Tell cargo to invalidate the built crate whenever the wrapper changes
+    println!("cargo:rerun-if-changed=src/wrapper.h");
+
+    // Add cargo metadata
+    println!("cargo:rustc-cfg=cargo");
+    println!("cargo:rustc-check-cfg=cfg(cargo)");
+
+    // The bindgen::Builder is the main entry point to bindgen
+    let bindings = bindgen::Builder::default()
+        // The input header we would like to generate bindings for
+        .header("src/wrapper.h")
+        // Allow the specified functions, types, and variables
+        .allowlist_function("ibv_.*")
+        .allowlist_function("mlx5dv_.*")
+        .allowlist_function("mlx5_wqe_.*")
+        .allowlist_type("ibv_.*")
+        .allowlist_type("mlx5dv_.*")
+        .allowlist_type("mlx5_wqe_.*")
+        .allowlist_var("MLX5_.*")
+        // Block specific types that are manually defined in lib.rs
+        .blocklist_type("ibv_wc")
+        .blocklist_type("mlx5_wqe_ctrl_seg")
+        // Apply the same bindgen flags as in the BUCK file
+        .bitfield_enum("ibv_access_flags")
+        .bitfield_enum("ibv_qp_attr_mask")
+        .bitfield_enum("ibv_wc_flags")
+        .bitfield_enum("ibv_send_flags")
+        .bitfield_enum("ibv_port_cap_flags")
+        .constified_enum_module("ibv_qp_type")
+        .constified_enum_module("ibv_qp_state")
+        .constified_enum_module("ibv_port_state")
+        .constified_enum_module("ibv_wc_opcode")
+        .constified_enum_module("ibv_wr_opcode")
+        .constified_enum_module("ibv_wc_status")
+        .derive_default(true)
+        .prepend_enum_name(false)
+        // Finish the builder and generate the bindings
+        .generate()
+        // Unwrap the Result and panic on failure
+        .expect("Unable to generate bindings");
+
+    // Write the bindings to the $OUT_DIR/bindings.rs file
+    let out_path = PathBuf::from(env::var("OUT_DIR").unwrap());
+    bindings
+        .write_to_file(out_path.join("bindings.rs"))
+        .expect("Couldn't write bindings!");
+}