Updating loopsets to better represent aliasing of certain symbols and memory locations, and prep to allow indices to have constant integers.

Author: chriselrod

src/constructors.jl

@@ -18,3 +18,13 @@ macro avx(q)
     esc(lower(LoopSet(q)))
 end
 
+#=
+@generated function vmaterialize(
+    dest::AbstractArray{T,N}, bc::BC
+) where {T,N,BC <: Base.Broadcast.Broadcasted}
+    # we have an N dimensional loop.
+    # need to construct the LoopSet
+    
+end
+=#
+

src/determinestrategy.jl

@@ -1,5 +1,6 @@
 
 # TODO: FIXME for general case
+# wrong for transposed matrices, and certain views/SubArrays.
 unitstride(op, s) = first(loopdependencies(op)) === s
 
 function cost(op::Operation, unrolled::Symbol, Wshift::Int, size_T::Int = op.elementbytes)

src/graphs.jl

@@ -82,7 +82,12 @@ struct LoopSet
     # stridesets::Dict{ShortVector{Symbol},ShortVector{Symbol}}
     preamble::Expr # TODO: add preamble to lowering
     includedarrays::Vector{Tuple{Symbol,Int}}
+    syms_aliasing_refs::Vector{Symbol} # O(N) search is faster at small sizes
+    refs_aliasing_syms::Vector{ArrayReference}
+    # sym_to_ref_aliases::Dict{Symbol,ArrayReference}
+    # ref_to_sym_aliases::Dict{ArrayReference,Symbol}
 end
+
 function includesarray(ls::LoopSet, array::Symbol)
     for (a,i) ∈ ls.includedarrays
         a === array && return i
@@ -97,7 +102,11 @@ function LoopSet()
         Int[],
         LoopOrder(),
         Expr(:block,),
-        Tuple{Symbol,Int}[]
+        Tuple{Symbol,Int}[],
+        Symbol[],
+        ArrayReference[]
+        # Dict{Symbol,ArrayReference}()
+        # Dict{ArrayReference,Symbol}()
     )
 end
 num_loops(ls::LoopSet) = length(ls.loops)
@@ -209,17 +218,33 @@ function add_vptr!(ls::LoopSet, indexed::Symbol, id::Int)
 end
 
 function add_load!(
-    ls::LoopSet, var::Symbol, indexed::Symbol, indices::AbstractVector, elementbytes::Int = 8
+    ls::LoopSet, var::Symbol, ref::ArrayReference, elementbytes::Int = 8
 )
-    op = Operation( length(operations(ls)), var, elementbytes, :getindex, memload, indices, [indexed], NOPARENTS )
+    if ref.loaded[] == true
+        op = getop(ls, var)
+        @assert var === op.variable
+        return op
+    end
+    push!(ls.syms_aliasing_refs, var)
+    push!(ls.refs_aliasing_syms, ref)
+    ref.loaded[] = true
+    # ls.sym_to_ref_aliases[ var ] = ref
+    # ls.ref_to_sym_aliases[ ref ] = var
+    op = Operation(
+        length(operations(ls)), var, elementbytes,
+        :getindex, memload, loopdependencies(ref),
+        NODEPENDENCY, NOPARENTS, ref
+    )
     add_vptr!(ls, indexed, identifier(op))
     pushop!(ls, op, var)
 end
 function add_load_ref!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::Int = 8)
-    add_load!(ls, var, ex.args[1], @view(ex.args[2:end]), elementbytes)
+    ref = ref_from_ref(ex)
+    add_load!(ls, var, ref, elementbytes)
 end
 function add_load_getindex!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::Int = 8)
-    add_load!(ls, var, ex.args[2], @view(ex.args[3:end]), elementbytes)
+    ref = ref_from_getindex(ex)
+    add_load!(ls, var, ref, elementbytes)
 end
 function instruction(x)
     x isa Symbol ? x : last(x.args).value
