Function TestBatchReduce

catboost/cuda/cuda_lib/ut/test_batch_reduce.cpp:18–63

    const int tries = 100;

    void TestBatchReduce(int batchSize,
                         const size_t partSize = 128 * 4 * 9,
                         bool compress = false) {
        SetDefaultProfileMode(EProfileMode::ImplicitLabelSync);
        TSetLogging inThisScope(ELoggingLevel::Debug);
        auto& profiler = GetProfiler();

        for (int tr = 0; tr < tries; ++tr) {
            {
                auto beforeReduceMapping = TStripeMapping::RepeatOnAllDevices(partSize);
                auto afterReduceMapping = TStripeMapping::SplitBetweenDevices(partSize);

                TVector<TStripeBuffer<float>> data;
                TVector<TComputationStream> streams;
                for (int i = 0; i < batchSize; ++i) {
                    data.emplace_back(TStripeBuffer<float>::Create(beforeReduceMapping));
                    FillBuffer(data.back(), 1.0f * i);
                    if (streams.size() < 100) {
                        streams.push_back(GetCudaManager().RequestStream());
                    }
                }

                {
                    TString label = compress ? "ReduceCompressed" : "Reduce";
                    auto guard = profiler.Profile(label);
                    for (int i = 0; i < batchSize; ++i) {
                        auto& toReduce = data[i];
                        TReducer<TStripeBuffer<float>> reducer(streams[i % 100].GetId());
                        reducer(toReduce, afterReduceMapping, compress);
                    }
                }
            }
            {
                auto beforeReduceMapping = TStripeMapping::RepeatOnAllDevices(batchSize * partSize);
                auto afterReduceMapping = TStripeMapping::SplitBetweenDevices(batchSize * partSize);

                TStripeBuffer<float> data = TStripeBuffer<float>::Create(beforeReduceMapping);
                {
                    TString label = compress ? "ReduceAllBatchCompressed" : "ReduceBatch";
                    auto guard = profiler.Profile(label);
                    TReducer<TStripeBuffer<float>> reducer;
                    reducer(data, afterReduceMapping, compress);
                }
            }
        }
    }

    Y_UNIT_TEST(TestBatchReduce1) {
        auto stopCudaManagerGuard = StartCudaManager();