aligned.hpp
1 /*
2  * Copyright (c) 2020-2021, NVIDIA CORPORATION.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #pragma once
18 
19 #include <cassert>
20 #include <cstddef>
21 #include <cstdint>
22 #include <memory>
23 #include <new>
24 
25 namespace rmm::detail {
26 
31 static constexpr std::size_t RMM_DEFAULT_HOST_ALIGNMENT{alignof(std::max_align_t)};
32 
37 static constexpr std::size_t CUDA_ALLOCATION_ALIGNMENT{256};
38 
43 constexpr bool is_pow2(std::size_t value) { return (0 == (value & (value - 1))); }
44 
49 constexpr bool is_supported_alignment(std::size_t alignment) { return is_pow2(alignment); }
50 
59 constexpr std::size_t align_up(std::size_t value, std::size_t alignment) noexcept
60 {
61  assert(is_supported_alignment(alignment));
62  return (value + (alignment - 1)) & ~(alignment - 1);
63 }
64 
73 constexpr std::size_t align_down(std::size_t value, std::size_t alignment) noexcept
74 {
75  assert(is_supported_alignment(alignment));
76  return value & ~(alignment - 1);
77 }
78 
87 constexpr bool is_aligned(std::size_t value, std::size_t alignment) noexcept
88 {
89  assert(is_supported_alignment(alignment));
90  return value == align_down(value, alignment);
91 }
92 
93 inline bool is_pointer_aligned(void* ptr, std::size_t alignment = CUDA_ALLOCATION_ALIGNMENT)
94 {
95  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-reinterpret-cast)
96  return rmm::detail::is_aligned(reinterpret_cast<ptrdiff_t>(ptr), alignment);
97 }
98 
126 template <typename Alloc>
127 void* aligned_allocate(std::size_t bytes, std::size_t alignment, Alloc alloc)
128 {
129  assert(is_pow2(alignment));
130 
131  // allocate memory for bytes, plus potential alignment correction,
132  // plus store of the correction offset
133  std::size_t padded_allocation_size{bytes + alignment + sizeof(std::ptrdiff_t)};
134 
135  char* const original = static_cast<char*>(alloc(padded_allocation_size));
136 
137  // account for storage of offset immediately prior to the aligned pointer
138  // NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-pointer-arithmetic)
139  void* aligned{original + sizeof(std::ptrdiff_t)};
140 
141  // std::align modifies `aligned` to point to the first aligned location
142  std::align(alignment, bytes, aligned, padded_allocation_size);
143 
144  // Compute the offset between the original and aligned pointers
145  std::ptrdiff_t offset = static_cast<char*>(aligned) - original;
146 
147  // Store the offset immediately before the aligned pointer
148  // NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-pointer-arithmetic)
149  *(static_cast<std::ptrdiff_t*>(aligned) - 1) = offset;
150 
151  return aligned;
152 }
153 
169 template <typename Dealloc>
170 // NOLINTNEXTLINE(bugprone-easily-swappable-parameters)
171 void aligned_deallocate(void* ptr, std::size_t bytes, std::size_t alignment, Dealloc dealloc)
172 {
173  (void)alignment;
174 
175  // Get offset from the location immediately prior to the aligned pointer
176  // NOLINTNEXTLINE
177  std::ptrdiff_t const offset = *(reinterpret_cast<std::ptrdiff_t*>(ptr) - 1);
178 
179  // NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-pointer-arithmetic)
180  void* const original = static_cast<char*>(ptr) - offset;
181 
182  dealloc(original);
183 }
184 } // namespace rmm::detail