@@ -274,20 +299,31 @@ function pushparent!(parents::Vector{Operation}, deps::Vector{Symbol}, reducedde
 end
 function maybe_cse_load!(ls::LoopSet, expr::Expr, elementbytes::Int = 8)
     if expr.head === :ref
-        array = first(expr.args)::Symbol
-        args = @view expr.args[2:end]
+        offset = 0
+        # array = first(expr.args)::Symbol
+        # args = @view expr.args[2:end]
+        # ref = ref_from_ref(expr)
     elseif expr.head === :call && first(expr.args) === :getindex
-        array = (expr.args[2])::Symbol
-        args = @view expr.args[3:end]
+        offset = 1
+        # array = (expr.args[2])::Symbol
+        # args = @view expr.args[3:end]
+        # ref = ref_from_getindex(expr)
     else
         return add_operation!(ls, gensym(:temporary), expr, elementbytes)
     end
-    id = includesarray(ls, array)
-    if id > 0
-        ls.operations[id]
-    else
+    ref = ArrayReference( ex.args[1+offset], @view(ex.args[2+offset:end]) )::ArrayReference
+    id = findfirst(r -> r == ref, ls.refs_aliasing_syms)
+    if id === nothing
         add_load!( ls, gensym(:temporary), array, args, elementbytes )
+    else
+        ls.syms_aliasing_refs[id]
     end
+    # id = includesarray(ls, array)
+    # if id > 0
+        # ls.operations[id]
+    # else
+        # add_load!( ls, gensym(:temporary), array, args, elementbytes )
+    # end
 end
 function add_parent!(
     parents::Vector{Operation}, deps::Vector{Symbol}, reduceddeps::Vector{Symbol}, ls::LoopSet, var, elementbytes::Int = 8
@@ -325,7 +361,7 @@ function add_reduction_update_parent!(
     parent.instruction === Symbol("##CONSTANT##") && push!(ls.outer_reductions, identifier(op))
     pushop!(ls, op, var) # note this overwrites the entry in the operations dict, but not the vector
 end
-function add_compute!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::Int = 8)
+function add_compute!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::Int = 8, ref = nothing)
     @assert ex.head === :call
     instr = instruction(first(ex.args))::Symbol
     args = @view(ex.args[2:end])
@@ -338,6 +374,9 @@ function add_compute!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::Int = 8)
         if arg === var
             reduction = true
             add_reduction!(parents, deps, reduceddeps, ls, arg, elementbytes)
+        elseif ref == arg
+            reduction = true
+            add_load!(ls, var, ref, elementbytes)
         else
             add_parent!(parents, deps, reduceddeps, ls, arg, elementbytes)
         end
@@ -402,9 +441,12 @@ function Base.push!(ls::LoopSet, ex::Expr, elementbytes::Int = 8)
             if RHS isa Symbol
                 lrhs = RHS
             elseif RHS isa Expr
+                # need to check of LHS appears in RHS
                 # assign RHS to lrhs
-                lrhs = gensym(:RHS)
-                add_operation!(ls, lrhs, RHS, elementbytes)
+                ref = ArrayReference(LHS)
+                id = findfirst(r -> r == ref, ls.refs_aliasing_syms)
+                lrhs = id === nothing ? gensym(:RHS) : ls.syms_aliasing_refs[id]
+                add_operation!(ls, lrhs, RHS, elementbytes, ref)
             end
             add_store_ref!(ls, lrhs, LHS, elementbytes)
         else

src/lowering.jl

@@ -522,7 +522,7 @@ end
 function gc_preserve(ls::LoopSet, q::Expr)
     length(ls.includedarrays) == 0 && return q # is this even possible?
     gcp = Expr(:macrocall, Expr(:(.), :GC, QuoteNode(Symbol("@preserve"))), LineNumberNode(@__LINE__, @__FILE__))
-    for (array,i) ∈ ls.includedarrays
+    for (array,_) ∈ ls.includedarrays
         push!(gcp.args, array)
     end
     push!(q.args, nothing)
