Search code examples
halide

Express time-step loops in Halide

I am trying to implement the following loop nest in Halide

for (t = 0; t < TSTEPS; t++)
  for (i = 1; i < N - 1; i++)
     A[i] = (A[i-1] + A[i] + A[i + 1])/3;

But I couldn't figure out how to express the "t" loop since it doesn't contribute in the memory accesses. I went through the Halide tutorials but I couldn't find an example that illustrate such pattern.

Can you please show me an example of how to express such free loops in Halide.

Thanks.

Solution

  • Use a two dimensional RDom:

    Var x{"x"};

Param<int> N, TSTEPS;
RDom i {{{0, N}, {0, TSTEPS}}, "i"};

Func A{"A"};

A(x) = ...;
A(i.x) = (A(i.x-1) + A(i.x) + A(i.x + 1))/3;

    The TSTEPS loop variable isn't mentioned in the update rule, but the loop will still be emitted. i.x refers to the N dimension, and i.y would refer to the TSTEPS dimension.

    See this tutorial for more information: https://halide-lang.org/tutorials/tutorial_lesson_09_update_definitions.html