From 808451680a31d11168c60b39c025880bbc7a6399 Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Mon, 15 Jun 2020 17:23:31 -0500
Subject: [PATCH] Test, fix allocator behavior for zero-size allocations

---
 src/wrap_mempool.cpp |  9 +++++++++
 test/test_wrapper.py | 20 +++++++++++++++++++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/src/wrap_mempool.cpp b/src/wrap_mempool.cpp
index dceb33dc..16662bc5 100644
--- a/src/wrap_mempool.cpp
+++ b/src/wrap_mempool.cpp
@@ -197,6 +197,15 @@ namespace
       alloc.try_release_blocks();
     }
 
+    if (!mem)
+    {
+      if (size == 0)
+        return nullptr;
+      else
+        throw pyopencl::error("Allocator", CL_INVALID_VALUE,
+            "allocator succeeded but returned NULL cl_mem");
+    }
+
     try
     {
       return new pyopencl::buffer(mem, false);
diff --git a/test/test_wrapper.py b/test/test_wrapper.py
index dc5772de..1b4a3b42 100644
--- a/test/test_wrapper.py
+++ b/test/test_wrapper.py
@@ -33,7 +33,8 @@ import pyopencl.array as cl_array
 import pyopencl.cltypes as cltypes
 import pyopencl.clrandom
 from pyopencl.tools import (  # noqa
-        pytest_generate_tests_for_pyopencl as pytest_generate_tests)
+        pytest_generate_tests_for_pyopencl as pytest_generate_tests,
+        ImmediateAllocator, DeferredAllocator)
 from pyopencl.characterize import get_pocl_version
 
 # Are CL implementations crashy? You be the judge. :)
@@ -573,6 +574,23 @@ def test_mempool_2(ctx_factory):
         assert asize < asize*(1+1/8)
 
 
+@pytest.mark.parametrize("allocator_cls", [ImmediateAllocator, DeferredAllocator])
+def test_allocator(ctx_factory, allocator_cls):
+    context = ctx_factory()
+    queue = cl.CommandQueue(context)
+
+    if allocator_cls is DeferredAllocator:
+        allocator = allocator_cls(context)
+    else:
+        allocator = allocator_cls(queue)
+
+    mem = allocator(15)
+    mem2 = allocator(0)
+
+    assert mem is not None
+    assert mem2 is None
+
+
 def test_vector_args(ctx_factory):
     context = ctx_factory()
     queue = cl.CommandQueue(context)
-- 
GitLab