Fix verify_arg_names when generated function has sizeof != 0 (#1321)

* fix `verify_arg_names` when generated function has sizeof != 0

* fix

* Apply suggestions from code review

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

* bump

* Update lib/ReactantCore/src/ReactantCore.jl

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

* woopsie

* test nested loop

---------

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

Author: jumerckx

Project.toml

@@ -86,7 +86,7 @@ Preferences = "1.4"
 PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
-ReactantCore = "0.1.9"
+ReactantCore = "0.1.10"
 Reactant_jll = "0.0.187"
 ScopedValues = "1.3.0"
 Scratch = "1.2"

lib/ReactantCore/Project.toml

@@ -1,7 +1,7 @@
 name = "ReactantCore"
 uuid = "a3311ec8-5e00-46d5-b541-4f83e724a433"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>"]
-version = "0.1.9"
+version = "0.1.10"
 
 [deps]
 ExpressionExplorer = "21656369-7473-754a-2065-74616d696c43"

lib/ReactantCore/src/ReactantCore.jl

@@ -215,10 +215,14 @@ function trace_for(mod, expr; track_numbers)
             end
         ) for (s, ref) in zip(external_syms, ref_syms)
     ]
+    body_fn_sym = gensym(:body_fn)
+    cond_fn_sym = gensym(:cond_fn)
+    args_sym = gensym(:args)
+    verify_arg_names_sym = gensym(:verify_arg_names)
 
     reactant_code_block = quote
-        let args = $(args_init)
-            cond_fn =
+        let $(args_sym) = $(args_init)
+            $(cond_fn_sym) =
                 $(arg_syms) -> begin
                     $(to_locals...)
                     local num_iters = div($limit - $start, $step, RoundDown)
@@ -227,7 +231,7 @@ function trace_for(mod, expr; track_numbers)
                     )
                     $counter[] < num_iters + 1
                 end
-            body_fn =
+            $(body_fn_sym) =
                 $(arg_syms) -> begin
                     local step_ = $step
                     local start_ = $start
@@ -238,13 +242,17 @@ function trace_for(mod, expr; track_numbers)
                     $counter[].mlir_data = ($counter[] + 1).mlir_data
                     nothing
                 end
-
+            $(verify_arg_names_sym) = if sizeof($(cond_fn_sym)) != 0
+                (Symbol($cond_fn_sym), $(QuoteNode.(args_names.args)...))
+            else
+                ($(QuoteNode.(args_names.args)...),)
+            end
             $(ReactantCore).traced_while(
-                cond_fn,
-                body_fn,
-                args;
+                $(cond_fn_sym),
+                $(body_fn_sym),
+                $(args_sym);
                 track_numbers=$(track_numbers),
-                verify_arg_names=$(QuoteNode(args_names)),
+                verify_arg_names=$(verify_arg_names_sym),
             )
         end
     end

src/TracedUtils.jl

@@ -495,7 +495,7 @@ function prepare_mlir_fn_args(
         stridx = if verify_arg_names isa Nothing
             "arg" * string(path[2])
         else
-            string(verify_arg_names.args[path[2]])
+            string(verify_arg_names[path[2]])
         end
         aval = args[path[2]]
         for (cidx, idx) in enumerate(path[3:end])
@@ -683,7 +683,7 @@ function finalize_mlir_fn(
         for (errs, prev, post) in ((err1, resis, argis), (err2, argis, resis))
             conflicts = setdiff(prev, post)
             for conflict in conflicts
-                stridx = string(verify_arg_names.args[conflict[1]])
+                stridx = string(verify_arg_names[conflict[1]])
                 aval = args[conflict[1]]
                 for (cidx, idx) in enumerate(Base.tail(conflict))
                     if aval isa Array || aval isa Dict

test/control_flow.jl

@@ -852,3 +852,28 @@ end
 
     @test @jit(loop_batched(x_ra)) ≈ loop_batched(x)
 end
+
+function loop!(h_mat::AbstractMatrix, η_mat::AbstractMatrix, H_mat::AbstractMatrix)
+    m,n = size(h_mat)
+    @inbounds @trace for i in 1:m
+        @trace for j in 1:n
+            @allowscalar h_mat[i,j] = η_mat[i,j] + H_mat[i,j]
+        end
+    end
+end
+
+@testset "loop! with nested traced loops and scalar setindex!" begin
+    h = zeros(Float64, 2, 3)
+    η = [1.0 2.0 3.0; 4.0 5.0 6.0]
+    H = [0.5 1.5 2.5; 3.5 4.5 5.5]
+
+    h_ra = Reactant.to_rarray(h)
+    η_ra = Reactant.to_rarray(η)
+    H_ra = Reactant.to_rarray(H)
+
+    @jit loop!(h_ra, η_ra, H_ra)
+
+    loop!(h, η, H)
+
+    @test h_ra ≈ h
+end