[Performance][CUDA] Labor UVA optimization (#5885)

Co-authored-by: Xin Yao <xiny@nvidia.com> Co-authored-by: Rhett Ying <85214957+Rhett-Ying@users.noreply.github.com>
dmlc · Aug 10, 2023 · 39b0a52 · 39b0a52
1 parent a7c8afb
commit 39b0a52
Show file tree

Hide file tree

Showing 5 changed files with 300 additions and 87 deletions.
diff --git a/src/array/cpu/labor_pick.h b/src/array/cpu/labor_pick.h
@@ -31,7 +31,6 @@
 #include <algorithm>
 #include <cmath>
 #include <functional>
-#include <iostream>
 #include <memory>
 #include <numeric>
 #include <string>
@@ -124,7 +123,10 @@ auto compute_importance_sampling_probabilities(
         var_1 = 0;
         if (weighted) {
           for (auto j = indptr[rid]; j < indptr[rid + 1]; j++)
-            var_1 += A[j] * A[j] / std::min(ONE, c * ps[j - indptr[rid]]);
+            // The check for zero is necessary for numerical stability
+            var_1 += A[j] > 0
+                         ? A[j] * A[j] / std::min(ONE, c * ps[j - indptr[rid]])
+                         : 0;
         } else {
           for (auto j = indptr[rid]; j < indptr[rid + 1]; j++)
             var_1 += ONE / std::min(ONE, c * ps[j - indptr[rid]]);
@@ -208,8 +210,9 @@ std::pair<COOMatrix, FloatArray> CSRLaborPick(
   IdArray picked_row = NDArray::Empty({hop_size}, vidtype, ctx);
   IdArray picked_col = NDArray::Empty({hop_size}, vidtype, ctx);
   IdArray picked_idx = NDArray::Empty({hop_size}, vidtype, ctx);
-  FloatArray picked_imp =
-      NDArray::Empty({importance_sampling ? hop_size : 0}, dtype, ctx);
+  FloatArray picked_imp = importance_sampling
+                              ? NDArray::Empty({hop_size}, dtype, ctx)
+                              : NullArray();
   IdxType* picked_rdata = picked_row.Ptr<IdxType>();
   IdxType* picked_cdata = picked_col.Ptr<IdxType>();
   IdxType* picked_idata = picked_idx.Ptr<IdxType>();