Halide - sort buffer/function in one dimension

Question

I am currently using Halide with the use of a generator and ahead of time compilation.

Somewhere in the pipeline I have a 3D buffer with limited extent (typically 3-6 values) in one of the dimensions. I would like to sort the values in that dimension. When I skip the processing at the beginning of the pipe line, it looks somewhat like this:

Input <  Buffer<uint16_t>> input  { "input" , 2}; // Dimensions: (y, x)
Input <  uint8_t>        > sizeZ  { "sizeZ"    }; // Size in Z-dimension
Output<  Buffer<uint16_t>> output { "output", 3}; // Dimensions: (z, y, x)

Var x,y,z;

Func input3D(z,y,x) = input(y,z+x*sizeZ);

output = 'sort input3D on Z dimension'.

I would be most helped if some sorting functionality is already available in Halide (is that so?). An alternative would be to call an external C implementation to sort all values in that dimension and assign them to the output buffer. That would be something like:

output(:, y, x) = external_sort(input(:, y, x))

In which I used the Python notation to express all elements in Z dimension. Is something like this possible in Halide?

Answer 1

There is an example of calling an external C function to sort a Halide func in our tests here: https://github.com/halide/Halide/blob/master/test/correctness/extern_sort.cpp

General purpose sorting algorithms cannot be expressed in Halide. However, sorting networks for small vectors can be. See here for an example of bitonic sorting: https://github.com/halide/Halide/blob/master/test/correctness/sort_exprs.cpp