Halide Func reshape

Question

I want to convert the dimensions of a Halide Func. For example consider the following,
func1(x, y, z) = some operation
I know the range of each dimension of the function. Lets assume they are
x = [0, 3], y = [0, 3], z = [0, 15] Now if I want to convert the func to 1 dimension, something like this
flatten(z * 16 + y * 4 + x) = func1(x, y, z)
But the above definition is not valid as the function definition does not have a pure Var. Is there anyway I can do this in Halide?

Answer 1

You'll need to express it as a pure function of one variable, e.g.:

flatten(x) = func1(x % 4, (x % 16) / 4, x / 16);

You can simplify away much of this added addressing arithmetic by scheduling flatten appropriately, for example:

// Require the bounds realized to be a multiple of 16.
flatten.align_bounds(x, 16);

// Split the loops at the "critical points" of the addressing arithmetic.
flatten
    .unroll(x, 4, TailStrategy::RoundUp)
    .split(x, xo, xi, 4, TailStrategy::RoundUp);