exo-lang/exo

Support for elementary bounds inference

Opened this issue · 0 comments

For a particular scope of code, we would like to determine all the possible reads/writes that may be made. For example, in the following code snippet:

x: i8[N+1]
for io in seq(0, N/32):
  for ii in seq(0, 32):
    x[32*io + ii] = 0.0
    x[32*io + ii + 1] = 0.0

we would like to infer that within the ii loop, x is read at [32*io, 32*io + 32]. This functionality would be useful in

  • automating stage_mem so that we don't have to manually specify the window we want to stage.
  • implementing Halide's version of compute_root. Exo can currently implement very simple versions, but in general needs bounds inference to reasonably abstract.

I think Exo already has the capabilities to do this. The task can be reduced to bounding individual index expressions (32*io + ii) within a particular scope (the ii loop). Essentially, we want to decompose the index expression into base + offset component where base can't be bounded within integer constants, but offset can. In the example above, base = 32 * io and offset = ii, which we can bound to be in [0, 32). Our output is of the form [base + lo, base + hi).

This can be achieved by using liveness analysis to determine which index variables are defined outside of the scope (and hence we can't reason about) and which index variables are defined within the scope (and can be reasoned about). After that, we can leverage existing interval analysis from simplify to reason about the latter.