Using groupBy/groupingBy/aggregate to sum into smaller buckets in parallel? - kotlin

Question

I've got a collection of "stuff", and I'd like to sum it into smaller buckets. (In my particular case, I'm downsampling a luma channel of an image by 8x.)
I'd like it to be as fast as possible on your average multi-core android device, which I think means coroutine-per-bucket. (because there isn't any reason to play with IntAdders if I don't have to)
The naive linear solution:
val SCALE = 8
image.planes[0].buffer.toByteArray().forEachIndexed { index, byte ->
val x1 = index % image.width
val y1 = index / image.width
val x2 = x1 / SCALE
val y2 = y1 / SCALE
val quadIdx = y2 * (image.width / SCALE) + x2
summedQuadLum[quadIdx] += (byte.toInt() and 0xFF)
}
That isn't great - needs to pre-declare the summedQuadLum collection, and doesn't have any chance of parallel work.
I'd love to use groupBy, or groupingBy? or aggregate?) but those all seem to use the values to determine the new keys, and I need to use the key to determine the new keys. I think the least overhead is withIndex which could be done as
val thumbSums = bufferArray.withIndex().groupingBy { (idx, _) ->
val x1 = idx % previewImageDimension.width
val y1 = idx / previewImageDimension.width
val x2 = x1 / SCALE
val y2 = y1 / SCALE
y2 * (previewImageDimension.width / SCALE) + x2
}.aggregate { _, acc: Int?, (_, lum), _ ->
(acc ?: 0) + (lum.toInt() and 0xFF)
}.values.toIntArray()
Much better, it is really close - if I could figure out how to sum each bucket in a coroutine, I think it would be as good as can be expected.

Answer 1

So after groupingBy we have a Grouping object, which we can use to aggregate values. It's important to notice the grouping itself has not been done yet, we basically have a description how to group the values and an iterator of the original array. From here we have a few options:
Create a Channel from the iterator and launch a few worker coroutines to consume it in parallel. Channels support fan-out, so every item in the source is processed by one worker only. The problem here is all the workers need to update different items in the resulting array, so synchronization is required and that's where it gets tricky and likely inefficient.
To avoid multiple workers to write to the same item, we need to tell each of them what items to process. That mean either each of the worker should process all the items, picking only suitable or we should precalculate the groups in advance and feed the workers with the groups. Both approaches have pretty much the same performance as the serial algorithm, so do not make any sense.
So to parallelize it efficiently we want to avoid having a shared mutable state, because it requires synchronization. Obviously we don't want to precalculate the groups also.
My suggestion here is to come from another side - instead of mapping original array to sampled one, let's map sampled array to the original. So we say
This approaches makes each value to be calculated independently by one worker, so no synchronization needed. Now we can implement it like this:
suspend fun sample() {
val asyncFactor = 8
val src = Image(bufferArray, width)
val dst = Image(src.width / SCALE, src.height / SCALE)
val chunkSize = dst.sizeBytes / asyncFactor
val jobs = Array(asyncFactor) { idx ->
async(Dispatchers.Default) {
val chunkStartIdx = chunkSize * idx
val chunkEndIdxExclusive = min(chunkStartIdx + chunkSize, dst.sizeBytes)
calculateSampledImageForIndexes(src, dst, chunkStartIdx, chunkEndIdxExclusive, SCALE)
}
}
awaitAll(*jobs)
}
private fun calculateSampledImageForIndexes(src: Image, dst: Image, startIdx: Int, exclusiveEndIdx: Int, scaleFactor: Int) {
for (i in startIdx until exclusiveEndIdx) {
val destX = i % dst.width
val destY = i / dst.width
val srcX = destX * scaleFactor
val srcY = destY * scaleFactor
var sum = 0
for (xi in 0 until scaleFactor) {
for (yi in 0 until scaleFactor) {
sum += src[srcX + xi, srcY + yi]
}
}
dst[destX, destY] = sum / (scaleFactor * scaleFactor)
}
}
Where Image is a convenient wrapper around the image data buffer:
class Image(val buffer: ByteArray, val width: Int) {
val height = buffer.size / width
val sizeBytes get() = buffer.size
constructor(w: Int, h: Int) : this(ByteArray(w * h), w)
operator fun get(x: Int, y: Int): Byte = buffer[clampX(x) * width + clampY(y)]
operator fun set(x: Int, y: Int, value: Int) {
buffer[x * width + y] = (value and 0xFF).toByte()
}
private fun clampX(x: Int) = max(min(x, width), 0)
private fun clampY(y: Int) = max(min(y, height), 0)
}
Also, with this approach you can easily implement many image processing functions, which based on convolution operation, like blur and edge detection.