From 8b737dc80228a24c079a3fcdda39e56c76c99fd1 Mon Sep 17 00:00:00 2001
From: Tim Holy <tim.holy@gmail.com>
Date: Mon, 16 Mar 2020 16:21:22 -0500
Subject: [PATCH] WIP: shorten compile time with nospecialize

---
 src/LoopVectorization.jl   |  1 +
 src/broadcast.jl           | 18 ++++++++++++------
 src/reconstruct_loopset.jl |  5 +++--
 src/typeutils.jl           | 21 +++++++++++++++++++++
 test/broadcast.jl          |  2 ++
 5 files changed, 39 insertions(+), 8 deletions(-)
 create mode 100644 src/typeutils.jl

diff --git a/src/LoopVectorization.jl b/src/LoopVectorization.jl
index 544bcfd72..4ceabebc7 100644
--- a/src/LoopVectorization.jl
+++ b/src/LoopVectorization.jl
@@ -25,6 +25,7 @@ export LowDimArray, stridedpointer, vectorizable,
 
 include("vectorizationbase_extensions.jl")
 include("predicates.jl")
+include("typeutils.jl")
 include("map.jl")
 include("filter.jl")
 include("costs.jl")
diff --git a/src/broadcast.jl b/src/broadcast.jl
index a4645d43c..8df17213e 100644
--- a/src/broadcast.jl
+++ b/src/broadcast.jl
@@ -60,8 +60,9 @@ the final operation.
 # TODO: Need to make this handle A or B being (1 or 2)-D broadcast objects.
 function add_broadcast!(
     ls::LoopSet, mC::Symbol, bcname::Symbol, loopsyms::Vector{Symbol},
-    ::Type{Product{A,B}}, elementbytes::Int
-) where {A, B}
+    @nospecialize(AB::Type{<:Product}), elementbytes::Int
+)
+    A, B = AB.parameters
     K = gensym(:K)
     mA = gensym(:Aₘₖ)
     mB = gensym(:Bₖₙ)
@@ -126,8 +127,9 @@ function LowDimArray{D}(data::A) where {D,T,N,A <: AbstractArray{T,N}}
 end
 function add_broadcast!(
     ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol},
-    ::Type{<:LowDimArray{D,T,N}}, elementbytes::Int
-) where {D,T,N}
+    @nospecialize(LDA::Type{<:LowDimArray}), elementbytes::Int
+)
+    D, T, N = LDA.parameters
     fulldims = Symbol[loopsyms[n] for n ∈ 1:N if D[n]]
     ref = ArrayReference(bcname, fulldims)
     add_simple_load!(ls, destname, ref, elementbytes, true, false )::Operation
@@ -180,9 +182,13 @@ function add_broadcast!(
 end
 function add_broadcast!(
     ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol},
-    ::Type{Broadcasted{S,Nothing,F,A}},
+    BT::Type{<:Broadcasted},
     elementbytes::Int
-) where {N,S<:Base.Broadcast.AbstractArrayStyle{N},F,A}
+)
+    @nospecialize BT
+    S,_,F,A = BT.parameters
+    N = abstractparameters(S, Broadcast.AbstractArrayStyle)[1]
+
     instr = get(FUNCTIONSYMBOLS, F) do
         # f = gensym(:func)
         # pushpreamble!(ls, Expr(:(=), f, Expr(:(.), bcname, QuoteNode(:f))))
diff --git a/src/reconstruct_loopset.jl b/src/reconstruct_loopset.jl
index e1e66db4c..0381c2846 100644
--- a/src/reconstruct_loopset.jl
+++ b/src/reconstruct_loopset.jl
@@ -83,8 +83,9 @@ function add_mref!(ls::LoopSet, ar::ArrayReferenceMeta, i::Int, ::Type{OffsetStr
 end
 
 function add_mref!(
-    ls::LoopSet, ar::ArrayReferenceMeta, i::Int, ::Type{S}
-) where {T, X <: Tuple, S <: VectorizationBase.AbstractStaticStridedPointer{T,X}}
+    ls::LoopSet, ar::ArrayReferenceMeta, i::Int, @nospecialize(S::Type{<:VectorizationBase.AbstractStaticStridedPointer})
+)
+    T, X = abstractparameters(S, VectorizationBase.AbstractStaticStridedPointer)
     if last(X.parameters)::Int == 1
         pushvarg′!(ls, ar, i)
     else
diff --git a/src/typeutils.jl b/src/typeutils.jl
new file mode 100644
index 000000000..8fef6660e
--- /dev/null
+++ b/src/typeutils.jl
@@ -0,0 +1,21 @@
+# @nospecialize  # don't specialize anything until @specialize
+
+"""
+    params = abstractparameters(T::Type, AT::Type)
+
+For a type `T` which is a subtype of `AT{params...}`, return `params`.
+The purpose of this function is to allow one to use `@nospecialize` on type-arguments
+and still extract the parameters of the corresponding abstract type.
+"""
+function abstractparameters(T::Type, AT::Type)
+    @nospecialize T AT
+    @assert T <: AT
+    Tst = supertype(T)
+    while Tst <: AT
+        T = Tst
+        Tst = supertype(T)
+    end
+    return T.parameters
+end
+
+# @specialize
diff --git a/test/broadcast.jl b/test/broadcast.jl
index 17f4e60c3..c381d2e69 100644
--- a/test/broadcast.jl
+++ b/test/broadcast.jl
@@ -1,3 +1,5 @@
+using LoopVectorization, LinearAlgebra, Test
+
 @testset "broadcast" begin
     M, N = 37, 47
     # M = 77;