Increase max alignment to 256.

This commit is contained in:
Antonio Sánchez 2025-02-04 20:06:28 +00:00
parent b1e74b1ccd
commit b73bb766a5

View File

@ -141,23 +141,24 @@ EIGEN_DEVICE_FUNC inline void throw_std_bad_alloc() {
*/
EIGEN_DEVICE_FUNC inline void* handmade_aligned_malloc(std::size_t size,
std::size_t alignment = EIGEN_DEFAULT_ALIGN_BYTES) {
eigen_assert(alignment >= sizeof(void*) && alignment <= 128 && (alignment & (alignment - 1)) == 0 &&
"Alignment must be at least sizeof(void*), less than or equal to 128, and a power of 2");
eigen_assert(alignment >= sizeof(void*) && alignment <= 256 && (alignment & (alignment - 1)) == 0 &&
"Alignment must be at least sizeof(void*), less than or equal to 256, and a power of 2");
check_that_malloc_is_allowed();
EIGEN_USING_STD(malloc)
void* original = malloc(size + alignment);
if (original == nullptr) return nullptr;
uint8_t offset = static_cast<uint8_t>(alignment - (reinterpret_cast<std::size_t>(original) & (alignment - 1)));
std::size_t offset = alignment - (reinterpret_cast<std::size_t>(original) & (alignment - 1));
void* aligned = static_cast<void*>(static_cast<uint8_t*>(original) + offset);
*(static_cast<uint8_t*>(aligned) - 1) = offset;
// Store offset - 1, since it is guaranteed to be at least 1.
*(static_cast<uint8_t*>(aligned) - 1) = static_cast<uint8_t>(offset - 1);
return aligned;
}
/** \internal Frees memory allocated with handmade_aligned_malloc */
EIGEN_DEVICE_FUNC inline void handmade_aligned_free(void* ptr) {
if (ptr != nullptr) {
uint8_t offset = static_cast<uint8_t>(*(static_cast<uint8_t*>(ptr) - 1));
std::size_t offset = static_cast<std::size_t>(*(static_cast<uint8_t*>(ptr) - 1)) + 1;
void* original = static_cast<void*>(static_cast<uint8_t*>(ptr) - offset);
check_that_malloc_is_allowed();
@ -174,7 +175,7 @@ EIGEN_DEVICE_FUNC inline void handmade_aligned_free(void* ptr) {
EIGEN_DEVICE_FUNC inline void* handmade_aligned_realloc(void* ptr, std::size_t new_size, std::size_t old_size,
std::size_t alignment = EIGEN_DEFAULT_ALIGN_BYTES) {
if (ptr == nullptr) return handmade_aligned_malloc(new_size, alignment);
uint8_t old_offset = *(static_cast<uint8_t*>(ptr) - 1);
std::size_t old_offset = static_cast<std::size_t>(*(static_cast<uint8_t*>(ptr) - 1)) + 1;
void* old_original = static_cast<uint8_t*>(ptr) - old_offset;
check_that_malloc_is_allowed();
@ -182,14 +183,15 @@ EIGEN_DEVICE_FUNC inline void* handmade_aligned_realloc(void* ptr, std::size_t n
void* original = realloc(old_original, new_size + alignment);
if (original == nullptr) return nullptr;
if (original == old_original) return ptr;
uint8_t offset = static_cast<uint8_t>(alignment - (reinterpret_cast<std::size_t>(original) & (alignment - 1)));
std::size_t offset = alignment - (reinterpret_cast<std::size_t>(original) & (alignment - 1));
void* aligned = static_cast<void*>(static_cast<uint8_t*>(original) + offset);
if (offset != old_offset) {
const void* src = static_cast<const void*>(static_cast<uint8_t*>(original) + old_offset);
std::size_t count = (std::min)(new_size, old_size);
std::memmove(aligned, src, count);
}
*(static_cast<uint8_t*>(aligned) - 1) = offset;
// Store offset - 1, since it is guaranteed to be at least 1.
*(static_cast<uint8_t*>(aligned) - 1) = static_cast<uint8_t>(offset - 1);
return aligned;
}