cumo 0.5.6 → 0.5.7

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 324cdfec0cf163d40b3b3b635fa69ecf516aa12e15a0bf614cbc34d3fe61a4c7
-  data.tar.gz: 3d0f6c8b66a9b51eca9fe6912c1a4e7c3534a028ed4c8d740467ae8d6e215824
+  metadata.gz: 3ddc7b756681afdf9fccd45652c422d2f67accfd1cf238d1fd05f6827fb5029d
+  data.tar.gz: 47848503fd8a388b7dfee79d2da0e627032fd4d8c8717645da70d8f09770f599
 SHA512:
-  metadata.gz: 9c82a985c471260943c4716630d6db803706ac50cccd5e1fc8a0a2bb02d80172ec38743f5c5f286467e97814843b8fbe1d6ec1f98e75081364fd9a5c0a4da832
-  data.tar.gz: d3ed8c3e8844d51676c4011139faa509de8cdb39ef325d84b4b8b41f6725fb63b9893fae8dbced5e546810ff4eb570f60a347ef248323f805bffa9fde60faab1
+  metadata.gz: f59ddc27c8f96a5eceec7fb4ae6665d1fae6dd20ef039a83563c0e0f5a76780fd2ab9b5ebbe06e116a3f01cd59c7493eee06dc6c1166a0bb82ba546c43331f03
+  data.tar.gz: 16543c06d74f79089799442abc30741ffb2a26556da953410023b19711fd25e15adb05548f9664c79bf8430a3be9c8f15c0813dccfae4b5e918ffc55a5fc68c0

data/CHANGELOG.md

