Function GetDataShape

src/data/quantile_dmatrix.cc:87–147

}

void GetDataShape(Context const* ctx, DMatrixProxy* proxy,
                  DataIterProxy<DataIterResetCallback, XGDMatrixCallbackNext>* iter, float missing,
                  ExternalDataInfo* p_info) {
  auto& info = *p_info;

  auto const is_valid = data::IsValidFunctor{missing};
  auto nnz_cnt = [&]() {
    return DispatchAny(proxy, [&](auto const& value) {
      bst_idx_t n_threads = ctx->Threads();
      bst_idx_t n_features = info.column_sizes.size();
      linalg::Tensor<bst_idx_t, 2> column_sizes_tloc({n_threads, n_features}, DeviceOrd::CPU());
      column_sizes_tloc.Data()->Fill(0ul);
      auto view = column_sizes_tloc.HostView();
      common::ParallelFor(value.Size(), n_threads, common::Sched::Static(256), [&](auto i) {
        auto const& line = value.GetLine(i);
        for (bst_idx_t j = 0; j < line.Size(); ++j) {
          data::COOTuple const& elem = line.GetElement(j);
          if (is_valid(elem)) {
            view(omp_get_thread_num(), elem.column_idx)++;
          }
        }
      });
      auto ptr = column_sizes_tloc.Data()->HostPointer();
      auto result = std::accumulate(ptr, ptr + column_sizes_tloc.Size(), static_cast<bst_idx_t>(0));
      for (bst_idx_t tidx = 0; tidx < n_threads; ++tidx) {
        for (bst_idx_t fidx = 0; fidx < n_features; ++fidx) {
          info.column_sizes[fidx] += view(tidx, fidx);
        }
      }
      return result;
    });
  };

  /**
   * CPU impl needs an additional loop for accumulating the column size.
   */
  do {
    // We use do while here as the first batch is fetched in ctor
    if (info.n_features == 0) {
      info.n_features = BatchColumns(proxy);
      collective::SafeColl(collective::Allreduce(ctx, &info.n_features, collective::Op::kMax));
      info.column_sizes.clear();
      info.column_sizes.resize(info.n_features, 0);
      p_info->cats =
          std::make_shared<CatContainer>(cpu_impl::BatchCats(proxy), BatchCatsIsRef(proxy));
    } else {
      CHECK_EQ(info.n_features, BatchColumns(proxy)) << "Inconsistent number of columns.";
      auto cats = cpu_impl::BatchCats(proxy);
      CHECK_EQ(cats.n_total_cats, p_info->cats->NumCatsTotal()) << error::InconsistentCategories();
    }
    bst_idx_t batch_size = BatchSamples(proxy);
    info.batch_nnz.push_back(nnz_cnt());
    info.base_rowids.push_back(batch_size);
    info.nnz += info.batch_nnz.back();
    info.accumulated_rows += batch_size;
    info.n_batches++;
  } while (iter->Next());
  iter->Reset();