Open hessammehr opened 4 years ago
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
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]
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.
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?
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.
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.
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
))))
The following minimal example complains about the indices of
a
being out of range even though1 <= a[1, j] <= 3
. I wonder if bounds checking could be modified to handle scenarios like this?