Use the aligned_resource_adaptor to allocate bloom filter device buffers

cpp/src/io/parquet/bloom_filter_reader.cu

@@ -32,6 +32,7 @@
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_buffer.hpp>
 #include <rmm/exec_policy.hpp>
+#include <rmm/mr/device/aligned_resource_adaptor.hpp>
 
 #include <cuco/bloom_filter_policies.cuh>
 #include <cuco/bloom_filter_ref.cuh>
@@ -378,14 +379,16 @@ class bloom_filter_expression_converter : public equality_literals_collector {
  * @param bloom_filter_sizes Bloom filter sizes for all chunks
  * @param chunk_source_map Association between each column chunk and its source
  * @param stream CUDA stream used for device memory operations and kernel launches
+ * @param aligned_mr Aligned device memory resource to allocate bloom filter buffers
  */
 void read_bloom_filter_data(host_span<std::unique_ptr<datasource> const> sources,
                             size_t num_chunks,
                             cudf::host_span<rmm::device_buffer> bloom_filter_data,
                             cudf::host_span<std::optional<int64_t>> bloom_filter_offsets,
                             cudf::host_span<std::optional<int32_t>> bloom_filter_sizes,
                             std::vector<size_type> const& chunk_source_map,
-                            rmm::cuda_stream_view stream)
+                            rmm::cuda_stream_view stream,
+                            rmm::device_async_resource_ref aligned_mr)
 {
   // Read tasks for bloom filter data
   std::vector<std::future<size_t>> read_tasks;
@@ -445,15 +448,15 @@ void read_bloom_filter_data(host_span<std::unique_ptr<datasource> const> sources
 
       // Check if we already read in the filter bitset in the initial read.
       if (initial_read_size >= bloom_filter_header_size + bitset_size) {
-        bloom_filter_data[chunk] =
-          rmm::device_buffer{buffer->data() + bloom_filter_header_size, bitset_size, stream};
+        bloom_filter_data[chunk] = rmm::device_buffer{
+          buffer->data() + bloom_filter_header_size, bitset_size, stream, aligned_mr};
       }
       // Read the bitset from datasource.
       else {
         auto const bitset_offset = bloom_filter_offset + bloom_filter_header_size;
         // Directly read to device if preferred
         if (source->is_device_read_preferred(bitset_size)) {
-          bloom_filter_data[chunk] = rmm::device_buffer{bitset_size, stream};
+          bloom_filter_data[chunk] = rmm::device_buffer{bitset_size, stream, aligned_mr};
           auto future_read_size =
             source->device_read_async(bitset_offset,
                                       bitset_size,
@@ -462,8 +465,9 @@ void read_bloom_filter_data(host_span<std::unique_ptr<datasource> const> sources
 
           read_tasks.emplace_back(std::move(future_read_size));
         } else {
-          buffer                   = source->host_read(bitset_offset, bitset_size);
-          bloom_filter_data[chunk] = rmm::device_buffer{buffer->data(), buffer->size(), stream};
+          buffer = source->host_read(bitset_offset, bitset_size);
+          bloom_filter_data[chunk] =
+            rmm::device_buffer{buffer->data(), buffer->size(), stream, aligned_mr};
         }
       }
     });
@@ -481,7 +485,8 @@ std::vector<rmm::device_buffer> aggregate_reader_metadata::read_bloom_filters(
   host_span<std::vector<size_type> const> row_group_indices,
   host_span<int const> column_schemas,
   size_type total_row_groups,
-  rmm::cuda_stream_view stream) const
+  rmm::cuda_stream_view stream,
+  rmm::device_async_resource_ref aligned_mr) const
 {
   // Descriptors for all the chunks that make up the selected columns
   auto const num_input_columns = column_schemas.size();
@@ -540,7 +545,8 @@ std::vector<rmm::device_buffer> aggregate_reader_metadata::read_bloom_filters(
                          bloom_filter_offsets,
                          bloom_filter_sizes,
                          chunk_source_map,
-                         stream);
+                         stream,
+                         aligned_mr);
 
   // Return bloom filter data
   return bloom_filter_data;
@@ -609,10 +615,16 @@ std::optional<std::vector<std::vector<size_type>>> aggregate_reader_metadata::ap
   // Return early if no column with equality predicate(s)
   if (equality_col_schemas.empty()) { return std::nullopt; }
 
+  using policy_type = cuco::arrow_filter_policy<cuda::std::byte, cudf::hashing::detail::XXHash_64>;
+  size_t constexpr alignment =
+    policy_type::words_per_block * sizeof(typename policy_type::word_type);
+  auto aligned_mr =
+    rmm::mr::aligned_resource_adaptor(cudf::get_current_device_resource(), alignment);
+
   // Read a vector of bloom filter bitset device buffers for all columns with equality
   // predicate(s) across all row groups
   auto bloom_filter_data = read_bloom_filters(
-    sources, input_row_group_indices, equality_col_schemas, total_row_groups, stream);
+    sources, input_row_group_indices, equality_col_schemas, total_row_groups, stream, aligned_mr);
 
   // No bloom filter buffers, return the original row group indices
   if (bloom_filter_data.empty()) { return std::nullopt; }

cpp/src/io/parquet/reader_impl_helpers.hpp

@@ -204,6 +204,7 @@ class aggregate_reader_metadata {
    * @param[out] bloom_filter_data List of bloom filter data device buffers
    * @param column_schemas Schema indices of columns whose bloom filters will be read
    * @param stream CUDA stream used for device memory operations and kernel launches
+   * @param aligned_mr Aligned device memory resource to allocate bloom filter buffers
    *
    * @return A flattened list of bloom filter bitset device buffers for each predicate column across
    * row group
@@ -213,7 +214,8 @@ class aggregate_reader_metadata {
     host_span<std::vector<size_type> const> row_group_indices,
     host_span<int const> column_schemas,
     size_type num_row_groups,
-    rmm::cuda_stream_view stream) const;
+    rmm::cuda_stream_view stream,
+    rmm::device_async_resource_ref aligned_mr) const;
 
   /**
    * @brief Collects Parquet types for the columns with the specified schema indices