Method ProceedDataSet

catboost/cuda/methods/tree_ctrs.cpp:468–515

}

void NCatboostCuda::TTreeCtrDataSetsHelper::ProceedDataSet(ui32 dataSetPermutationKey,
                                                           NCatboostCuda::TTreeCtrDataSet& dataSet,
                                                           NCatboostCuda::TTreeCtrDataSetsHelper::TDataSetVisitor& visitor) {
    const ui32 deviceId = dataSet.GetDeviceId();
    auto ctrTargets = DeviceView(DataSet.GetCtrTargets(), deviceId);

    if (!dataSet.HasCompressedIndex()) {
        NCudaLib::GetCudaManager().WaitComplete();

        auto guard = NCudaLib::GetCudaManager().GetProfiler().Profile(
            TStringBuilder() << "Build  #" << dataSet.GetCtrs().size() << " ctrs dataset");

        using TTensorBuilder = TBatchFeatureTensorBuilder;
        const ui32 tensorBuilderStreams = PackSizeEstimators[deviceId]->GetStreamCountForCtrCalculation();

        TTreeCtrDataSetBuilder builder(DepthPermutations[dataSetPermutationKey].DeviceView(deviceId),
                                       dataSet,
                                       tensorBuilderStreams > 1,
                                       !LevelBasedCompressedIndex);

        NCudaLib::GetCudaManager().WaitComplete();

        TTensorBuilder batchFeatureTensorBuilder(FeaturesManager,
                                                 DataSet.GetCatFeatures(),
                                                 tensorBuilderStreams);

        TVector<ui32> catFeatureIds(dataSet.GetCatFeatures().begin(),
                                    dataSet.GetCatFeatures().end());
        TCtrFromTensorCalcer ctrFromTensorCalcer(builder,
                                                 dataSet.GetCtrConfigs(),
                                                 ctrTargets);

        TBatchFeatureTensorBuilder::TFeatureTensorVisitor ctrFromTensorCalcerFunc = [&](const TFeatureTensor& tensor, TCtrBinBuilder<NCudaLib::TSingleMapping>& builder) {
            ctrFromTensorCalcer(tensor, builder);
        };
        batchFeatureTensorBuilder.VisitCtrBinBuilders(dataSet.GetBaseTensorIndices(),
                                                      dataSet.GetBaseTensor(),
                                                      catFeatureIds,
                                                      ctrFromTensorCalcerFunc);
        NCudaLib::GetCudaManager().WaitComplete();
    }

    visitor(dataSet);

    if (NeedToDropDataSetAfterVisit(deviceId)) {
        TTreeCtrDataSetBuilder::DropCache(dataSet);
    }
}

bool NCatboostCuda::TTreeCtrDataSetsHelper::NeedToDropDataSetAfterVisit(ui32 deviceId) const {
    if (IsLastLevel()) {