Compile parametric methods that contain llvmcall (fixes #112, fixes #288) (#289)

Author: timholy

Project.toml

@@ -17,7 +17,8 @@ Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Mmap = "a63ad114-7e13-5084-954f-fe012c677804"
 SHA = "ea8e919c-243c-51af-8825-aaa63cd721ce"
+Tensors = "48a634ad-e948-5137-8d70-aa71f2a747f4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "Distributed", "Dates", "SHA", "Mmap"]
+test = ["Test", "Dates", "Distributed", "Mmap", "SHA", "Tensors"]

src/construct.jl

@@ -165,6 +165,11 @@ function prepare_framecode(method::Method, @nospecialize(argtypes); enter_genera
                 generator = false
             end
         end
+        # Currenly, our strategy to deal with llvmcall can't handle parametric functions
+        # (the "mini interpreter" runs in module scope, not method scope)
+        if !isempty(lenv) && (hasarg(isequal(:llvmcall), code.code) || hasarg(a->is_global_ref(a, Base, :llvmcall), code.code))
+            return Compiled()
+        end
         framecode = FrameCode(method, code; generator=generator)
         if is_generated(method) && !enter_generated
             genframedict[(method, argtypes)] = framecode

src/utils.jl

@@ -87,6 +87,19 @@ function scan_ssa_use!(used::BitSet, @nospecialize(stmt))
     end
 end
 
+function hasarg(predicate, args)
+    predicate(args) && return true
+    for a in args
+        predicate(a) && return true
+        if isa(a, Expr)
+            hasarg(predicate, a.args) && return true
+        elseif isa(a, QuoteNode)
+            predicate(a.value) && return true
+        end
+    end
+    return false
+end
+
 ## Predicates
 
 is_goto_node(@nospecialize(node)) = isa(node, GotoNode) || isexpr(node, :gotoifnot)

test/interpret.jl

@@ -513,6 +513,36 @@ function f_mmap()
 end
 @interpret f_mmap()
 
+# parametric llvmcall (issues #112 and #288)
+module VecTest
+    using Tensors
+    Vec{N,T} = NTuple{N,VecElement{T}}
+    # The following test mimic SIMD.jl
+    const _llvmtypes = Dict{DataType, String}(
+        Float64 => "double",
+        Float32 => "float",
+        Int32 => "i32",
+        Int64 => "i64"
+    )
+    @generated function vecadd(x::Vec{N, T}, y::Vec{N, T}) where {N, T}
+        llvmT = _llvmtypes[T]
+        func = T <: AbstractFloat ? "fadd" : "add"
+        exp = """
+        %3 = $(func) <$(N) x $(llvmT)> %0, %1
+        ret <$(N) x $(llvmT)> %3
+        """
+        return quote
+            Base.@_inline_meta
+            Core.getfield(Base, :llvmcall)($exp, Vec{$N, $T}, Tuple{Vec{$N, $T}, Vec{$N, $T}}, x, y)
+        end
+    end
+    f() = 1.0 * one(Tensor{2,3})
+end
+let a = (VecElement{Float64}(1.0), VecElement{Float64}(2.0))
+    @test @interpret(VecTest.vecadd(a, a)) == VecTest.vecadd(a, a)
+end
+@test @interpret(VecTest.f()) == [1 0 0; 0 1 0; 0 0 1]
+
 # Test exception type for undefined variables
 f() = s = s + 1
 @test_throws UndefVarError @interpret f()