Compiler optimization challenge: Bayer pattern image unpacking

I’ve found an interesting example of a slightly convoluted for-loop that might be a good showcase of Julia’s ability to generate efficient code. I need help optimizing this code, though, because I’m not sure I have either the fastest or the most elegant code. Pretty sure I got neither!

This is a quite relevant function, with a non-trivial operation, and I’m actually taking this from real-world Python project

The way this works is: every 4 x UInt16 pixels in the output array come from 5 bytes of the input array, with the 5th byte providing the lowest two bits in every group of pixels.

Here’s the code I’ve come up with, first a direct porting of the Python code, with Numpy-style vectorial operations, and then an explicit for-loop. It already results in a speed-up of 80% in my machine. How much further can it be improved?

using BenchmarkTools

function pythonicconv5_8to4_16!(output_data::Array{UInt16}, pixel_bytes_2d::Array{UInt8})
    # First, set aside cohort 4: the bytes that will be unpacked into the low bits to go with the first 4 bytes
    cohort_4 = UInt16.(pixel_bytes_2d[:, 5:5:end])
    for byte_cohort_index in 1:4
        # High bits come from the input array,
        # shifted left by two bits to make room for the low 2-bits which will come from the 5th byte
        high_bits = UInt16.(pixel_bytes_2d[:, byte_cohort_index:5:end]) .<< 2

        # Now process bits from cohort 4 and unpack the appropriate ones to be our low 2 bits:
        # Shift the bits over so that the relevant ones are in the rightmost (lowest 2 bits) position
        # eg. for byte 1, 0b00001100 -> 0b11
        shifted_bits_from_cohort_4 = cohort_4 .>> ((byte_cohort_index-1) * 2)
        # Mask the relevant ones (the lowest 2)
        masked_low_bits = shifted_bits_from_cohort_4 .& 0x0003

        # Finally, put the masked low bits into their place in the output array
        output_data[:, byte_cohort_index:4:end] .= high_bits .| masked_low_bits

function conv5_8to4_16!(output_data::Array{UInt16}, pixel_bytes_2d::Array{UInt8})
    ll = length(pixel_bytes_2d)
    for n in 5:5:ll
        low = pixel_bytes_2d[n]

        output_data[m-3] = UInt16(pixel_bytes_2d[n-4]) << 2
        output_data[m-2] = UInt16(pixel_bytes_2d[n-3]) << 2
        output_data[m-1] = UInt16(pixel_bytes_2d[n-2]) << 2
        output_data[m] = UInt16(pixel_bytes_2d[n-1]) << 2

        output_data[m-3] |= low & 0x03
        output_data[m-2] |= (low >> 2) & 0x03
        output_data[m-1] |= (low >> 4) & 0x03
        output_data[m] |= (low >> 6) & 0x03
        m += 4

len = 1024
px = rand(UInt8, len, len*5÷4)
out = Array{UInt16, 2}(undef, len, len)

@btime conv5_8to4_16!(out, px)
@btime pythonicconv5_8to4_16!(out, px)