restore.

dmlc · Jul 9, 2024 · 478c608 · 478c608
1 parent f21ab04
commit 478c608
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 11 deletions.
diff --git a/src/common/hist_util.cc b/src/common/hist_util.cc
@@ -188,14 +188,14 @@ class GHistBuildingManager {
 
 template <bool do_prefetch, class BuildingManager>
 void RowsWiseBuildHistKernel(Span<GradientPair const> gpair,
-                             common::Span<bst_idx_t const> row_indices,
-                             const GHistIndexMatrix &gmat, GHistRow hist) {
+                             RowSetCollection::Elem const row_indices, const GHistIndexMatrix &gmat,
+                             GHistRow hist) {
   constexpr bool kAnyMissing = BuildingManager::kAnyMissing;
   constexpr bool kFirstPage = BuildingManager::kFirstPage;
   using BinIdxType = typename BuildingManager::BinIdxType;
 
-  const size_t size = row_indices.size();
-  bst_idx_t const *rid = row_indices.data();
+  const size_t size = row_indices.Size();
+  bst_idx_t const *rid = row_indices.begin();
   auto const *p_gpair = reinterpret_cast<const float *>(gpair.data());
   const BinIdxType *gradient_index = gmat.index.data<BinIdxType>();
 
@@ -312,7 +312,7 @@ void ColsWiseBuildHistKernel(Span<GradientPair const> gpair,
 }
 
 template <class BuildingManager>
-void BuildHistDispatch(Span<GradientPair const> gpair, const RowSetCollection::Elem row_indices,
+void BuildHistDispatch(Span<GradientPair const> gpair, RowSetCollection::Elem row_indices,
                        const GHistIndexMatrix &gmat, GHistRow hist) {
   if (BuildingManager::kReadByColumn) {
     ColsWiseBuildHistKernel<BuildingManager>(gpair, row_indices, gmat, hist);
@@ -325,11 +325,12 @@ void BuildHistDispatch(Span<GradientPair const> gpair, const RowSetCollection::E
 
     if (contiguousBlock) {
       // contiguous memory access, built-in HW prefetching is enough
-      RowsWiseBuildHistKernel<false, BuildingManager>(
-          gpair, Span{row_indices.begin(), row_indices.Size()}, gmat, hist);
+      RowsWiseBuildHistKernel<false, BuildingManager>(gpair, row_indices, gmat, hist);
     } else {
-      auto span1 = Span{row_indices.begin(), row_indices.end() - no_prefetch_size};
-      auto span2 = Span{row_indices.end() - no_prefetch_size, row_indices.end()};
+      RowSetCollection::Elem const span1{row_indices.begin(), row_indices.end() - no_prefetch_size};
+      RowSetCollection::Elem const span2{row_indices.end() - no_prefetch_size, row_indices.end()};
+      // auto span1 = Span{row_indices.begin(), row_indices.end() - no_prefetch_size};
+      // auto span2 = Span{row_indices.end() - no_prefetch_size, row_indices.end()};
 
       RowsWiseBuildHistKernel<true, BuildingManager>(gpair, span1, gmat, hist);
       // no prefetching to avoid loading extra memory

diff --git a/src/common/partition_builder.h b/src/common/partition_builder.h
@@ -112,7 +112,7 @@ class PartitionBuilder {
                  const common::Range1d range, const bst_bin_t split_cond,
                  GHistIndexMatrix const& gmat, const common::ColumnMatrix& column_matrix,
                  const RegTree& tree, bst_idx_t const* rid) {
-    common::Span<bst_idx_t const> rid_span(rid + range.begin(), rid + range.end());
+    common::Span<bst_idx_t const> rid_span{rid + range.begin(), rid + range.end()};
     common::Span<bst_idx_t> left = GetLeftBuffer(node_in_set, range.begin(), range.end());
     common::Span<bst_idx_t> right = GetRightBuffer(node_in_set, range.begin(), range.end());
     std::size_t nid = nodes[node_in_set].nid;
@@ -183,7 +183,7 @@ class PartitionBuilder {
   }
 
   template <bool any_missing, typename ColumnType, typename Predicate>
-  void MaskKernel(ColumnType* p_column, common::Span<const bst_idx_t> row_indices,
+  void MaskKernel(ColumnType* p_column, common::Span<bst_idx_t const> row_indices,
                   bst_idx_t base_rowid, BitVector* decision_bits, BitVector* missing_bits,
                   Predicate&& pred) {
     auto& column = *p_column;