feat: support kwargs in macros

Author: avik-pal

src/Compiler.jl

@@ -588,6 +588,7 @@ function compile_mlir!(
     optimize::Union{Bool,Symbol}=true,
     no_nan::Bool=false,
     backend="gpu",
+    fn_kwargs=(),
 )
     # Explicitly don't use block! to avoid creating a closure, which creates
     # both compile-time and relocatability issues
@@ -597,7 +598,7 @@ function compile_mlir!(
     activate_callcache!(callcache)
 
     mlir_fn_res = try
-        Reactant.TracedUtils.make_mlir_fn(f, args, (), "main", true)
+        Reactant.TracedUtils.make_mlir_fn(f, args, fn_kwargs, "main", true)
     finally
         deactivate_callcache!(callcache)
         MLIR.IR.deactivate!(MLIR.IR.body(mod))
@@ -984,6 +985,7 @@ function compile_call_expr(mod, compiler, options, args...)
     call = only(args)
     f_symbol = gensym(:f)
     args_symbol = gensym(:args)
+    kwargs_symbol = gensym(:kwargs)
     compiled_symbol = gensym(:compiled)
 
     if Meta.isexpr(call, :call)
@@ -999,19 +1001,32 @@ function compile_call_expr(mod, compiler, options, args...)
         else
             :($(fname))
         end
-        args_rhs = Expr(:tuple, call.args[2:end]...)
+        args_rhs = call.args[2:end]
     elseif Meta.isexpr(call, :(.), 2) && Meta.isexpr(call.args[2], :tuple)
         fname = :($(Base.Broadcast.BroadcastFunction)($(call.args[1])))
-        args_rhs = only(call.args[2:end])
+        @assert length(call.args[2:end]) == 1
+        args_rhs = (call.args[2],)
     else
         error("Invalid function call: $(call)")
     end
 
+    if length(args_rhs) ≥ 1 && Meta.isexpr(args_rhs[1], :parameters)
+        kwargs_rhs = args_rhs[1].args
+        args_rhs = args_rhs[2:end]
+    else
+        kwargs_rhs = []
+    end
+    args_rhs = Expr(:tuple, args_rhs...)
+
     return quote
         $(f_symbol) = $(fname)
         $(args_symbol) = $(args_rhs)
+        $(kwargs_symbol) = (; $(kwargs_rhs...))
         $(compiled_symbol) = $(compiler)(
-            $(f_symbol), $(args_symbol); $(Expr.(:kw, keys(options), values(options))...)
+            $(f_symbol),
+            $(args_symbol);
+            fn_kwargs=$(kwargs_symbol),
+            $(Expr.(:kw, keys(options), values(options))...),
         )
     end,
     (; compiled=compiled_symbol, args=args_symbol)

src/TracedUtils.jl

@@ -251,14 +251,19 @@ function make_mlir_fn(
     # Explicitly don't use block! to avoid creating a closure, which creates
     # both compile-time and relocatability issues
     MLIR.IR.activate!(fnbody)
+
     result = try
         for (i, arg) in enumerate(linear_args)
             raw_arg = MLIR.IR.argument(fnbody, i)
             row_maj_arg = do_transpose ? transpose_val(raw_arg) : raw_arg
             set_mlir_data!(arg, row_maj_arg)
         end
 
-        Reactant.call_with_reactant(f, traced_args...)
+        if isempty(kwargs)
+            Reactant.call_with_reactant(f, traced_args...)
+        else
+            Reactant.call_with_reactant(Core.kwcall, kwargs, f, traced_args...)
+        end
     finally
         MLIR.IR.deactivate!(fnbody)
     end