fixed_size_memory_resource.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2020-2025, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once
 
#include <rmm/aligned.hpp>
#include <rmm/cuda_stream_view.hpp>
#include <rmm/detail/error.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/detail/logging_assert.hpp>
#include <rmm/detail/thrust_namespace.h>
#include <rmm/mr/device/detail/fixed_size_free_list.hpp>
#include <rmm/mr/device/detail/stream_ordered_memory_resource.hpp>
#include <rmm/resource_ref.hpp>
 
#include <cuda_runtime_api.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/transform_iterator.h>
 
#include <algorithm>
#include <cstddef>
#include <utility>
#include <vector>
 
namespace RMM_NAMESPACE {
namespace mr {
template <typename Upstream>
class fixed_size_memory_resource
  : public detail::stream_ordered_memory_resource<fixed_size_memory_resource<Upstream>,
                                                  detail::fixed_size_free_list> {
 public:
  friend class detail::stream_ordered_memory_resource<fixed_size_memory_resource<Upstream>,
                                                      detail::fixed_size_free_list>;
 
  static constexpr std::size_t default_block_size = 1 << 20;  
 
  static constexpr std::size_t default_blocks_to_preallocate = 128;
 
  explicit fixed_size_memory_resource(
    device_async_resource_ref upstream_mr,
    // NOLINTNEXTLINE bugprone-easily-swappable-parameters
    std::size_t block_size            = default_block_size,
    std::size_t blocks_to_preallocate = default_blocks_to_preallocate)
    : upstream_mr_{upstream_mr},
      block_size_{align_up(block_size, CUDA_ALLOCATION_ALIGNMENT)},
      upstream_chunk_size_{block_size_ * blocks_to_preallocate}
  {
    // allocate initial blocks and insert into free list
    this->insert_blocks(std::move(blocks_from_upstream(cuda_stream_legacy)), cuda_stream_legacy);
  }
 
  explicit fixed_size_memory_resource(
    Upstream* upstream_mr,
    // NOLINTNEXTLINE bugprone-easily-swappable-parameters
    std::size_t block_size            = default_block_size,
    std::size_t blocks_to_preallocate = default_blocks_to_preallocate)
    : upstream_mr_{to_device_async_resource_ref_checked(upstream_mr)},
      block_size_{align_up(block_size, CUDA_ALLOCATION_ALIGNMENT)},
      upstream_chunk_size_{block_size_ * blocks_to_preallocate}
  {
    // allocate initial blocks and insert into free list
    this->insert_blocks(std::move(blocks_from_upstream(cuda_stream_legacy)), cuda_stream_legacy);
  }
 
  ~fixed_size_memory_resource() override { release(); }
 
  fixed_size_memory_resource()                                             = delete;
  fixed_size_memory_resource(fixed_size_memory_resource const&)            = delete;
  fixed_size_memory_resource(fixed_size_memory_resource&&)                 = delete;
  fixed_size_memory_resource& operator=(fixed_size_memory_resource const&) = delete;
  fixed_size_memory_resource& operator=(fixed_size_memory_resource&&)      = delete;
 
  [[nodiscard]] device_async_resource_ref get_upstream_resource() const noexcept
  {
    return upstream_mr_;
  }
 
  [[nodiscard]] std::size_t get_block_size() const noexcept { return block_size_; }
 
 protected:
  using free_list  = detail::fixed_size_free_list;  
  using block_type = free_list::block_type;         
  using typename detail::stream_ordered_memory_resource<fixed_size_memory_resource<Upstream>,
                                                        detail::fixed_size_free_list>::split_block;
  using lock_guard = std::lock_guard<std::mutex>;  
 
  [[nodiscard]] std::size_t get_maximum_allocation_size() const { return get_block_size(); }
 
  block_type expand_pool(std::size_t size, free_list& blocks, cuda_stream_view stream)
  {
    blocks.insert(std::move(blocks_from_upstream(stream)));
    return blocks.get_block(size);
  }
 
  free_list blocks_from_upstream(cuda_stream_view stream)
  {
    void* ptr = get_upstream_resource().allocate_async(upstream_chunk_size_, stream);
    block_type block{ptr};
    upstream_blocks_.push_back(block);
 
    auto num_blocks = upstream_chunk_size_ / block_size_;
 
    auto block_gen = [ptr, this](int index) {
      // NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-pointer-arithmetic)
      return block_type{static_cast<char*>(ptr) + index * block_size_};
    };
    auto first =
      thrust::make_transform_iterator(thrust::make_counting_iterator(std::size_t{0}), block_gen);
    return free_list(first, first + num_blocks);
  }
 
  split_block allocate_from_block(block_type const& block, std::size_t size)
  {
    return {block, block_type{nullptr}};
  }
 
  block_type free_block(void* ptr, std::size_t size) noexcept
  {
    // Deallocating a fixed-size block just inserts it in the free list, which is
    // handled by the parent class
    RMM_LOGGING_ASSERT(align_up(size, CUDA_ALLOCATION_ALIGNMENT) <= block_size_);
    return block_type{ptr};
  }
 
  void release()
  {
    lock_guard lock(this->get_mutex());
 
    for (auto block : upstream_blocks_) {
      get_upstream_resource().deallocate(block.pointer(), upstream_chunk_size_);
    }
    upstream_blocks_.clear();
  }
 
#ifdef RMM_DEBUG_PRINT
  void print()
  {
    lock_guard lock(this->get_mutex());
 
    auto const [free, total] = rmm::available_device_memory();
    std::cout << "GPU free memory: " << free << " total: " << total << "\n";
 
    std::cout << "upstream_blocks: " << upstream_blocks_.size() << "\n";
    std::size_t upstream_total{0};
 
    for (auto blocks : upstream_blocks_) {
      blocks.print();
      upstream_total += upstream_chunk_size_;
    }
    std::cout << "total upstream: " << upstream_total << " B\n";
 
    this->print_free_blocks();
  }
#endif
 
  std::pair<std::size_t, std::size_t> free_list_summary(free_list const& blocks)
  {
    return blocks.is_empty() ? std::make_pair(std::size_t{0}, std::size_t{0})
                             : std::make_pair(block_size_, blocks.size() * block_size_);
  }
 
 private:
  device_async_resource_ref upstream_mr_;  // The resource from which to allocate new blocks
 
  std::size_t block_size_;           // size of blocks this MR allocates
  std::size_t upstream_chunk_size_;  // size of chunks allocated from heap MR
 
  // blocks allocated from heap: so they can be easily freed
  std::vector<block_type> upstream_blocks_;
};
  // end of group
}  // namespace mr
}  // namespace RMM_NAMESPACE