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 }