Fixed bug where Amn = A[m,n] caused an error when m and n are constants.

Author: chriselrod

src/add_compute.jl

@@ -51,7 +51,14 @@ function add_parent!(
     vparents::Vector{Operation}, deps::Vector{Symbol}, reduceddeps::Vector{Symbol}, ls::LoopSet, var, elementbytes::Int, position::Int
 )
     parent = if var isa Symbol
+        # if var === :kern_1_1
+        #     @show operations(ls) ls.preamble_symsym
+        # end
         opp = getop(ls, var, elementbytes)
+        # if var === :kern_1_1
+        #     @show operations(ls) ls.preamble_symsym
+        # end 
+        # @show var opp first(operations(ls)) opp === first(operations(ls))
         if iscompute(opp) && instruction(opp).instr === :identity && length(loopdependencies(opp)) < position && isone(length(parents(opp))) && name(opp) === name(first(parents(opp)))
             first(parents(opp))
         else
@@ -220,6 +227,7 @@ function add_compute!(
     deps = Symbol[]
     reduceddeps = Symbol[]
     reduction_ind = 0
+    # @show ex first(operations(ls)) === getop(ls, :kern_1_1, elementbytes) first(operations(ls)) getop(ls, :kern_1_1, elementbytes)
     for (ind,arg) ∈ enumerate(args)
         if var === arg
             reduction_ind = ind

src/add_constants.jl

@@ -1,7 +1,8 @@
 function add_constant!(ls::LoopSet, var::Symbol, elementbytes::Int)
     op = Operation(length(operations(ls)), var, elementbytes, LOOPCONSTANT, constant, NODEPENDENCY, Symbol[], NOPARENTS)
-    pushpreamble!(ls, op, var)
-    pushop!(ls, op, var)
+    rop = pushop!(ls, op, var)
+    rop === op && pushpreamble!(ls, op, var)
+    rop
 end
 # function add_constant!(ls::LoopSet, var, elementbytes::Int = 8)
 #     sym = gensym(:loopconstant)
@@ -12,6 +13,8 @@ function add_constant!(ls::LoopSet, var::Number, elementbytes::Int = 8)
     op = Operation(length(operations(ls)), gensym(:loopconstnumber), elementbytes, LOOPCONSTANT, constant, NODEPENDENCY, Symbol[], NOPARENTS)
     ops = operations(ls)
     typ = var isa Integer ? HardInt : HardFloat
+    rop = pushop!(ls, op)
+    rop !== op && return rop
     if iszero(var)
         for (id,typ_) ∈ ls.preamble_zeros
             (instruction(ops[id]) === LOOPCONSTANT && typ == typ_) && return ops[id]
@@ -28,7 +31,7 @@ function add_constant!(ls::LoopSet, var::Number, elementbytes::Int = 8)
         end
         push!(ls.preamble_symfloat, (identifier(op), var))
     end
-    pushop!(ls, op)
+    rop
 end
 function add_constant!(ls::LoopSet, mpref::ArrayReferenceMetaPosition, elementbytes::Int)
     op = Operation(length(operations(ls)), varname(mpref), elementbytes, LOOPCONSTANT, constant, NODEPENDENCY, Symbol[], NOPARENTS, mpref.mref)

src/add_loads.jl

@@ -23,7 +23,7 @@ end
 function add_load!(
     ls::LoopSet, mpref::ArrayReferenceMetaPosition, elementbytes::Int
 )
-    length(mpref.loopdependencies) == 0 && return add_constant!(ls, mpref, elementbytes)
+    iszero(length(mpref.loopdependencies)) && return add_constant!(ls, mpref, elementbytes)
     op = Operation( ls, varname(mpref), elementbytes, :getindex, memload, mpref )
     add_load!(ls, op, true, false)
 end

src/graphs.jl

@@ -522,9 +522,9 @@ function add_loop!(ls::LoopSet, loop::Loop, itersym::Symbol = loop.itersymbol)
     nothing
 end
 
-function instruction(x)
-    x isa Symbol ? x : last(x.args).value
-end
+# function instruction(x)
+#     x isa Symbol ? x : last(x.args).value
+# end
 # instruction(ls::LoopSet, f::Symbol) = instruction!(ls, f)
 function instruction!(ls::LoopSet, x::Expr)
     x isa Symbol && return x

src/lower_compute.jl

@@ -1,18 +1,11 @@
 
+
 function load_constrained(op, u₁loop, u₂loop, innermost_loop, forprefetch = false)
     loopdeps = loopdependencies(op)
     dependsonu₁ = u₁loop ∈ loopdeps
-    if u₂loop === Symbol("##undefined##")
-        if forprefetch
-            dependsonu₁ || return false
-        end
-        # unrolleddeps = [ u₁loop ]
-    else        
-        dependsonu₂ = u₂loop ∈ loopdeps
-        if forprefetch
-            (dependsonu₁ & dependsonu₂) || return false
-        end
-        # unrolleddeps = [ u₁loop, u₂loop ]
+    dependsonu₂ = u₂loop ∈ loopdeps
+    if forprefetch
+        (dependsonu₁ & dependsonu₂) || return false
     end
     unrolleddeps = Symbol[]
     dependsonu₁ && push!(unrolleddeps, u₁loop)

src/operations.jl

@@ -206,7 +206,7 @@ function matches(op1::Operation, op2::Operation)
     op1.instruction === op2.instruction || return false
     op1.node_type == op2.node_type || return false
     if isconstant(op1)
-        return false
+        return iszero(length(loopdependencies(op1))) && iszero(length(loopdependencies(op2))) && (mangledvar(op1) === mangledvar(op2))
     end
     op1.dependencies == op2.dependencies || return false
     op2.reduced_deps == op2.reduced_deps || return false

src/reconstruct_loopset.jl

@@ -306,7 +306,7 @@ function add_op!(
     # If it's a CartesianIndex add or subtract, we may have to add multiple operations
     expanded = expandedv[i]# isexpanded(ls, ops, nopsv, i)
     opoffsets = ls.operation_offsets
-    offsets = ls.loopsymbol_offsets
+    # offsets = ls.loopsymbol_offsets
     optyp = optype(os)
     if !expanded
         op = Operation(

test/gemm.jl

@@ -47,7 +47,22 @@
     if LoopVectorization.REGISTER_COUNT != 8
         @test LoopVectorization.choose_order(lsAmulB3) == (Symbol[:n,:m,:k], :m, :n, :m, Unum, Tnum)
     end
-    
+    if LoopVectorization.REGISTER_COUNT != 8
+        for (fA,fB,v,Un,Tn) ∈ [(identity,identity,:m,Unum,Tnum),(adjoint,identity,:k,Unumt,Tnumt),(identity,adjoint,:m,Unum,Tnum),(adjoint,adjoint,:n,Unum,Tnum)]
+            A = fA(rand(2,2))
+            B = fB(rand(2,2))
+            C = similar(A)
+            ls = LoopVectorization.@avx_debug for m ∈ axes(A,1), n ∈ axes(B,2)
+                ΔCₘₙ = zero(eltype(C))
+                for k ∈ axes(A,2)
+                    ΔCₘₙ += A[m,k] * B[k,n]
+                end
+                C[m,n] += ΔCₘₙ
+            end
+            (m, n) = v === :n ? (:n, :m) : (:m, :n)
+            @test LoopVectorization.choose_order(ls) == (Symbol[:n,:m,:k], m, n, v, Un, Tn)
+        end
+    end
     function AmulB!(C, A, B)
         C .= 0
         for k ∈ axes(A,2), j ∈ axes(B,2)

test/offsetarrays.jl

@@ -110,18 +110,19 @@ using LoopVectorization.VectorizationBase: StaticUnitRange
         # Manually unpack the OffsetArray
         @avx for j in rng2, i in rng1
             tmp_0 = zero(eltype(out))
-            Base.Cartesian.@nexprs 3 jk -> Base.Cartesian.@nexprs 3 ik -> tmp_{ik+(jk-1)*3} = A[(ik-2)+i,(jk-2)+j] * kern_ik_jk + tmp_{ik+(jk-1)*3-1}
+            Base.Cartesian.@nexprs 3 jk -> Base.Cartesian.@nexprs 3 ik -> tmp_{ik+(jk-1)*3} = A[(ik-2)+i,(jk-2) + j*1] * kern_ik_jk + tmp_{ik+(jk-1)*3-1}
             out[i,j] = tmp_9
         end
         out
     end
     function avx2dunrolled2x2!(out::AbstractMatrix, A::AbstractMatrix, kern::SizedOffsetMatrix{T,-1,1,-1,1}) where {T}
         rng1,  rng2  = axes(out)
-        Base.Cartesian.@nexprs 3 jk -> Base.Cartesian.@nexprs 3 ik -> kern_ik_jk = kern[ik-2,jk-2]
         # Manually unpack the OffsetArray
         @avx unroll=(2,2) for j in rng2, i in rng1
+            Base.Cartesian.@nexprs 3 jk -> Base.Cartesian.@nexprs 3 ik -> kern_ik_jk = kern[ik - 2, jk + (-2)]
             tmp_0 = zero(eltype(out))
-            Base.Cartesian.@nexprs 3 jk -> Base.Cartesian.@nexprs 3 ik -> tmp_{ik+(jk-1)*3} = A[i+(ik-2),j+(jk-2)] * kern_ik_jk + tmp_{ik+(jk-1)*3-1}
+            j1 = j * 1
+            Base.Cartesian.@nexprs 3 jk -> Base.Cartesian.@nexprs 3 ik -> tmp_{ik+(jk-1)*3} = A[i + (ik-2), (jk-2) + j1] * kern_ik_jk + tmp_{ik+(jk-1)*3-1}
             out[i,j] = tmp_9
         end
         out
@@ -132,7 +133,7 @@ using LoopVectorization.VectorizationBase: StaticUnitRange
         # Manually unpack the OffsetArray
         @avx unroll=(3,3) for j in rng2, i in rng1
             tmp_0 = zero(eltype(out))
-            Base.Cartesian.@nexprs 3 jk -> Base.Cartesian.@nexprs 3 ik -> tmp_{ik+(jk-1)*3} = A[i+(ik-2),j+(jk-2)] * kern_ik_jk + tmp_{ik+(jk-1)*3-1}
+            Base.Cartesian.@nexprs 3 jk -> Base.Cartesian.@nexprs 3 ik -> tmp_{ik+(jk-1)*3} = A[(ik-2) + i, j*1 + (jk-2)] * kern_ik_jk + tmp_{ik+(jk-1)*3-1}
             out[i,j] = tmp_9
         end
         out