diff --git a/Tests/test_arrow.py b/Tests/test_arrow.py index 231db6512..ee6134001 100644 --- a/Tests/test_arrow.py +++ b/Tests/test_arrow.py @@ -170,8 +170,7 @@ def test_multiblock_l_image(): with pytest.raises(ValueError): (schema, arr) = img.__arrow_c_array__() - -def test_multiblock__rgba_image(): +def test_multiblock_rgba_image(): block_size = Image.core.get_block_size() # check a 2 block image in 4 channel mode @@ -192,8 +191,7 @@ def test_multiblock_l_schema(): with pytest.raises(ValueError): schema = img.__arrow_c_schema__() - -def test_multiblock__rgba_schema(): +def test_multiblock_rgba_schema(): block_size = Image.core.get_block_size() # check a 2 block image in 4 channel mode @@ -201,4 +199,62 @@ def test_multiblock__rgba_schema(): img = Image.new("RGBA", size, (128, 127, 126, 125)) with pytest.raises(ValueError): - schema = img.__arrow_c_schema__() + schema= img.__arrow_c_schema__() + + +def test_singleblock_l_image(): + Image.core.set_use_block_allocator(1) + + block_size = Image.core.get_block_size() + + # check a 2 block image in 4 channel mode + size = (4096, 2* (block_size//4096)) + img = Image.new('L', size, 128) + assert img.im.isblock() + + (schema, arr) = img.__arrow_c_array__() + assert schema + assert arr + + Image.core.set_use_block_allocator(0) + +def test_singleblock_rgba_image(): + Image.core.set_use_block_allocator(1) + block_size = Image.core.get_block_size() + + # check a 2 block image in 4 channel mode + size = (4096, (block_size//4096) //2) + img = Image.new('RGBA', size, (128,127,126,125)) + assert img.im.isblock() + + (schema, arr) = img.__arrow_c_array__() + assert schema + assert arr + Image.core.set_use_block_allocator(0) + + +def test_singleblock_l_schema(): + Image.core.set_use_block_allocator(1) + block_size = Image.core.get_block_size() + + # check a 2 block image in single channel mode + size = (4096, 2*block_size//4096) + img = Image.new('L', size, 128) + assert img.im.isblock() + + schema = img.__arrow_c_schema__() + assert schema + Image.core.set_use_block_allocator(0) + +def test_singleblock_rgba_schema(): + Image.core.set_use_block_allocator(1) + block_size = Image.core.get_block_size() + + # check a 2 block image in 4 channel mode + size = (4096, (block_size//4096) //2) + img = Image.new('RGBA', size, (128,127,126,125)) + assert img.im.isblock() + + schema= img.__arrow_c_schema__() + assert schema + Image.core.set_use_block_allocator(0) diff --git a/src/_imaging.c b/src/_imaging.c index b776267bc..54ec59ac7 100644 --- a/src/_imaging.c +++ b/src/_imaging.c @@ -4192,6 +4192,22 @@ _set_blocks_max(PyObject *self, PyObject *args) { return Py_None; } +static PyObject * +_set_use_block_allocator(PyObject *self, PyObject *args) { + int use_block_allocator; + if (!PyArg_ParseTuple(args, "i:set_use_block_allocator", &use_block_allocator)) { + return NULL; + } + ImagingMemorySetBlockAllocator(&ImagingDefaultArena, use_block_allocator); + Py_RETURN_NONE; +} + + +static PyObject * +_get_use_block_allocator(PyObject *self, PyObject *args) { + return PyLong_FromLong(ImagingDefaultArena.use_block_allocator); +} + static PyObject * _clear_cache(PyObject *self, PyObject *args) { int i = 0; @@ -4399,9 +4415,11 @@ static PyMethodDef functions[] = { {"get_alignment", (PyCFunction)_get_alignment, METH_VARARGS}, {"get_block_size", (PyCFunction)_get_block_size, METH_VARARGS}, {"get_blocks_max", (PyCFunction)_get_blocks_max, METH_VARARGS}, + {"get_use_block_allocator", (PyCFunction)_get_use_block_allocator, METH_VARARGS}, {"set_alignment", (PyCFunction)_set_alignment, METH_VARARGS}, {"set_block_size", (PyCFunction)_set_block_size, METH_VARARGS}, {"set_blocks_max", (PyCFunction)_set_blocks_max, METH_VARARGS}, + {"set_use_block_allocator", (PyCFunction)_set_use_block_allocator, METH_VARARGS}, {"clear_cache", (PyCFunction)_clear_cache, METH_VARARGS}, {NULL, NULL} /* sentinel */ diff --git a/src/libImaging/Imaging.h b/src/libImaging/Imaging.h index 09155765b..fb3d31a99 100644 --- a/src/libImaging/Imaging.h +++ b/src/libImaging/Imaging.h @@ -178,6 +178,7 @@ typedef struct ImagingMemoryArena { int stats_reallocated_blocks; /* Number of blocks which were actually reallocated after retrieving */ int stats_freed_blocks; /* Number of freed blocks */ + int use_block_allocator; /* don't use arena, use block allocator */ #ifdef Py_GIL_DISABLED PyMutex mutex; #endif @@ -191,6 +192,8 @@ extern int ImagingMemorySetBlocksMax(ImagingMemoryArena arena, int blocks_max); extern void ImagingMemoryClearCache(ImagingMemoryArena arena, int new_size); +extern void +ImagingMemorySetBlockAllocator(ImagingMemoryArena arena, int use_block_allocator); extern Imaging ImagingNew(const char *mode, int xsize, int ysize); diff --git a/src/libImaging/Storage.c b/src/libImaging/Storage.c index 6b1937a56..a1241098e 100644 --- a/src/libImaging/Storage.c +++ b/src/libImaging/Storage.c @@ -341,6 +341,7 @@ struct ImagingMemoryArena ImagingDefaultArena = { 0, 0, 0, // Stats + 0, // use_block_allocator #ifdef Py_GIL_DISABLED {0}, #endif @@ -373,6 +374,12 @@ ImagingMemorySetBlocksMax(ImagingMemoryArena arena, int blocks_max) { return 1; } +void +ImagingMemorySetBlockAllocator(ImagingMemoryArena arena, int use_block_allocator) { + arena->use_block_allocator = use_block_allocator; +} + + void ImagingMemoryClearCache(ImagingMemoryArena arena, int new_size) { while (arena->blocks_cached > new_size) { @@ -649,11 +656,17 @@ ImagingNewInternal(const char *mode, int xsize, int ysize, int dirty) { Imaging ImagingNew(const char *mode, int xsize, int ysize) { + if (ImagingDefaultArena.use_block_allocator) { + return ImagingNewBlock(mode, xsize, ysize); + } return ImagingNewInternal(mode, xsize, ysize, 0); } Imaging ImagingNewDirty(const char *mode, int xsize, int ysize) { + if (ImagingDefaultArena.use_block_allocator) { + return ImagingNewBlock(mode, xsize, ysize); + } return ImagingNewInternal(mode, xsize, ysize, 1); }