Merge pull request #88 from vchuravy/vc/vacall

correctly call pure function with varargs

src/context.jl

@@ -36,11 +36,22 @@ function transform(ctx, ref)
 
     # don't overdub pure functions
     if CI.pure
-        Cassette.insert_statements!(CI.code, CI.codelocs,
-          (x, i) -> i == 1 ?  2 : nothing,
-          (x, i) -> i == 1 ? [
-              Expr(:call, Expr(:nooverdub, Core.SlotNumber(1)), (Core.SlotNumber(i) for i in 2:ref.method.nargs)...),
-              Expr(:return, Core.SSAValue(i))] : nothing)
+        n_method_args = Int(ref.method.nargs)
+        if ref.method.isva
+            Cassette.insert_statements!(CI.code, CI.codelocs,
+                (x, i) -> i == 1 ?  3 : nothing,
+                (x, i) -> i == 1 ? [
+                    # this could run into troubles when the function is @pure f(x...) since then n_method_args==2, but this seems to work sofar.
+                    Expr(:call, Expr(:nooverdub, GlobalRef(Core, :tuple)), (Core.SlotNumber(i) for i in 2:(n_method_args-1))...),
+                    Expr(:call, Expr(:nooverdub, GlobalRef(Core, :_apply)), Core.SlotNumber(1), Core.SSAValue(i), Core.SlotNumber(n_method_args)),
+                    Expr(:return, Core.SSAValue(i+1))] : nothing)
+        else
+            Cassette.insert_statements!(CI.code, CI.codelocs,
+                (x, i) -> i == 1 ?  2 : nothing,
+                (x, i) -> i == 1 ? [
+                    Expr(:call, Expr(:nooverdub, Core.SlotNumber(1)), (Core.SlotNumber(i) for i in 2:n_method_args)...)
+                    Expr(:return, Core.SSAValue(i))] : nothing)
+        end
         CI.ssavaluetypes = length(CI.code)
         return CI
     end

test/gpuenv/Project.toml

@@ -2,6 +2,7 @@
 CUDAnative = "be33ccc6-a3ff-5ff2-a52e-74243cff1e17"
 CuArrays = "3a865a2d-5b23-5a0f-bc46-62713ec82fae"
 GPUifyLoops = "ba82f77b-6841-5d2e-bd9f-4daf811aec27"
+InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/test.jl

@@ -1,5 +1,6 @@
 using GPUifyLoops
 using Test
+using InteractiveUtils
 
 function kernel(A)
     @loop for i in (1:size(A,1);
@@ -236,3 +237,29 @@ end
   @launch CPU() threads=(3,3) kernel_MArray!(A)
   @launch CPU() threads=(3,3) kernel_similar_MArray!(A)
 end
+
+# For the next three test we check that the _apply got inlined correctly
+Base.@pure pure_f1(x, y) = x + y
+let
+    CI, rt = @code_typed GPUifyLoops.Cassette.overdub(GPUifyLoops.ctx, pure_f1, 1, 2)
+    expr = CI.code[end-1]
+    @test expr.head === :call || expr.head === :invoke
+    @test expr.args[1] === GlobalRef(Base, :add_int)
+end
+
+Base.@pure pure_f2(x, ys...) = x + sum(ys)
+let
+    CI, rt = @code_typed GPUifyLoops.Cassette.overdub(GPUifyLoops.ctx, pure_f2, 1, 2, 3, 4)
+    expr = CI.code[end-1]
+    @test expr.head === :call || expr.head === :invoke
+    @test expr.args[1] === GlobalRef(Base, :add_int)
+end
+
+Base.@pure pure_f3(ys...) = sum(ys)
+let
+    CI, rt = @code_typed GPUifyLoops.Cassette.overdub(GPUifyLoops.ctx, pure_f3, 1, 2, 3, 4)
+    expr = CI.code[end-1]
+    @test expr.head === :call || expr.head === :invoke
+    @test expr.args[1] === GlobalRef(Base, :add_int)
+end
+

test/testenv/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 GPUifyLoops = "ba82f77b-6841-5d2e-bd9f-4daf811aec27"
+InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"