Is there a non-allocating way to modify elements of single-reference packed arrays?

Question

I need to repeatedly increment cuboidal chunks of a 3D image. The image is large - on the order of 1G elements.

The only way I know of doing it using built-in Image functions is:

accumulator = accumulator ~ ImageAdd ~ ImagePad[ImageConstant[1,size], padding]`

This of course constantly allocates the constant images and the padded outputs. Tens of thousands them, and at these sizes, it's slow.

So, instead, I can use a packed array to represent the image data, and add to that using Part and AddTo:

regions = { Splice[{1;;2, 3;;4, 5;;6}, Part], ... };
size = {1100, 1100, 1100};

accumulator = ConstantArray[0, size];
Assert[Developer`PackedArrayQ[accumulator]];

(* accumulator has a single reference here *)
Map[{region} |-> accumulator[[region]] += 1, regions];

result = Image3D[accumulator, "Bit16"];

Unfortunately, the AddTo function doesn't seem to have a specialization that would not reallocate, in spite of the modified array having just one reference*. Judging by the system memory use, each AddTo clones the image, modifies it, sets accumulator to it, then deallocates accumulator.

Is there some incantation that would do it without allocating memory?

* Bonus Question: Is there any way to inspect an object refcount in Mathematica?

---

The exact code I'm using is in a Module, and is:

gridCuboidImageData = {gcuboids} |->
   Module[{gridSize, i, gcuboid, start, size, ex, ey, ez, 
     n = Length[gcuboids], imageData},
    gridSize = Plus @@ gcuboids[[1]];
    (* if compiled to C, replace with: imageData=PadLeft[{{{}}},
    Reverse@gridSize] *)
    imageData = ConstantArray[0, Reverse@gridSize];
    Assert[Developer`PackedArrayQ[imageData]];
    For[i = 1, i <= n, i++,
     gcuboid = gcuboids[[i]];
     size = gcuboid[[3]];
     start = gcuboid[[1]];
     start = {start[[1]], gridSize[[2]] - start[[2]] - size[[2]], 
       gridSize[[3]] - start[[3]] - size[[3]]};
     ex = start[[1]] + size[[1]];
     ey = start[[2]] + size[[2]];
     ez = start[[3]] + size[[3]];
     imageData[[1 + start[[3]] ;; ez, 1 + start[[2]] ;; ey, 
       1 + start[[1]] ;; ex]] += 1;
     ];
    imageData
    ];

Answer 1

After I wrote the C code below, it came to my mind that the first thing you should do in your Mathematica session is to set

$HistoryLength = 0;

Otherwise, at least on my machine, repeated calls of something even so innocent as accumulator = ConstantArray[0, {1100,1100,1100}]; will quickly fill all RAM and lead to memory compression or swapping.

---

Sometimes, when one needs particular control of memory, it may be better to write plain C code and use LibraryLink to compile and link it to the Mathematica session. Even if the same can be done with pure Mathematica code, this way you don't have to bother how Mathematica handles the memory and you are guaranteed to get nearly optimal performance (on the CPU).

Here a simple example that takes the 3-tensor as "Shared" array and modifies it in-place.

Needs["CCompilerDriver`"]
Quiet[LibraryFunctionUnload[incCuboid]];
ClearAll[incCuboid];
incCuboid = Module[{lib, code, name},

    name = "incCuboid";
    
    code=StringJoin[
"
#include \"WolframLibrary.h\"

EXTERN_C DLLEXPORT int "<>name<>"(WolframLibraryData libData, mint Argc, MArgument *Args, MArgument Res)
{
    MTensor acc_    = MArgument_getMTensor(Args[0]);
    MTensor ranges_ = MArgument_getMTensor(Args[1]);

          mint * const acc    = libData->MTensor_getIntegerData(acc_);
    const mint * const ranges = libData->MTensor_getIntegerData(ranges_);
    const mint * const d      = libData->MTensor_getDimensions(acc_);

    const mint n = libData->MTensor_getDimensions(ranges_)[0];

    for( mint l = 0; l < n; ++l )
    {
        const mint r [3][2] = {
            { ranges[6 * l + 0]-1, ranges[6 * l + 1] },
            { ranges[6 * l + 2]-1, ranges[6 * l + 3] },
            { ranges[6 * l + 4]-1, ranges[6 * l + 5] }
        };

        for( mint i = r[0][0]; i < r[0][1]; ++i )
        {
            for( mint j = r[1][0]; j < r[1][1]; ++j )
            {
                for( mint k = r[2][0]; k < r[2][1]; ++k )
                {
                    ++acc[ (i * d[1] + j ) * d[2] + k];
                }
            }
        }
    }

    libData->MTensor_disown(acc_);

    return LIBRARY_NO_ERROR;
}"];

    lib=CreateLibrary[code, name, "Language"->"C++"];

    LibraryFunctionLoad[ lib, name,{{Integer,3,"Shared"},{Integer,3,"Constant"}},"Void"]
];

Here a simple usage example:

Initialization:

size = {1100, 1100, 1100};
accumulator = ConstantArray[0, size];

Running the actual incrementation:

regions = {{{1, 2}, {3, 4}, {5, 6}}, {{2, 3}, {4, 5}, {6, 6}}};
incCuboid[accumulator, regions];

Here is a timing and correctness test:

$HistoryLength = 0;
size = {1100, 1100, 1100};
a = ConstantArray[0, size];
b = ConstantArray[0, size];

n = 10000;
maxdist = 10;

regions = With[{r1 = RandomInteger[{1, 1100 - maxdist}, {n, 3}]},
   Transpose[{r1, r1 + RandomInteger[{1, maxdist}, {n, 3}]}, {3, 1, 2}]
   ];

Do[
  ++a[[r[[1, 1]] ;; r[[1, 2]], r[[2, 1]] ;; r[[2, 2]], r[[3, 1]] ;; r[[3, 2]]]]
  , {r, regions}]; // AbsoluteTiming // First
incCuboid[b, regions]; // AbsoluteTiming // First

0.14132

0.090627

True

Apparently, the speed difference is not that big. And with $HistoryLength = 0 I also don't see any memory leaks. If you write large blocks, then both methods will be roughly at the same speed. The only difference is the uncompiled Do loop after all; the blockwise incrementation code for ++a[[r[[1, 1]] ;; r[[1, 2]], r[[2, 1]] ;; r[[2, 2]], r[[3, 1]] ;; r[[3, 2]]]] should already be done in a compiled library with the nested loop like my i-j-k loop. But for small block sizes, incCuboid for to be faster. For example, it is 20 times faster then Do if maxdist = 4.

Potential improvements

If the blocks are many and very small (and thus fit into L1-cache), then you might get a 20 % improvement if regions is lexicographically sorted. The reason is cache locality. However, sorting regions with Sort will take longer than what you will save. But maybe you can generate regions in a sorted manner?

If the each region is sufficiently large, then it might be worthwhile to parallelize the loop over i. Or if you know that your regions are disjoint, then the loop over l can be parallelized easily, instead.