@@ -697,4 +697,5 @@ function lower(ls::LoopSet)
 end
 
 Base.convert(::Type{Expr}, ls::LoopSet) = lower(ls)
+Base.show(io::IO, ls::LoopSet) = println(io, lower(ls))
 

src/operations.jl

@@ -1,3 +1,52 @@
+struct ArrayReference
+    array::Symbol
+    ref::Vector{Union{Symbol,Int}}
+    loaded::Base.RefValue{Bool}
+end
+function Base.hash(x::ArrayReference, h::UInt)
+    @inbounds for n ∈ eachindex(x)
+        h = hash(x.ref[n], h)
+    end
+    hash(x.array, h)
+end
+loopdependencies(ref::ArrayReference) = filter(i -> i isa Symbol, ref.ref)
+function Base.isequal(x::ArrayReference, y::ArrayReference)
+    x.array === y.array || return false
+    nrefs = length(x.ref)
+    nrefs == length(y.ref) || return false
+    all(n -> x.ref[n] === y.ref[n], 1:nrefs)
+    # for n ∈ 1:nrefs
+        # x.ref[n] === y.ref[n] || return false
+    # end
+    # true
+end
+
+Base.:(==)(x::ArrayReference, y::ArrayReference) = isequal(x, y)
+
+function ref_from_ref(ex::Expr)
+    ArrayReference( ex.args[1], @view(ex.args[2:end]), Ref(false) )
+end
+function ref_from_getindex(ex::Expr)
+    ArrayReference( ex.args[2], @view(ex.args[3:end]), Ref(false) )
+end
+function ArrayReference(ex::Expr)
+    ex.head === :ref ? ref_from_ref(ex) : ref_from_getindex(ex)
+end
+function Base.:(==)(x::ArrayReference, y::Expr)
+    if y.head === :ref
+        isequal(x, ref_from_ref(y))
+    elseif y.head === :call && first(y.args) === :getindex
+        isequal(x, ref_from_getindex(y))
+    else
+        false
+    end
+end
+Base.:(==)(x::ArrayReference, y) = false
+
+
+
+# Avoid memory allocations by accessing this
+const NOTAREFERENCE = ArrayReference(Symbol(""), Union{Symbol,Int}[])
 
 @enum OperationType begin
     constant
@@ -17,8 +66,7 @@ symbolic metadata contains info on direct dependencies / placement within loop.
 if isload(op) -> Symbol(:vptr_, first(op.reduced_deps))
 if istore(op) -> Symbol(:vptr_, op.variable)
 is how we access the memory.
-If numerical_metadata[i] == -1
-Symbol(:stride_, op.variable, :_, op.symbolic_metadata[i])
+
 is the stride for loop index
 symbolic_metadata[i]
 """
@@ -28,9 +76,10 @@ struct Operation
     elementbytes::Int
     instruction::Symbol
     node_type::OperationType
-    dependencies::Vector{Symbol}
+    dependencies::Vector{Symbol}#::Vector{Symbol}
     reduced_deps::Vector{Symbol}
     parents::Vector{Operation}
+    ref::ArrayReference
     # children::Vector{Operation}
     # numerical_metadata::Vector{Int} # stride of -1 indicates dynamic
     # symbolic_metadata::Vector{Symbol}
@@ -42,19 +91,21 @@ struct Operation
         node_type,
         dependencies = Symbol[],
         reduced_deps = Symbol[],
-        parents = Operation[]
+        parents = Operation[],
+        ref::ArrayReference = NOTAREFERENCE
     )
         new(
             identifier, variable, elementbytes, instruction, node_type,
             convert(Vector{Symbol},dependencies),
             convert(Vector{Symbol},reduced_deps),
-            convert(Vector{Operation},parents)
+            convert(Vector{Operation},parents),
+            ref
         )
     end
 end
 
  # negligible save on allocations for operations that don't need these (eg, constants).
-const NODEPENDENCY = Symbol[]
+const NODEPENDENCY = Union{Symbol,Int}[]
 const NOPARENTS = Operation[]