From a552afb88c6496814326651555bb9dde2a87028d Mon Sep 17 00:00:00 2001 From: Karthikeyan <6488848+karthikeyann@users.noreply.github.com> Date: Tue, 25 Jan 2022 11:10:13 +0530 Subject: [PATCH] Limit benchmark iterations using environment variable (#10060) To address part of https://github.com/rapidsai/cudf/issues/5773 This allows to run benchmarks for only specific iterations using environment variable `CUDF_BENCHMARK_ITERATIONS`. except when benchmark definition itself specifies iteration count. Also, makes pool as static to allocate pool memory resource only once per binary. Authors: - Karthikeyan (https://github.com/karthikeyann) Approvers: - Conor Hoekstra (https://github.com/codereport) - Vukasin Milovanovic (https://github.com/vuule) - Nghia Truong (https://github.com/ttnghia) URL: https://github.com/rapidsai/cudf/pull/10060 --- cpp/benchmarks/fixture/benchmark_fixture.hpp | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/cpp/benchmarks/fixture/benchmark_fixture.hpp b/cpp/benchmarks/fixture/benchmark_fixture.hpp index 83f79bd68c5..ca3a748ccad 100644 --- a/cpp/benchmarks/fixture/benchmark_fixture.hpp +++ b/cpp/benchmarks/fixture/benchmark_fixture.hpp @@ -29,9 +29,11 @@ namespace { // memory resource factory helpers inline auto make_cuda() { return std::make_shared(); } -inline auto make_pool() +inline auto make_pool_instance() { - return rmm::mr::make_owning_wrapper(make_cuda()); + static rmm::mr::cuda_memory_resource cuda_mr; + static rmm::mr::pool_memory_resource pool_mr{&cuda_mr}; + return std::shared_ptr(&pool_mr); } } // namespace @@ -68,9 +70,15 @@ inline auto make_pool() */ class benchmark : public ::benchmark::Fixture { public: + benchmark() : ::benchmark::Fixture() + { + const char* env_iterations = std::getenv("CUDF_BENCHMARK_ITERATIONS"); + if (env_iterations != nullptr) { this->Iterations(std::max(0L, atol(env_iterations))); } + } + void SetUp(const ::benchmark::State& state) override { - mr = make_pool(); + mr = make_pool_instance(); rmm::mr::set_current_device_resource(mr.get()); // set default resource to pool }