@@ -1,3 +1,28 @@
+# 0.5.7 (2026/08/13)
+
+Breaking changes:
+
+* `reshape` raises `ArgumentError` instead of `RangeError` for a dimension above `INT_MAX` (PR #183)
+* `parse` accepts only numeric literals and `true`/`false`/`nil`, so `parse("3/4")` raises `ArgumentError` instead of returning `[[0]]` (PR #190)
+
+Fixes:
+
+* Fix segfaults from CUDA handles the process did not create, and from destroying one twice (PR #192)
+* Fix process abort from freeing NArray data with the wrong allocator, including a raise out of GC for `Cumo::RObject` (PR #191)
+* Fix arbitrary code execution in `parse`, which ran `eval` on every token (PR #190)
+* Fix stack buffer overflow in `Cumo::RObject#format` with an element longer than 47 characters (PR #189)
+* Fix a freed chunk returning to the arena of the wrong stream (PR #187)
+* Fix `Malloc` searching the free list with the default stream's index (PR #186)
+* Fix data race on the memory pool chunk graph and on the per-device pool table (PR #185)
+* Fix a SIGFPE on fractional range steps and an out-of-bounds read on newaxis (PR #184)
+* Fix a SIGFPE and a silently corrupt shape from unchecked `reshape` arguments (PR #183)
+* Fix segfault when `store_binary` receives a non-String (PR #182)
+* Fix missing range check on NArray indices, which segfaulted on a negative one (PR #181)
+
+Changes:
+
+* Compile the memory pool test harness in CI, and run it from `rake test` (PR #188)
+
 # 0.5.6 (2026/08/09)
 
 Fixes:

data/Rakefile

@@ -12,10 +12,19 @@ end
 require "rake/extensiontask"
 Rake::ExtensionTask.new("cumo")
 
-task :ctest do
-  sh 'cd ext/cumo && ruby extconf.rb && make && make build-ctest && make run-ctest'
+desc 'Build the C++ memory pool test harness without running it (running it needs a GPU)'
+task :build_ctest do
+  sh 'cd ext/cumo && ruby extconf.rb && make build-ctest'
 end
 
+task :ctest => :build_ctest do
+  sh 'cd ext/cumo && make run-ctest'
+end
+
+# Both need a GPU, so CI runs neither. Tie them together so the one place
+# they do run does not skip half of them.
+task :test => :ctest
+
 task :docs do
   dir = "ext/cumo"
   srcs = %w[array.c data.c index.c math.c narray.c rand.c struct.c].map { |s| File.join(dir, "narray", s) }

data/ext/cumo/cuda/driver.c

@@ -3,12 +3,16 @@
 #include <cuda.h>
 #include <cuda_runtime.h>
 #include "cumo/cuda/driver.h"
+#include "cumo/cuda/handle.h"
 
 VALUE cumo_cuda_eDriverError;
 VALUE cumo_cuda_mDriver;
 #define eDriverError cumo_cuda_eDriverError
 #define mDriver cumo_cuda_mDriver
 
+static cumo_cuda_handle_set_t link_states;
+static cumo_cuda_handle_set_t modules;
+
 static void
 check_status(CUresult status)
 {
@@ -100,13 +104,13 @@ cuLinkAddData_without_gvl_cb(void *param)
 static VALUE
 rb_cuLinkAddData(VALUE self, VALUE state, VALUE type, VALUE data, VALUE name)
 {
-    CUlinkState _state = (CUlinkState)NUM2SIZET(state);
     CUjitInputType _type = (CUjitInputType)NUM2INT(type);
     // The image may be a cubin, so it is taken by length and allowed to hold
     // NUL bytes; the name is a plain C string cuLinkAddData reports errors with.
     void* _data = (void *)StringValuePtr(data);
     size_t _size = RSTRING_LEN(data);
     const char* _name = StringValueCStr(name);
+    CUlinkState _state = (CUlinkState)cumo_cuda_handle_get(&link_states, state, "CUlinkState");
     CUresult status;
 
     struct cuLinkAddDataParam param = {_state, _type, _data, _size, _name, 0, (CUjit_option*)0, (void**)0};
@@ -141,9 +145,9 @@ cuLinkAddFile_without_gvl_cb(void *param)
 static VALUE
 rb_cuLinkAddFile(VALUE self, VALUE state, VALUE type, VALUE path)
 {
-    CUlinkState _state = (CUlinkState)NUM2SIZET(state);
     CUjitInputType _type = (CUjitInputType)NUM2INT(type);
     const char* _path = StringValueCStr(path);
+    CUlinkState _state = (CUlinkState)cumo_cuda_handle_get(&link_states, state, "CUlinkState");
     CUresult status;
 
     struct cuLinkAddFileParam param = {_state, _type, _path, 0, (CUjit_option*)0, (void **)0};
@@ -173,7 +177,7 @@ cuLinkComplete_without_gvl_cb(void *param)
 static VALUE
 rb_cuLinkComplete(VALUE self, VALUE state)
 {
-    CUlinkState _state = (CUlinkState)NUM2SIZET(state);
+    CUlinkState _state = (CUlinkState)cumo_cuda_handle_get(&link_states, state, "CUlinkState");
     void* _cubinOut;
     size_t _sizeOut;
     CUresult status;
@@ -214,6 +218,7 @@ rb_cuLinkCreate(VALUE self)
     //status = cuLinkCreate(0, (CUjit_option*)0, (void**)0, &state);
 
     check_status(status);
+    cumo_cuda_handle_set_add(&link_states, (size_t)state);
     return SIZET2NUM((size_t)state);
 }
 
@@ -233,7 +238,7 @@ cuLinkDestroy_without_gvl_cb(void *param)
 static VALUE
 rb_cuLinkDestroy(VALUE self, VALUE state)
 {
-    CUlinkState _state = (CUlinkState)NUM2SIZET(state);
+    CUlinkState _state = (CUlinkState)cumo_cuda_handle_take(&link_states, state, "CUlinkState");
     CUresult status;
 
     struct cuLinkDestroyParam param = {_state};
@@ -263,8 +268,8 @@ static VALUE
 rb_cuModuleGetFunction(VALUE self, VALUE hmod, VALUE name)
 {
     CUfunction _hfunc;
-    CUmodule _hmod = (CUmodule)NUM2SIZET(hmod);
     const char* _name = StringValueCStr(name);
+    CUmodule _hmod = (CUmodule)cumo_cuda_handle_get(&modules, hmod, "CUmodule");
     CUresult status;
 
     struct cuModuleGetFunctionParam param = {&_hfunc, _hmod, _name};
@@ -297,8 +302,8 @@ rb_cuModuleGetGlobal(VALUE self, VALUE hmod, VALUE name)
 {
     CUdeviceptr _dptr;
     size_t _bytes;
-    CUmodule _hmod = (CUmodule)NUM2SIZET(hmod);
     const char* _name = StringValueCStr(name);
+    CUmodule _hmod = (CUmodule)cumo_cuda_handle_get(&modules, hmod, "CUmodule");
     CUresult status;
     VALUE ret;
 
@@ -342,6 +347,7 @@ rb_cuModuleLoad(VALUE self, VALUE fname)
 
     RB_GC_GUARD(fname);
     check_status(status);
+    cumo_cuda_handle_set_add(&modules, (size_t)_module);
     return SIZET2NUM((size_t)_module);
 }
 
@@ -373,6 +379,7 @@ rb_cuModuleLoadData(VALUE self, VALUE image)
 
     RB_GC_GUARD(image);
     check_status(status);
+    cumo_cuda_handle_set_add(&modules, (size_t)_module);
     return SIZET2NUM((size_t)_module);
 }
 
@@ -392,7 +399,7 @@ cuModuleUnload_without_gvl_cb(void *param)
 static VALUE
 rb_cuModuleUnload(VALUE self, VALUE hmod)
 {
-    CUmodule _hmod = (CUmodule)NUM2SIZET(hmod);
+    CUmodule _hmod = (CUmodule)cumo_cuda_handle_take(&modules, hmod, "CUmodule");
     CUresult status;
 
     struct cuModuleUnloadParam param = {_hmod};
@@ -414,6 +421,9 @@ Init_cumo_cuda_driver()
     mDriver = rb_define_module_under(mCUDA, "Driver");
     eDriverError = rb_define_class_under(mCUDA, "DriverError", rb_eStandardError);
 
+    cumo_cuda_handle_set_init(&link_states);
+    cumo_cuda_handle_set_init(&modules);
+
     rb_define_singleton_method(mDriver, "cuCtxGetCurrent", rb_cuCtxGetCurrent, 0);
     rb_define_singleton_method(mDriver, "cuLinkAddData",   rb_cuLinkAddData,   4);
     rb_define_singleton_method(mDriver, "cuLinkAddFile", rb_cuLinkAddFile, 3);

data/ext/cumo/cuda/memory_pool.cpp

@@ -56,7 +56,6 @@ cumo_cuda_runtime_free(char *ptr)
     // releases memory the pool still hands out, and fails outright for a chunk
     // which is not at the head of its buffer.
     try {
-        // TODO(sonots): Get current CUDA stream and pass it
         if (pool.Free(reinterpret_cast<intptr_t>(ptr))) {
             return;
         }

data/ext/cumo/cuda/memory_pool_impl.cpp

@@ -132,7 +132,7 @@ intptr_t SingleDeviceMemoryPool::Malloc(size_t size, cudaStream_t stream_ptr) {
 
         // find best-fit, or a smallest larger allocation
         Arena& arena = GetArena(stream_ptr);
-        size_t arena_index = GetArenaIndex(size);
+        size_t arena_index = GetArenaIndex(size, stream_ptr);
         size_t arena_length = arena.size();
         for (size_t i = arena_index; i < arena_length; ++i) {
             FreeList& free_list = arena[i];
@@ -143,15 +143,22 @@ intptr_t SingleDeviceMemoryPool::Malloc(size_t size, cudaStream_t stream_ptr) {
             // TODO(sonots): compact_index
             break;
         }
-    }
 
-    if (chunk != nullptr) {
-        std::shared_ptr<Chunk> remaining = Split(chunk, size);
-        if (remaining != nullptr) {
-            AppendToFreeList(remaining->size(), remaining, stream_ptr);
+        // Splitting rewrites the prev/next pointers of the neighbouring
+        // chunks, which other threads reach through the free lists, so it has
+        // to happen under the same lock as the search above.
+        if (chunk != nullptr) {
+            std::shared_ptr<Chunk> remaining = Split(chunk, size);
+            if (remaining != nullptr) {
+                AppendToFreeList(remaining->size(), remaining, stream_ptr);
+            }
         }
-    } else {
-        // cudaMalloc if a cache is not found
+    }
+
+    if (chunk == nullptr) {
+        // cudaMalloc if a cache is not found. This stays outside the lock: it
+        // is slow, and a chunk of a fresh allocation has no neighbours for
+        // another thread to reach it through.
         std::shared_ptr<Memory> mem = nullptr;
         try {
             mem = std::make_shared<Memory>(size);
@@ -186,12 +193,23 @@ intptr_t SingleDeviceMemoryPool::Malloc(size_t size, cudaStream_t stream_ptr) {
     return chunk->ptr();
 }
 
-bool SingleDeviceMemoryPool::Free(intptr_t ptr, cudaStream_t stream_ptr) {
+bool SingleDeviceMemoryPool::Free(intptr_t ptr) {
     std::shared_ptr<Chunk> chunk = nullptr;
 
-    {
-        std::lock_guard<std::recursive_mutex> lock{mutex_};
+    // The whole body runs under the lock. Walking prev/next and merging is a
+    // read-modify-write of the chunk graph, which is shared: a neighbour can
+    // be split by a concurrent Malloc while this reads its prev/next.
+    //
+    // The lost merge is the part that shows: two threads freeing neighbouring
+    // chunks each look for the other in the free list before the other has
+    // appended itself, so both merges fail and the two chunks stay split for
+    // good. Fragmentation accumulates from there.
+    //
+    // mutex_ is recursive, so the RemoveFromFreeList/AppendToFreeList calls
+    // below can take it again.
+    std::lock_guard<std::recursive_mutex> lock{mutex_};
 
+    {
         // find rather than operator[], which would insert an empty entry for
         // every pointer this pool does not own.
         auto it = in_use_.find(ptr);
@@ -206,6 +224,14 @@ bool SingleDeviceMemoryPool::Free(intptr_t ptr, cudaStream_t stream_ptr) {
         chunk->set_in_use(false);
     }
 
+    // A chunk belongs to the arena of the stream it was allocated on. That is
+    // what makes reuse safe: the next allocation from that arena is queued on
+    // the same stream as the work that last used the chunk, so the two are
+    // ordered. Returning it to another stream's arena hands it to work with no
+    // such ordering. Malloc and Merge each state the invariant in an assert,
+    // but ruby.h defines NDEBUG, so neither ever runs.
+    const cudaStream_t stream_ptr = chunk->stream_ptr();
+
     if (chunk->next() != nullptr && !chunk->next()->in_use()) {
         if (RemoveFromFreeList(chunk->next()->size(), chunk->next(), stream_ptr)) {
             Merge(chunk, chunk->next());

data/ext/cumo/cuda/memory_pool_impl.hpp

@@ -174,7 +174,10 @@ public:
 
     // Returns false if this pool did not allocate the pointer, in which case
     // nothing was freed.
-    bool Free(intptr_t ptr, cudaStream_t stream_ptr = 0);
+    //
+    // There is no stream parameter: the chunk records the stream it was
+    // allocated on, and that is the only arena it may be returned to.
+    bool Free(intptr_t ptr);
 
     // Free all **non-split** chunks in all arenas
     void FreeAllBlocks();
@@ -216,7 +219,12 @@ public:
         return (size - 1) / kRoundSize;
     }
 
-    size_t GetArenaIndex(size_t size, cudaStream_t stream_ptr = 0) {
+    // Returns where a chunk of `size` belongs in the arena of `stream_ptr`.
+    //
+    // The stream is not optional: an index taken from one stream's index map
+    // means nothing in another stream's arena, and defaulting it silently
+    // started the free list search at the wrong bin.
+    size_t GetArenaIndex(size_t size, cudaStream_t stream_ptr) {
         size_t bin_index = GetBinIndex(size);
         ArenaIndexMap& arena_index_map = GetArenaIndexMap(stream_ptr);
         return std::lower_bound(arena_index_map.begin(), arena_index_map.end(), bin_index) - arena_index_map.begin();
@@ -267,14 +275,14 @@ public:
         return true;
     }
 
-    void AppendToFreeList(size_t size, std::shared_ptr<Chunk>& chunk, cudaStream_t stream_ptr = 0);
+    void AppendToFreeList(size_t size, std::shared_ptr<Chunk>& chunk, cudaStream_t stream_ptr);
 
     // Removes the chunk from the free list.
     //
     // @return true if the chunk can successfully be removed from
     //         the free list. false` otherwise (e.g., the chunk could not
     //         be found in the free list as the chunk is allocated.)
-    bool RemoveFromFreeList(size_t size, std::shared_ptr<Chunk>& chunk, cudaStream_t stream_ptr = 0);
+    bool RemoveFromFreeList(size_t size, std::shared_ptr<Chunk>& chunk, cudaStream_t stream_ptr);
 
     void CompactIndex(cudaStream_t stream_ptr, bool free);
 };
@@ -304,6 +312,20 @@ private:
     }
 
     std::unordered_map<int, SingleDeviceMemoryPool> pools_;
+    std::mutex pools_mutex_;
+
+    // Returns the pool of the current device, creating it on first use.
+    //
+    // The lock only has to cover the lookup: pools_ is never erased from
+    // before destruction, and unordered_map keeps references to its elements
+    // valid across a rehash, so the returned reference outlives the lock.
+    // Without it, the insert operator[] performs on first use could rehash the
+    // map underneath another thread walking it.
+    SingleDeviceMemoryPool& GetPool() {
+        int id = device_id();  // a CUDA call, so keep it out of the lock
+        std::lock_guard<std::mutex> lock{pools_mutex_};
+        return pools_[id];  // find or create
+    }
 
 public:
     MemoryPool() {}
@@ -318,35 +340,48 @@ public:
     // Returns:
     //     intptr_t: Pointer address to the allocated buffer.
     intptr_t Malloc(size_t size, cudaStream_t stream_ptr = 0) {
-        auto& mp = pools_[device_id()];
-        return mp.Malloc(size, stream_ptr);
+        return GetPool().Malloc(size, stream_ptr);
     }
 
     // Frees the memory, to the pool
     //
     // Args:
     //     ptr (intptr_t): Pointer of the memory buffer
-    //     stream_ptr (cudaStream_t): Return the memory to the arena of given stream
     // Returns:
     //     bool: false if no pool allocated the pointer, in which case nothing
     //           was freed and the caller has to free it by its own means.
-    bool Free(intptr_t ptr, cudaStream_t stream_ptr = 0) {
-        if (pools_.empty()) {  // nothing has ever been allocated from a pool
-            return false;
-        }
+    //
+    // The stream is not a parameter: the chunk records the stream it was
+    // allocated on, and that is the only arena it may be returned to.
+    bool Free(intptr_t ptr) {
         int current_device_id = device_id();
-        auto it = pools_.find(current_device_id);
-        if (it != pools_.end() && it->second.Free(ptr, stream_ptr)) {
+
+        // Take the pointers out under the lock and release it before freeing:
+        // an element of pools_ stays put once inserted, so the pointers remain
+        // valid, and this keeps a Free of one device's pool from blocking a
+        // Malloc on another. `others` does not allocate in the usual
+        // single-device case, where it stays empty.
+        SingleDeviceMemoryPool* current = nullptr;
+        std::vector<SingleDeviceMemoryPool*> others;
+        {
+            std::lock_guard<std::mutex> lock{pools_mutex_};
+            for (auto& entry : pools_) {
+                if (entry.first == current_device_id) {
+                    current = &entry.second;
+                } else {
+                    others.emplace_back(&entry.second);
+                }
+            }
+        }
+
+        if (current != nullptr && current->Free(ptr)) {
             return true;
         }
         // The current device may have been switched since the allocation.
         // cudaMallocManaged hands out addresses which are unique across the
         // host and every device, so the pointer identifies its pool on its own.
-        for (auto& entry : pools_) {
-            if (entry.first == current_device_id) {
-                continue;
-            }
-            if (entry.second.Free(ptr, stream_ptr)) {
+        for (SingleDeviceMemoryPool* mp : others) {
+            if (mp->Free(ptr)) {
                 return true;
             }
         }
@@ -355,8 +390,7 @@ public:
 
     // Free all **non-split** chunks in all arenas
     void FreeAllBlocks() {
-        auto& mp = pools_[device_id()];
-        return mp.FreeAllBlocks();
+        return GetPool().FreeAllBlocks();
     }
 
     // Free all **non-split** chunks in specified arena
@@ -364,8 +398,7 @@ public:
     // Args:
     //     stream_ptr (cudaStream_t): Release free blocks in the arena of given stream
     void FreeAllBlocks(cudaStream_t stream_ptr) {
-        auto& mp = pools_[device_id()];
-        return mp.FreeAllBlocks(stream_ptr);
+        return GetPool().FreeAllBlocks(stream_ptr);
     }
 
     // Count the total number of free blocks.
@@ -373,8 +406,7 @@ public:
     // Returns:
     //     size_t: The total number of free blocks.
     size_t GetNumFreeBlocks() {
-        auto& mp = pools_[device_id()];
-        return mp.GetNumFreeBlocks();
+        return GetPool().GetNumFreeBlocks();
     }
 
     // Get the total number of bytes used.
@@ -382,8 +414,7 @@ public:
     // Returns:
     //     size_t: The total number of bytes used.
     size_t GetUsedBytes() {
-        auto& mp = pools_[device_id()];
-        return mp.GetUsedBytes();
+        return GetPool().GetUsedBytes();
     }
 
     // Get the total number of bytes acquired but not used in the pool.
@@ -391,8 +422,7 @@ public:
     // Returns:
     //     size_t: The total number of bytes acquired but not used in the pool.
     size_t GetFreeBytes() {
-        auto& mp = pools_[device_id()];
-        return mp.GetFreeBytes();
+        return GetPool().GetFreeBytes();
     }
 
     // Get the total number of bytes acquired in the pool.
@@ -400,8 +430,7 @@ public:
     // Returns:
     //     size_t: The total number of bytes acquired in the pool.
     size_t GetTotalBytes() {
-        auto& mp = pools_[device_id()];
-        return mp.GetTotalBytes();
+        return GetPool().GetTotalBytes();
     }
 };
 

data/ext/cumo/cuda/memory_pool_impl_test.cpp

@@ -1,8 +1,12 @@
 #include "memory_pool_impl.hpp"
 
+#include <atomic>
 #include <cassert>
+#include <cstring>
 #include <memory>
 #include <iostream>
+#include <thread>
+#include <vector>
 
 // TODO(sonots): Use googletest?
 // TODO(sonots): Provide clean way to build this test outside extconf.rb
@@ -135,6 +139,8 @@ public:
         TearDown(); SetUp(); TestGetRoundedSize();
         TearDown(); SetUp(); TestGetBinIndex();
         TearDown(); SetUp(); TestGetArenaIndexWithHugeSize();
+        TearDown(); SetUp(); TestMallocOnAnotherStream();
+        TearDown(); SetUp(); TestFreeReturnsChunkToItsOwnStream();
         TearDown(); SetUp(); TestAppendToFreeList();
         TearDown(); SetUp(); TestRemoveFromFreeList();
         TearDown(); SetUp(); TestMalloc();
@@ -183,6 +189,53 @@ public:
         assert(pool_->GetArenaIndex(kMaxAllocationSize, stream_ptr_) == 1);
     }
 
+    // Malloc used to take the starting arena index from stream 0's index map
+    // and then index the requested stream's arena with it. The two index maps
+    // are unrelated, so the search could start below the requested size and
+    // hand out a chunk too small for it.
+    void TestMallocOnAnotherStream() {
+        cudaStream_t other = reinterpret_cast<cudaStream_t>(1);
+
+        auto small_mem = std::make_shared<Memory>(kRoundSize);
+        auto small = std::make_shared<Chunk>(small_mem, 0, small_mem->size(), other);
+        pool_->AppendToFreeList(small->size(), small, other);
+
+        auto big_mem = std::make_shared<Memory>(kRoundSize * 4);
+        auto big = std::make_shared<Chunk>(big_mem, 0, big_mem->size(), other);
+        pool_->AppendToFreeList(big->size(), big, other);
+
+        // stream 0 has no bins at all here, so its index is 0 for any size.
+        assert(pool_->GetArenaIndex(kRoundSize * 4, stream_ptr_) == 0);
+        assert(pool_->GetArenaIndex(kRoundSize * 4, other) == 1);
+
+        auto p = pool_->Malloc(kRoundSize * 4, other);
+        assert(p == big->ptr());
+        assert(pool_->GetUsedBytes() == kRoundSize * 4);
+        // The small chunk is untouched, and no chunk was split to a size its
+        // buffer cannot back -- which used to underflow GetFreeBytes.
+        assert(pool_->GetFreeBytes() == kRoundSize);
+        assert(small->size() == kRoundSize);
+    }
+
+    // Free used to append the chunk to the arena of the stream it was passed,
+    // and cumo_cuda_runtime_free passes none, so anything allocated on another
+    // stream was returned to stream 0's arena -- where a later Malloc on
+    // stream 0 would reuse memory it never ordered against.
+    void TestFreeReturnsChunkToItsOwnStream() {
+        cudaStream_t other = reinterpret_cast<cudaStream_t>(1);
+
+        auto p = pool_->Malloc(kRoundSize, other);
+        pool_->Free(p);
+        assert(pool_->GetNumFreeBlocks() == 1);
+        assert(!pool_->HasArena(stream_ptr_));
+
+        // The owning stream gets its own chunk back rather than cudaMalloc'ing
+        // a second one, which is the same thing seen from the other side.
+        auto q = pool_->Malloc(kRoundSize, other);
+        assert(q == p);
+        assert(pool_->GetTotalBytes() == kRoundSize);
+    }
+
     void TestAppendToFreeList() {
         Arena& arena = pool_->GetArena(stream_ptr_);
         ArenaIndexMap& arena_index_map = pool_->GetArenaIndexMap(stream_ptr_);
@@ -620,6 +673,95 @@ public:
     }
 };
 
+// The pool is reachable from several threads at once: cumo.c declares the
+// extension Ractor-safe, and a Ractor does not hold the GVL that serializes
+// ordinary Ruby threads.
+class TestConcurrency {
+private:
+    static const int kThreads = 8;
+    static const int kRounds = 400;
+    static const int kPerThread = 8;
+
+public:
+    void Run() {
+        TestConcurrentSplitAndMerge();
+        TestConcurrentPools();
+    }
+
+    // Chunks are only neighbours when they were Split from a common
+    // allocation, so prime the pool with one large free chunk: every Malloc
+    // below splits off it, and each thread then frees chunks sitting next to
+    // chunks other threads still hold. That is the state Free walks over.
+    void TestConcurrentSplitAndMerge() {
+        SingleDeviceMemoryPool pool;
+        const size_t total = size_t{kRoundSize} * kThreads * kPerThread * 2;
+        intptr_t big = pool.Malloc(total);
+        pool.Free(big);
+        assert(pool.GetNumFreeBlocks() == 1);
+
+        std::atomic<int> corrupted{0};
+        std::vector<std::thread> threads;
+        for (int t = 0; t < kThreads; ++t) {
+            threads.emplace_back([&pool, &corrupted, t] {
+                std::vector<intptr_t> ptrs(kPerThread);
+                for (int round = 0; round < kRounds; ++round) {
+                    for (int k = 0; k < kPerThread; ++k) {
+                        ptrs[k] = pool.Malloc(kRoundSize);
+                        // Stamp the chunk with a value no other thread uses.
+                        // Merging across a chunk which is still in use hands
+                        // the same bytes out twice, and the second holder's
+                        // stamp overwrites the first one's.
+                        std::memset(reinterpret_cast<void*>(ptrs[k]),
+                                    Stamp(t, k), kRoundSize);
+                    }
+                    for (int k = 0; k < kPerThread; ++k) {
+                        auto* p = reinterpret_cast<unsigned char*>(ptrs[k]);
+                        for (size_t i = 0; i < kRoundSize; ++i) {
+                            if (p[i] != Stamp(t, k)) {
+                                ++corrupted;
+                                break;
+                            }
+                        }
+                        pool.Free(ptrs[k]);
+                    }
+                }
+            });
+        }
+        for (auto& thread : threads) thread.join();
+
+        assert(corrupted == 0);
+        assert(pool.GetUsedBytes() == 0);
+        // Free merges a chunk with any free neighbour, so once every chunk is
+        // freed they all collapse back into the single block they came from.
+        // A lost merge leaves the pool fragmented forever.
+        assert(pool.GetNumFreeBlocks() == 1);
+        assert(pool.GetFreeBytes() == total);
+    }
+
+    // MemoryPool::pools_ is inserted into on first use of a device.
+    void TestConcurrentPools() {
+        MemoryPool pool;
+        std::vector<std::thread> threads;
+        for (int t = 0; t < kThreads; ++t) {
+            threads.emplace_back([&pool] {
+                for (int round = 0; round < kRounds; ++round) {
+                    intptr_t p = pool.Malloc(kRoundSize);
+                    assert(p != 0);
+                    pool.Free(p);
+                }
+            });
+        }
+        for (auto& thread : threads) thread.join();
+
+        assert(pool.GetUsedBytes() == 0);
+    }
+
+private:
+    static unsigned char Stamp(int t, int k) {
+        return static_cast<unsigned char>(t * kPerThread + k + 1);
+    }
+};
+
 // Resets the CUDA device, so it has to run after every other test.
 class TestMemoryDestructor {
 public:
@@ -650,6 +792,7 @@ int main() {
     cumo::internal::TestChunk{}.Run();
     cumo::internal::TestSingleDeviceMemoryPool{}.Run();
     cumo::internal::TestMemoryPool{}.Run();
+    cumo::internal::TestConcurrency{}.Run();
     cumo::internal::TestMemoryDestructor{}.Run();
     return 0;
 }