diff --git a/cpp/include/cudf_test/base_fixture.hpp b/cpp/include/cudf_test/base_fixture.hpp index eee2d396d9f..5fa07fd5568 100644 --- a/cpp/include/cudf_test/base_fixture.hpp +++ b/cpp/include/cudf_test/base_fixture.hpp @@ -24,7 +24,9 @@ #include #include +#include #include +#include #include #include #include @@ -217,6 +219,8 @@ class TempDirTestEnvironment : public ::testing::Environment { /// MR factory functions inline auto make_cuda() { return std::make_shared(); } +inline auto make_async() { return std::make_shared(); } + inline auto make_managed() { return std::make_shared(); } inline auto make_pool() @@ -224,6 +228,11 @@ inline auto make_pool() return rmm::mr::make_owning_wrapper(make_cuda()); } +inline auto make_arena() +{ + return rmm::mr::make_owning_wrapper(make_cuda()); +} + inline auto make_binning() { auto pool = make_pool(); @@ -253,7 +262,9 @@ inline std::shared_ptr create_memory_resource( { if (allocation_mode == "binning") return make_binning(); if (allocation_mode == "cuda") return make_cuda(); + if (allocation_mode == "async") return make_async(); if (allocation_mode == "pool") return make_pool(); + if (allocation_mode == "arena") return make_arena(); if (allocation_mode == "managed") return make_managed(); CUDF_FAIL("Invalid RMM allocation mode: " + allocation_mode); }