diff --git a/libc/src/__support/GPU/allocator.cpp b/libc/src/__support/GPU/allocator.cpp index 5ea27a9c44b66..b2f2953e4f285 100644 --- a/libc/src/__support/GPU/allocator.cpp +++ b/libc/src/__support/GPU/allocator.cpp @@ -22,6 +22,7 @@ #include "src/__support/GPU/utils.h" #include "src/__support/RPC/rpc_client.h" #include "src/__support/threads/sleep.h" +#include "src/string/memory_utils/inline_memcpy.h" namespace LIBC_NAMESPACE_DECL { @@ -550,5 +551,26 @@ void deallocate(void *ptr) { release_slab(slab); } +void *reallocate(void *ptr, uint64_t size) { + if (ptr == nullptr) + return gpu::allocate(size); + + // Non-slab allocations are considered foreign pointers so we fail. + if ((reinterpret_cast(ptr) & SLAB_ALIGNMENT) == 0) + return nullptr; + + // The original slab pointer is the 2MiB boundary using the given pointer. + Slab *slab = cpp::launder(reinterpret_cast( + (reinterpret_cast(ptr) & ~SLAB_ALIGNMENT))); + if (slab->get_chunk_size() >= size) + return ptr; + + // If we need a new chunk we reallocate and copy it over. + void *new_ptr = gpu::allocate(size); + inline_memcpy(new_ptr, ptr, slab->get_chunk_size()); + gpu::deallocate(ptr); + return new_ptr; +} + } // namespace gpu } // namespace LIBC_NAMESPACE_DECL diff --git a/libc/src/__support/GPU/allocator.h b/libc/src/__support/GPU/allocator.h index 466009aa71981..757f3a406015b 100644 --- a/libc/src/__support/GPU/allocator.h +++ b/libc/src/__support/GPU/allocator.h @@ -17,6 +17,7 @@ namespace gpu { void *allocate(uint64_t size); void deallocate(void *ptr); +void *reallocate(void *ptr, uint64_t size); } // namespace gpu } // namespace LIBC_NAMESPACE_DECL diff --git a/libc/src/stdlib/gpu/realloc.cpp b/libc/src/stdlib/gpu/realloc.cpp index 4fd4d6b278179..97ad1b3bbeff2 100644 --- a/libc/src/stdlib/gpu/realloc.cpp +++ b/libc/src/stdlib/gpu/realloc.cpp @@ -16,17 +16,15 @@ namespace LIBC_NAMESPACE_DECL { LLVM_LIBC_FUNCTION(void *, realloc, (void *ptr, size_t size)) { - if (ptr == nullptr) - return gpu::allocate(size); - - void *newmem = gpu::allocate(size); - if (newmem == nullptr) - return nullptr; - - // This will copy garbage if it goes beyond the old allocation size. - inline_memcpy(newmem, ptr, size); - gpu::deallocate(ptr); - return newmem; + // FIXME: NVIDIA targets currently use the built-in 'malloc' which we cannot + // reason with. But we still need to provide this function for compatibility. +#ifndef LIBC_TARGET_ARCH_IS_NVPTX + return gpu::reallocate(ptr, size); +#else + (void)ptr; + (void)size; + return nullptr; +#endif } } // namespace LIBC_NAMESPACE_DECL diff --git a/libc/test/integration/src/stdlib/gpu/CMakeLists.txt b/libc/test/integration/src/stdlib/gpu/CMakeLists.txt index 26c877b1b6ae6..69e1909fe78ed 100644 --- a/libc/test/integration/src/stdlib/gpu/CMakeLists.txt +++ b/libc/test/integration/src/stdlib/gpu/CMakeLists.txt @@ -17,6 +17,21 @@ if(NOT LIBC_TARGET_ARCHITECTURE_IS_NVPTX) --blocks 1024 ) + add_integration_test( + realloc + SUITE + stdlib-gpu-integration-tests + SRCS + realloc.cpp + DEPENDS + libc.src.stdlib.malloc + libc.src.stdlib.free + libc.src.stdlib.realloc + LOADER_ARGS + --threads 256 + --blocks 1024 + ) + add_integration_test( malloc_stress SUITE diff --git a/libc/test/integration/src/stdlib/gpu/realloc.cpp b/libc/test/integration/src/stdlib/gpu/realloc.cpp new file mode 100644 index 0000000000000..7cdf29489928a --- /dev/null +++ b/libc/test/integration/src/stdlib/gpu/realloc.cpp @@ -0,0 +1,44 @@ +#include "test/IntegrationTest/test.h" + +#include "src/__support/GPU/utils.h" +#include "src/stdlib/free.h" +#include "src/stdlib/malloc.h" +#include "src/stdlib/realloc.h" + +using namespace LIBC_NAMESPACE; + +TEST_MAIN(int, char **, char **) { + // realloc(nullptr, size) is equivalent to malloc. + int *alloc = reinterpret_cast(LIBC_NAMESPACE::realloc(nullptr, 32)); + EXPECT_NE(alloc, nullptr); + *alloc = 42; + EXPECT_EQ(*alloc, 42); + + // realloc to same size returns the same pointer. + void *same = LIBC_NAMESPACE::realloc(alloc, 32); + EXPECT_EQ(same, alloc); + EXPECT_EQ(reinterpret_cast(same)[0], 42); + + // realloc to smaller size returns same pointer. + void *smaller = LIBC_NAMESPACE::realloc(same, 16); + EXPECT_EQ(smaller, alloc); + EXPECT_EQ(reinterpret_cast(smaller)[0], 42); + + // realloc to larger size returns new pointer and preserves contents. + int *larger = reinterpret_cast(LIBC_NAMESPACE::realloc(smaller, 128)); + EXPECT_NE(larger, nullptr); + EXPECT_EQ(larger[0], 42); + + // realloc works when called with a divergent size. + int *div = reinterpret_cast( + LIBC_NAMESPACE::malloc((gpu::get_thread_id() + 1) * 16)); + EXPECT_NE(div, nullptr); + div[0] = static_cast(gpu::get_thread_id()); + int *div_realloc = reinterpret_cast( + LIBC_NAMESPACE::realloc(div, ((gpu::get_thread_id() + 1) * 32))); + EXPECT_NE(div_realloc, nullptr); + EXPECT_EQ(div_realloc[0], static_cast(gpu::get_thread_id())); + LIBC_NAMESPACE::free(div_realloc); + + return 0; +}