aligned.hpp
1 /*
2  * Copyright (c) 2020, 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 {
26 namespace detail {
27 
32 static constexpr std::size_t RMM_DEFAULT_HOST_ALIGNMENT{alignof(std::max_align_t)};
33 
38 constexpr bool is_pow2(std::size_t n) { return (0 == (n & (n - 1))); }
39 
44 constexpr bool is_supported_alignment(std::size_t alignment) { return is_pow2(alignment); }
45 
54 constexpr std::size_t align_up(std::size_t v, std::size_t align_bytes) noexcept
55 {
56  assert(is_supported_alignment(align_bytes));
57  return (v + (align_bytes - 1)) & ~(align_bytes - 1);
58 }
59 
68 constexpr std::size_t align_down(std::size_t v, std::size_t align_bytes) noexcept
69 {
70  assert(is_supported_alignment(align_bytes));
71  return v & ~(align_bytes - 1);
72 }
73 
82 constexpr bool is_aligned(std::size_t v, std::size_t align_bytes) noexcept
83 {
84  assert(is_supported_alignment(align_bytes));
85  return v == align_down(v, align_bytes);
86 }
87 
115 template <typename Alloc>
116 void *aligned_allocate(std::size_t bytes, std::size_t alignment, Alloc alloc)
117 {
118  assert(is_pow2(alignment));
119 
120  // allocate memory for bytes, plus potential alignment correction,
121  // plus store of the correction offset
122  std::size_t padded_allocation_size{bytes + alignment + sizeof(std::ptrdiff_t)};
123 
124  char *const original = static_cast<char *>(alloc(padded_allocation_size));
125 
126  // account for storage of offset immediately prior to the aligned pointer
127  void *aligned{original + sizeof(std::ptrdiff_t)};
128 
129  // std::align modifies `aligned` to point to the first aligned location
130  std::align(alignment, bytes, aligned, padded_allocation_size);
131 
132  // Compute the offset between the original and aligned pointers
133  std::ptrdiff_t offset = static_cast<char *>(aligned) - original;
134 
135  // Store the offset immediately before the aligned pointer
136  *(static_cast<std::ptrdiff_t *>(aligned) - 1) = offset;
137 
138  return aligned;
139 }
140 
156 template <typename Dealloc>
157 void aligned_deallocate(void *p, std::size_t bytes, std::size_t alignment, Dealloc dealloc)
158 {
159  (void)alignment;
160 
161  // Get offset from the location immediately prior to the aligned pointer
162  std::ptrdiff_t const offset = *(reinterpret_cast<std::ptrdiff_t *>(p) - 1);
163 
164  void *const original = static_cast<char *>(p) - offset;
165 
166  dealloc(original);
167 }
168 } // namespace detail
169 } // namespace rmm