Add nvbench environment class for initializing RMM in benchmarks

cpp/CMakeLists.txt

@@ -788,8 +788,7 @@ if(CUDF_BUILD_BENCHMARKS)
   # do not support the latest version of fmt, which is automatically pulled into our conda
   # environments by mamba.
   set(CPM_DOWNLOAD_fmt TRUE)
-  include(${rapids-cmake-dir}/cpm/nvbench.cmake)
-  rapids_cpm_nvbench()
+  include(cmake/thirdparty/get_nvbench.cmake)
   add_subdirectory(benchmarks)
 endif()
 

cpp/benchmarks/fixture/rmm_pool_raii.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -69,4 +69,15 @@ class rmm_pool_raii {
   std::shared_ptr<rmm::mr::device_memory_resource> mr;
 };
 
+/**
+ * Base fixture for cudf benchmarks using nvbench.
+ *
+ * Initializes the default memory resource to use the RMM pool device resource.
+ */
+struct nvbench_base_fixture {
+  rmm_pool_raii _mr;
+};
+
 }  // namespace cudf
+
+#define NVBENCH_ENVIRONMENT cudf::nvbench_base_fixture

cpp/benchmarks/groupby/group_max.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -24,7 +24,6 @@
 template <typename Type>
 void bench_groupby_max(nvbench::state& state, nvbench::type_list<Type>)
 {
-  cudf::rmm_pool_raii pool_raii;
   const auto size = static_cast<cudf::size_type>(state.get_int64("num_rows"));
 
   auto const keys = [&] {

cpp/benchmarks/groupby/group_nunique.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -40,7 +40,6 @@ auto make_aggregation_request_vector(cudf::column_view const& values, Args&&...
 template <typename Type>
 void bench_groupby_nunique(nvbench::state& state, nvbench::type_list<Type>)
 {
-  cudf::rmm_pool_raii pool_raii;
   const auto size = static_cast<cudf::size_type>(state.get_int64("num_rows"));
 
   auto const keys = [&] {

cpp/benchmarks/groupby/group_rank.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -30,7 +30,6 @@ static void nvbench_groupby_rank(nvbench::state& state,
 {
   using namespace cudf;
   constexpr auto dtype = type_to_id<int64_t>();
-  cudf::rmm_pool_raii pool_raii;
 
   bool const is_sorted              = state.get_int64("is_sorted");
   cudf::size_type const column_size = state.get_int64("data_size");

cpp/benchmarks/groupby/group_struct_keys.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -29,8 +29,6 @@
 
 void bench_groupby_struct_keys(nvbench::state& state)
 {
-  cudf::rmm_pool_raii pool_raii;
-
   using Type           = int;
   using column_wrapper = cudf::test::fixed_width_column_wrapper<Type>;
   std::default_random_engine generator;

cpp/benchmarks/io/csv/csv_reader_input.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -67,8 +67,6 @@ void csv_read_common(DataType const& data_types,
 template <data_type DataType>
 void BM_csv_read_input(nvbench::state& state, nvbench::type_list<nvbench::enum_type<DataType>>)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   auto const d_type      = get_type_or_group(static_cast<int32_t>(DataType));
   auto const source_type = io_type::FILEPATH;
 
@@ -78,8 +76,6 @@ void BM_csv_read_input(nvbench::state& state, nvbench::type_list<nvbench::enum_t
 template <cudf::io::io_type IO>
 void BM_csv_read_io(nvbench::state& state, nvbench::type_list<nvbench::enum_type<IO>>)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   auto const d_type      = get_type_or_group({static_cast<int32_t>(data_type::INTEGRAL),
                                          static_cast<int32_t>(data_type::FLOAT),
                                          static_cast<int32_t>(data_type::DECIMAL),

cpp/benchmarks/io/csv/csv_reader_options.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -32,8 +32,6 @@ void BM_csv_read_varying_options(
   nvbench::state& state,
   nvbench::type_list<nvbench::enum_type<ColSelection>, nvbench::enum_type<RowSelection>>)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   auto const data_types =
     dtypes_for_column_selection(get_type_or_group({static_cast<int32_t>(data_type::INTEGRAL),
                                                    static_cast<int32_t>(data_type::FLOAT),

cpp/benchmarks/io/fst.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -78,9 +78,6 @@ constexpr std::size_t single_item = 1;
 
 void BM_FST_JSON(nvbench::state& state)
 {
-  // TODO: to be replaced by nvbench fixture once it's ready
-  cudf::rmm_pool_raii rmm_pool;
-
   CUDF_EXPECTS(state.get_int64("string_size") <= std::numeric_limits<size_type>::max(),
                "Benchmarks only support up to size_type's maximum number of items");
   auto const string_size{size_type(state.get_int64("string_size"))};
@@ -116,9 +113,6 @@ void BM_FST_JSON(nvbench::state& state)
 
 void BM_FST_JSON_no_outidx(nvbench::state& state)
 {
-  // TODO: to be replaced by nvbench fixture once it's ready
-  cudf::rmm_pool_raii rmm_pool;
-
   CUDF_EXPECTS(state.get_int64("string_size") <= std::numeric_limits<size_type>::max(),
                "Benchmarks only support up to size_type's maximum number of items");
   auto const string_size{size_type(state.get_int64("string_size"))};
@@ -154,9 +148,6 @@ void BM_FST_JSON_no_outidx(nvbench::state& state)
 
 void BM_FST_JSON_no_out(nvbench::state& state)
 {
-  // TODO: to be replaced by nvbench fixture once it's ready
-  cudf::rmm_pool_raii rmm_pool;
-
   CUDF_EXPECTS(state.get_int64("string_size") <= std::numeric_limits<size_type>::max(),
                "Benchmarks only support up to size_type's maximum number of items");
   auto const string_size{size_type(state.get_int64("string_size"))};
@@ -190,9 +181,6 @@ void BM_FST_JSON_no_out(nvbench::state& state)
 
 void BM_FST_JSON_no_str(nvbench::state& state)
 {
-  // TODO: to be replaced by nvbench fixture once it's ready
-  cudf::rmm_pool_raii rmm_pool;
-
   CUDF_EXPECTS(state.get_int64("string_size") <= std::numeric_limits<size_type>::max(),
                "Benchmarks only support up to size_type's maximum number of items");
   auto const string_size{size_type(state.get_int64("string_size"))};

cpp/benchmarks/io/json/nested_json.cpp

@@ -157,9 +157,6 @@ auto make_test_json_data(cudf::size_type string_size, rmm::cuda_stream_view stre
 
 void BM_NESTED_JSON(nvbench::state& state)
 {
-  // TODO: to be replaced by nvbench fixture once it's ready
-  cudf::rmm_pool_raii rmm_pool;
-
   auto const string_size{cudf::size_type(state.get_int64("string_size"))};
   auto const default_options = cudf::io::json_reader_options{};
 
@@ -189,9 +186,6 @@ NVBENCH_BENCH(BM_NESTED_JSON)
 
 void BM_NESTED_JSON_DEPTH(nvbench::state& state)
 {
-  // TODO: to be replaced by nvbench fixture once it's ready
-  cudf::rmm_pool_raii rmm_pool;
-
   auto const string_size{cudf::size_type(state.get_int64("string_size"))};
   auto const depth{cudf::size_type(state.get_int64("depth"))};
 

cpp/benchmarks/io/orc/orc_reader_input.cpp

@@ -60,8 +60,6 @@ void orc_read_common(cudf::io::orc_writer_options const& opts,
 template <data_type DataType>
 void BM_orc_read_data(nvbench::state& state, nvbench::type_list<nvbench::enum_type<DataType>>)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   auto const d_type                 = get_type_or_group(static_cast<int32_t>(DataType));
   cudf::size_type const cardinality = state.get_int64("cardinality");
   cudf::size_type const run_length  = state.get_int64("run_length");
@@ -84,8 +82,6 @@ void BM_orc_read_io_compression(
   nvbench::state& state,
   nvbench::type_list<nvbench::enum_type<IO>, nvbench::enum_type<Compression>>)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   auto const d_type = get_type_or_group({static_cast<int32_t>(data_type::INTEGRAL_SIGNED),
                                          static_cast<int32_t>(data_type::FLOAT),
                                          static_cast<int32_t>(data_type::DECIMAL),

cpp/benchmarks/io/orc/orc_reader_options.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -56,8 +56,6 @@ void BM_orc_read_varying_options(nvbench::state& state,
                                                     nvbench::enum_type<UsesNumpyDType>,
                                                     nvbench::enum_type<Timestamp>>)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   auto const num_chunks = RowSelection == row_selection::ALL ? 1 : chunked_read_num_chunks;
 
   auto const use_index     = UsesIndex == uses_index::YES;

cpp/benchmarks/io/orc/orc_writer.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -46,8 +46,6 @@ constexpr cudf::size_type num_cols = 64;
 template <data_type DataType>
 void BM_orc_write_encode(nvbench::state& state, nvbench::type_list<nvbench::enum_type<DataType>>)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   auto const d_type                 = get_type_or_group(static_cast<int32_t>(DataType));
   cudf::size_type const cardinality = state.get_int64("cardinality");
   cudf::size_type const run_length  = state.get_int64("run_length");
@@ -90,8 +88,6 @@ void BM_orc_write_io_compression(
   nvbench::state& state,
   nvbench::type_list<nvbench::enum_type<IO>, nvbench::enum_type<Compression>>)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   auto const d_type = get_type_or_group({static_cast<int32_t>(data_type::INTEGRAL_SIGNED),
                                          static_cast<int32_t>(data_type::FLOAT),
                                          static_cast<int32_t>(data_type::DECIMAL),
@@ -141,8 +137,6 @@ void BM_orc_write_statistics(
   nvbench::state& state,
   nvbench::type_list<nvbench::enum_type<Statistics>, nvbench::enum_type<Compression>>)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   auto const d_type = get_type_or_group({static_cast<int32_t>(data_type::INTEGRAL_SIGNED),
                                          static_cast<int32_t>(data_type::FLOAT),
                                          static_cast<int32_t>(data_type::DECIMAL),

cpp/benchmarks/io/orc/orc_writer_chunks.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -35,8 +35,6 @@ constexpr int64_t data_size = 512 << 20;
 
 void nvbench_orc_write(nvbench::state& state)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   cudf::size_type num_cols = state.get_int64("num_columns");
 
   auto tbl = create_random_table(
@@ -79,8 +77,6 @@ void nvbench_orc_write(nvbench::state& state)
 
 void nvbench_orc_chunked_write(nvbench::state& state)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   cudf::size_type num_cols   = state.get_int64("num_columns");
   cudf::size_type num_tables = state.get_int64("num_chunks");
 

cpp/benchmarks/io/parquet/parquet_reader_input.cpp

@@ -60,8 +60,6 @@ void parquet_read_common(cudf::io::parquet_writer_options const& write_opts,
 template <data_type DataType>
 void BM_parquet_read_data(nvbench::state& state, nvbench::type_list<nvbench::enum_type<DataType>>)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   auto const d_type                 = get_type_or_group(static_cast<int32_t>(DataType));
   cudf::size_type const cardinality = state.get_int64("cardinality");
   cudf::size_type const run_length  = state.get_int64("run_length");
@@ -87,8 +85,6 @@ void BM_parquet_read_io_compression(
   nvbench::state& state,
   nvbench::type_list<nvbench::enum_type<IO>, nvbench::enum_type<Compression>>)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   auto const d_type = get_type_or_group({static_cast<int32_t>(data_type::INTEGRAL),
                                          static_cast<int32_t>(data_type::FLOAT),
                                          static_cast<int32_t>(data_type::DECIMAL),

cpp/benchmarks/io/parquet/parquet_reader_options.cpp

@@ -57,8 +57,6 @@ void BM_parquet_read_options(nvbench::state& state,
                                                 nvbench::enum_type<UsesPandasMetadata>,
                                                 nvbench::enum_type<Timestamp>>)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   auto constexpr str_to_categories = ConvertsStrings == converts_strings::YES;
   auto constexpr uses_pd_metadata  = UsesPandasMetadata == uses_pandas_metadata::YES;
 

cpp/benchmarks/io/parquet/parquet_writer.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -46,8 +46,6 @@ constexpr cudf::size_type num_cols = 64;
 template <data_type DataType>
 void BM_parq_write_encode(nvbench::state& state, nvbench::type_list<nvbench::enum_type<DataType>>)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   auto const data_types             = get_type_or_group(static_cast<int32_t>(DataType));
   cudf::size_type const cardinality = state.get_int64("cardinality");
   cudf::size_type const run_length  = state.get_int64("run_length");
@@ -90,8 +88,6 @@ void BM_parq_write_io_compression(
   nvbench::state& state,
   nvbench::type_list<nvbench::enum_type<IO>, nvbench::enum_type<Compression>>)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   auto const data_types = get_type_or_group({static_cast<int32_t>(data_type::INTEGRAL),
                                              static_cast<int32_t>(data_type::FLOAT),
                                              static_cast<int32_t>(data_type::DECIMAL),

cpp/benchmarks/io/parquet/parquet_writer_chunks.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -33,8 +33,6 @@ constexpr int64_t data_size = 512 << 20;
 
 void PQ_write(nvbench::state& state)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   cudf::size_type const num_cols = state.get_int64("num_cols");
 
   auto const tbl  = create_random_table(cycle_dtypes({cudf::type_id::INT32}, num_cols),
@@ -67,8 +65,6 @@ void PQ_write(nvbench::state& state)
 
 void PQ_write_chunked(nvbench::state& state)
 {
-  cudf::rmm_pool_raii rmm_pool;
-
   cudf::size_type const num_cols   = state.get_int64("num_cols");
   cudf::size_type const num_tables = state.get_int64("num_chunks");
 

cpp/benchmarks/io/text/multibyte_split.cpp

@@ -116,8 +116,6 @@ template <data_chunk_source_type source_type>
 static void bench_multibyte_split(nvbench::state& state,
                                   nvbench::type_list<nvbench::enum_type<source_type>>)
 {
-  cudf::rmm_pool_raii pool_raii;
-
   auto const delim_size         = state.get_int64("delim_size");
   auto const delim_percent      = state.get_int64("delim_percent");
   auto const file_size_approx   = state.get_int64("size_approx");