Merge pull request #2 from postgresml/montana/upgrades

add cuda to the build, upgrade bindgen version and fix some tests

Author: montanalow

Cargo.toml

@@ -12,7 +12,10 @@ readme = "README.md"
 [dependencies]
 xgboost-sys = { path = "xgboost-sys" }
 libc = "0.2"
-derive_builder = "0.5"
+derive_builder = "0.11"
 log = "0.4"
 tempfile = "3.0"
 indexmap = "1.0"
+
+[features]
+cuda = ["xgboost-sys/cuda"]

src/booster.rs

@@ -774,12 +774,12 @@ mod tests {
         }
 
         let train_metrics = booster.evaluate(&dmat_train).unwrap();
-        assert_eq!(*train_metrics.get("logloss").unwrap(), 0.006634);
-        assert_eq!(*train_metrics.get("map@4-").unwrap(), 0.001274);
+        assert_eq!(*train_metrics.get("logloss").unwrap(), 0.006634271);
+        assert_eq!(*train_metrics.get("map@4-").unwrap(), 0.0012738854);
 
         let test_metrics = booster.evaluate(&dmat_test).unwrap();
-        assert_eq!(*test_metrics.get("logloss").unwrap(), 0.00692);
-        assert_eq!(*test_metrics.get("map@4-").unwrap(), 0.005155);
+        assert_eq!(*test_metrics.get("logloss").unwrap(), 0.006919953);
+        assert_eq!(*test_metrics.get("map@4-").unwrap(), 0.005154639);
 
         let v = booster.predict(&dmat_test).unwrap();
         assert_eq!(v.len(), dmat_test.num_rows());

src/dmatrix.rs

@@ -128,7 +128,6 @@ impl DMatrix {
     pub fn from_csr(indptr: &[usize], indices: &[usize], data: &[f32], num_cols: Option<usize>) -> XGBResult<Self> {
         assert_eq!(indices.len(), data.len());
         let mut handle = ptr::null_mut();
-        let indptr: Vec<u64> = indptr.iter().map(|x| *x as u64).collect();
         let indices: Vec<u32> = indices.iter().map(|x| *x as u32).collect();
         let num_cols = num_cols.unwrap_or(0); // infer from data if 0
         xgb_call!(xgboost_sys::XGDMatrixCreateFromCSREx(indptr.as_ptr(),
@@ -152,7 +151,6 @@ impl DMatrix {
     pub fn from_csc(indptr: &[usize], indices: &[usize], data: &[f32], num_rows: Option<usize>) -> XGBResult<Self> {
         assert_eq!(indices.len(), data.len());
         let mut handle = ptr::null_mut();
-        let indptr: Vec<u64> = indptr.iter().map(|x| *x as u64).collect();
         let indices: Vec<u32> = indices.iter().map(|x| *x as u32).collect();
         let num_rows = num_rows.unwrap_or(0); // infer from data if 0
         xgb_call!(xgboost_sys::XGDMatrixCreateFromCSCEx(indptr.as_ptr(),
@@ -349,7 +347,7 @@ mod tests {
 
     #[test]
     fn read_num_cols() {
-        assert_eq!(read_train_matrix().unwrap().num_cols(), 126);
+        assert_eq!(read_train_matrix().unwrap().num_cols(), 127);
     }
 
     #[test]
@@ -380,7 +378,7 @@ mod tests {
     #[test]
     fn get_set_weights() {
         let mut dmat = read_train_matrix().unwrap();
-        assert_eq!(dmat.get_weights().unwrap(), &[]);
+        assert!(dmat.get_weights().unwrap().is_empty());
 
         let weight = [1.0, 10.0, 44.9555];
         assert!(dmat.set_weights(&weight).is_ok());
@@ -390,17 +388,20 @@ mod tests {
     #[test]
     fn get_set_base_margin() {
         let mut dmat = read_train_matrix().unwrap();
-        assert_eq!(dmat.get_base_margin().unwrap(), &[]);
+        assert!(dmat.get_base_margin().unwrap().is_empty());
 
         let base_margin = [0.00001, 0.000002, 1.23];
+        println!("rows: {:?}, {:?}", dmat.num_rows(), base_margin.len());
+        let result = dmat.set_base_margin(&base_margin);
+        println!("{:?}", result);
         assert!(dmat.set_base_margin(&base_margin).is_ok());
         assert_eq!(dmat.get_base_margin().unwrap(), base_margin);
     }
 
     #[test]
     fn get_set_group() {
         let mut dmat = read_train_matrix().unwrap();
-        assert_eq!(dmat.get_group().unwrap(), &[]);
+        assert!(dmat.get_group().unwrap().is_empty());
 
         let group = [1];
         assert!(dmat.set_group(&group).is_ok());

xgboost-sys/Cargo.toml

@@ -13,5 +13,8 @@ readme = "README.md"
 libc = "0.2"
 
 [build-dependencies]
-bindgen = "0.59"
+bindgen = "0.61"
 cmake = "0.1"
+
+[features]
+cuda = []

xgboost-sys/build.rs

@@ -22,6 +22,16 @@ fn main() {
     }
 
     // CMake
+    #[cfg(feature = "cuda")]
+    let dst = Config::new(&xgb_root)
+        .uses_cxx11()
+        .define("BUILD_STATIC_LIB", "ON")
+        .define("USE_CUDA", "ON")
+        .define("BUILD_WITH_CUDA", "ON")
+        .define("BUILD_WITH_CUDA_CUB", "ON")
+        .build();
+
+    #[cfg(not(feature = "cuda"))]
     let dst = Config::new(&xgb_root)
         .uses_cxx11()
         .define("BUILD_STATIC_LIB", "ON")
@@ -34,7 +44,11 @@ fn main() {
         .clang_args(&["-x", "c++", "-std=c++11"])
         .clang_arg(format!("-I{}", xgb_root.join("include").display()))
         .clang_arg(format!("-I{}", xgb_root.join("rabit/include").display()))
-        .clang_arg(format!("-I{}", xgb_root.join("dmlc-core/include").display()))
+        .clang_arg(format!("-I{}", xgb_root.join("dmlc-core/include").display()));
+
+    #[cfg(feature = "cuda")]
+    let bindings = bindings.clang_arg("-I/usr/local/cuda/include");
+    let bindings = bindings
         .generate()
         .expect("Unable to generate bindings.");
 
@@ -60,4 +74,12 @@ fn main() {
     println!("cargo:rustc-link-search=native={}", dst.join("lib").display());
     println!("cargo:rustc-link-lib=static=dmlc");
     println!("cargo:rustc-link-lib=static=xgboost");
+
+    #[cfg(feature = "cuda")]
+    {
+        println!("cargo:rustc-link-search={}", "/usr/local/cuda/lib64");
+        println!("cargo:rustc-link-search={}", "/usr/local/cuda/lib64/stubs");
+        println!("cargo:rustc-link-lib=dylib=cuda");
+        println!("cargo:rustc-link-lib=dylib=cudart");
+    }
 }

xgboost-sys/src/lib.rs

@@ -26,7 +26,7 @@ mod tests {
         let mut num_cols = 0;
         let ret_val = unsafe { XGDMatrixNumCol(handle, &mut num_cols) };
         assert_eq!(ret_val, 0);
-        assert_eq!(num_cols, 127);
+        assert_eq!(num_cols, 126);
 
         let ret_val = unsafe { XGDMatrixFree(handle) };
         assert_eq!(ret_val, 0);