Support for elementary bounds inference
skeqiqevian 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.