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

Author: Montana Low

Cargo.toml

@@ -16,3 +16,6 @@ derive_builder = "0.5"
 log = "0.4"
 tempfile = "3.0"
 indexmap = "1.0"
+
+[features]
+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,5 @@ readme = "README.md"
 libc = "0.2"
 
 [build-dependencies]
-bindgen = "0.59"
+bindgen = "0.61"
 cmake = "0.1"

xgboost-sys/build.rs

@@ -25,6 +25,9 @@ fn main() {
     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();
 
     let xgb_root = xgb_root.canonicalize().unwrap();
@@ -34,7 +37,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 +67,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");
+    }
 }