Bounds checking too conservative on sub-arrays

[hessammehr]

The following minimal example complains about the indices of a being out of range even though 1 <= a[1, j] <= 3. I wonder if bounds checking could be modified to handle scenarios like this?

using Tullio

a = [1 2 3
     4 5 6]
b = [1, 2, 3]
@tullio x := b[a[1, j]]

# "extrema of index a[1, j] must fit within b"

[mcabbott]

Yes it's a bit conservative, I hadn't thought of this particular case but I think it's easy to address.

This one is a bug I meant to figure out:

@tullio x := b[a[1, j] + 1]  # junk

Cases where the index doesn't span the whole array are tricky, and still too conservative:

@tullio x := b[vec(a)[i+0]]  i in 1:3  # error

[mcabbott]

Another case: scatter operations have their own code for bounds checks, so this complains about a column of inds which it won't read:

inds = rand(1:10, 333, 100); inds[:,50] .= 99;
vals = randn(333);
out = zeros(10,10)
@tullio out[inds[r,1], inds[r,100]] += vals[r]

[hessammehr]

I get the impression Tullio eventually needs to implement/duplicate a lot of Julia, not just indexing logic but also figuring out dataflow. I wonder if operating on IR/SSA rather than AST input would lead to fewer edge cases.

[mcabbott]

That's an interesting thought. But the naiive execution has a bounds check on every getindex, whereas Tullio wants to check extrema(inds) once, up front. Is there a way to go from one to the other?

[hessammehr]

I must admit I'm probably out of my depth here but as a rough sketch perhaps something like this:

julia> struct Index
       end

julia> f(a,b,c) = begin
           i = Index()
           j = Index()
           @inbounds a[j] = b[i, j+1] + c[i, j-1]
       end
f (generic function with 2 methods)

julia> @code_lowered f(a,b,c)
CodeInfo(
1 ─       i = Main.Index()
│         j = Main.Index()
│         $(Expr(:inbounds, true))
│   %4  = i
│   %5  = j + 1
│   %6  = Base.getindex(b, %4, %5)
│   %7  = i
│   %8  = j - 1
│   %9  = Base.getindex(c, %7, %8)
│   %10 = %6 + %9
│         Base.setindex!(a, %10, j)
│         val = %10
│         $(Expr(:inbounds, :pop))
└──       return val
)

I haven't looked at packages like IRTools or Mjolnir too closely but wonder if inferring loop bounds from the getindex and setindex ops here would be more robust.

[mcabbott]

Now I had another look at https://github.com/shashi/ArrayMeta.jl, which is perhaps a more Julian approach than mine... although it's still parsing the expression itself, it just builds up another tree structure & passes this to a generated function to make the loops at compile-time, instead of making them directly in the macro during parsing.

[hessammehr]

Interesting, I wonder to what extent Meta.@lower overlaps with the functionality implemented in ArrayMeta.

julia> Meta.@lower a[j] = b[i, j+1] + c[i, j-1]
:($(Expr(:thunk, CodeInfo(
    @ none within `top-level scope'
1 ─ %1 = j + 1
│   %2 = Base.getindex(b, i, %1)
│   %3 = j - 1
│   %4 = Base.getindex(c, i, %3)
│   %5 = %2 + %4
│        Base.setindex!(a, %5, j)
└──      return %5
))))