Add a few convenience utilities

The main advantage these have is a small improvement in performance
due to not needing to construct new objects.

For Dicts of types, this switches to IdDict. This is both faster 
(particularly on construction), and safe for precompilation (see 
https://docs.julialang.org/en/latest/manual/modules/#Module-initialization-and-precompilation-1).

Author: timholy

src/LoopVectorization.jl

@@ -23,6 +23,7 @@ export LowDimArray, stridedpointer, vectorizable,
 
 
 include("vectorizationbase_extensions.jl")
+include("predicates.jl")
 include("map.jl")
 include("filter.jl")
 include("costs.jl")

src/constructors.jl

@@ -30,9 +30,9 @@ function substitute_broadcast(q::Expr, mod::Symbol)
         f = first(ciₙargs)
         if ciₙ.head === :(=)
             push!(ex.args, Expr(:(=), f, syms[((ciₙargs[2])::Core.SSAValue).id]))
-        elseif f === GlobalRef(Base, :materialize!)
+        elseif isglobalref(f, Base, :materialize!)
             add_ci_call!(ex, lv(:vmaterialize!), ciₙargs, syms, n, mod)
-        elseif f === GlobalRef(Base, :materialize)
+        elseif isglobalref(f, Base, :materialize)
             add_ci_call!(ex, lv(:vmaterialize), ciₙargs, syms, n, mod)
         else
             add_ci_call!(ex, f, ciₙargs, syms, n)

src/costs.jl

@@ -289,7 +289,7 @@ function isreductcombineinstr(instr::Symbol)
 end
 isreductcombineinstr(instr::Instruction) = isreductcombineinstr(instr.instr)
 
-const FUNCTIONSYMBOLS = Dict{Type{<:Function},Instruction}(
+const FUNCTIONSYMBOLS = IdDict{Type{<:Function},Instruction}(
     typeof(+) => :(+),
     typeof(SIMDPirates.vadd) => :(+),
     typeof(SIMDPirates.vadd!) => :(+),

src/graphs.jl

@@ -110,7 +110,7 @@ function terminatecondition(
     looprange(loop, 1, mangledname)
 end
 
-# load/compute/store × isunroled × istiled × pre/post loop × Loop number
+# load/compute/store × isunrolled × istiled × pre/post loop × Loop number
 struct LoopOrder <: AbstractArray{Vector{Operation},5}
     oporder::Vector{Vector{Operation}}
     loopnames::Vector{Symbol}
@@ -402,7 +402,7 @@ function register_single_loop!(ls::LoopSet, looprange::Expr)
             L = add_loop_bound!(ls, itersym, Expr(:call, :first, N), false)
             U = add_loop_bound!(ls, itersym, Expr(:call, :last, N), true)
             Loop(itersym, L, U)
-        elseif f === :OneTo || f == Expr(:(.), :Base, QuoteNode(:OneTo))
+        elseif f === :OneTo || isscopedname(f, :Base, :OneTo)
             otN = r.args[2]
             if otN isa Integer
                 Loop(itersym, 0, otN)

src/predicates.jl

@@ -0,0 +1,23 @@
+"""
+    isscopedname(ex, modpath, name::Symbol)
+
+Test whether expression `ex` is a module-scoped `name`. Both of these return `true`:
+
+```
+isscopedname(:(Base.OneTo), :Base, :OneTo)
+isscopedname(:(Base.Checked.checked_add), (:Base, :Checked), :checked_add)
+```
+"""
+function isscopedname(ex, modpath, name::Symbol)
+    isexpr(ex, :(.), 2) && (a = ex.args[2]; isa(a, QuoteNode) && a.value === name) && hasscope(ex.args[1], modpath)
+end
+hasscope(modex, mod::Symbol) = modex === mod
+hasscope(modex, mod::Tuple{Symbol}) = hasscope(modex, mod[1])
+hasscope(modex, modpath::Tuple{Vararg{Symbol}}) = isscopedname(modex, Base.front(modpath), modpath[end])
+
+"""
+    isglobalref(g, mod, name)
+
+Return true if `g` is equal to `GlobalRef(mod, name)`.
+"""
+isglobalref(g, mod, name) = isa(g, GlobalRef) && g.mod === mod && g.name === name

test/runtests.jl

@@ -18,8 +18,10 @@ end
 
     @time include("printmethods.jl")
 
+    @time include("utils.jl")
+
     @time include("offsetarrays.jl")
-    
+
     @time include("map.jl")
 
     @time include("filter.jl")

test/utils.jl

@@ -0,0 +1,15 @@
+using LoopVectorization
+using Test
+
+@testset "Utilities" begin
+    @test LoopVectorization.isscopedname(:(Base.OneTo), :Base, :OneTo)
+    @test LoopVectorization.isscopedname(:(A.B.C.D), (:A, :B, :C), :D)
+    @test !LoopVectorization.isscopedname(:(A.B.D),  (:A, :B, :C), :D)
+    @test !LoopVectorization.isscopedname(:(A.B.C.D), (:A, :B, :C), :E)
+    @test !LoopVectorization.isscopedname(:hello,  :Base, :OneTo)
+
+    @test LoopVectorization.isglobalref(GlobalRef(Base, :getindex), Base, :getindex)
+    @test !LoopVectorization.isglobalref(GlobalRef(Base, :getindex), Base, :setindex!)
+    @test !LoopVectorization.isglobalref(GlobalRef(Core, :getindex), Base, :getindex)
+    @test !LoopVectorization.isglobalref(:getindex, Base, :getindex)
+end