Ranking metric acceleration on the gpu (#5398)

dmlc · Mar 22, 2020 · d2231fc · d2231fc
1 parent cd7d6f7
commit d2231fc
Show file tree

Hide file tree

Showing 5 changed files with 747 additions and 14 deletions.
diff --git a/src/metric/metric.cc b/src/metric/metric.cc
@@ -81,5 +81,8 @@ namespace metric {
 DMLC_REGISTRY_LINK_TAG(elementwise_metric);
 DMLC_REGISTRY_LINK_TAG(multiclass_metric);
 DMLC_REGISTRY_LINK_TAG(rank_metric);
+#ifdef XGBOOST_USE_CUDA
+DMLC_REGISTRY_LINK_TAG(rank_metric_gpu);
+#endif
 }  // namespace metric
 }  // namespace xgboost
diff --git a/src/metric/rank_metric.cc b/src/metric/rank_metric.cc
@@ -260,7 +260,7 @@ struct EvalAuc : public Metric {
       !info.group_ptr_.empty() && info.weights_.Size() != info.num_row_;
 
     // Check if we have a GPU assignment; else, revert back to CPU
-    if (tparam_->gpu_id >= 0 && is_ranking_task) {
+    if (tparam_->gpu_id >= 0) {
       if (!auc_gpu_) {
         // Check and see if we have the GPU metric registered in the internal registry
         auc_gpu_.reset(GPUMetric::CreateGPUMetric(this->Name(), tparam_));