Fix vararg (#158)

* Fix vararg

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: wsmoses

src/utils.jl

@@ -87,17 +87,24 @@ function make_mlir_fn(f, args, kwargs, name="main", concretein=true; toscalar=fa
         end
 
         # TODO replace with `Base.invoke_within` if julia#52964 lands
-        ir = first(only(
+        ir, ty = only(
             # TODO fix it for kwargs
             Base.code_ircode(f, map(typeof, traced_args); interp),
-        ))
-
+        )
         oc = Core.OpaqueClosure(ir)
 
         if f === Reactant.apply
             oc(traced_args[1], (traced_args[2:end]...,))
         else
-            oc(traced_args...)
+            if length(traced_args) + 1 != length(ir.argtypes)
+                @assert ir.argtypes[end] <: Tuple
+                oc(
+                    traced_args[1:(length(ir.argtypes) - 2)]...,
+                    (traced_args[(length(ir.argtypes) - 1):end]...,),
+                )
+            else
+                oc(traced_args...)
+            end
         end
     end
 

test/basic.jl

@@ -105,6 +105,17 @@ bcast_cos(x) = cos.(x)
     @test r ≈ cos.(x)
 end
 
+f_var(args...) = sum(args)
+
+@testset "Vararg" begin
+    x = Reactant.to_rarray(ones(3))
+    y = Reactant.to_rarray(3 * ones(3))
+    z = Reactant.to_rarray(2.6 * ones(3))
+
+    f2 = @compile f_var(x, y, z)
+    @test f2(x, y, z) ≈ [6.6, 6.6, 6.6]
+end
+
 function sumcos(x)
     return sum(cos.(x))
 end