nicer error for type mutation in traced loop body (#1035)

Author: jumerckx

lib/ReactantCore/src/ReactantCore.jl

@@ -231,7 +231,7 @@ function trace_for(mod, expr; track_numbers)
                 end
 
             $(ReactantCore).traced_while(
-                cond_fn, body_fn, args; track_numbers=$(track_numbers)
+                cond_fn, body_fn, args; track_numbers=$(track_numbers), verify_arg_names=$(QuoteNode(args_init))
             )
         end
     end

src/ControlFlow.jl

@@ -9,7 +9,8 @@ function ReactantCore.traced_call(f::Function, args...)
 end
 
 function ReactantCore.traced_while(
-    cond_fn::CFn, body_fn::BFn, args; track_numbers=Number
+    cond_fn::CFn, body_fn::BFn, args; track_numbers=Number, verify_arg_names=nothing
 ) where {CFn,BFn}
-    return Ops.while_loop(cond_fn, body_fn, args...; track_numbers)
+    @warn verify_arg_names
+    return Ops.while_loop(cond_fn, body_fn, args...; track_numbers, verify_arg_names)
 end

src/Ops.jl

@@ -1746,7 +1746,7 @@ use [`MLIR.Dialects.stablehlo.dynamic_slice`](@ref) instead.
 end
 
 @noinline function while_loop(
-    cond_fn::CFn, body_fn::BFn, args...; track_numbers
+    cond_fn::CFn, body_fn::BFn, args...; track_numbers, verify_arg_names=nothing
 ) where {CFn,BFn}
     # TODO: detect and prevent mutation within the condition
 
@@ -1780,6 +1780,7 @@ end
             do_transpose=false,
         ).f
 
+    @warn verify_arg_names
     body_fn_compiled =
         Reactant.TracedUtils.make_mlir_fn(
             body_fn,
@@ -1790,6 +1791,7 @@ end
             return_dialect=:stablehlo,
             args_in_result=:none,
             do_transpose=false,
+            verify_arg_names
         ).f
 
     cond_reg = Reactant.TracedUtils.__take_region(cond_fn_compiled)

src/TracedUtils.jl

@@ -189,6 +189,7 @@ function make_mlir_fn(
     input_shardings=nothing,  # This is not meant to be used by the user.
     output_shardings=nothing, # This is not meant to be used by the user.
     runtime=nothing,
+    verify_arg_names=nothing,
 )
     if sizeof(typeof(f)) != 0 || f isa Base.BroadcastFunction
         mlir_fn_res = make_mlir_fn(
@@ -347,6 +348,18 @@ function make_mlir_fn(
     if args_in_result == :mutated
         append!(linear_results, linear_args[mutated_args])
     end
+    if !isnothing(verify_arg_names) && typeof.(linear_args) != typeof.(linear_results)
+        @assert length(linear_args) <= length(linear_results)
+        argis = first.(get_argidx.(linear_args))
+        resis = Set(getindex.(get_residx.(linear_results), Ref(2)))
+        # this can be more efficient
+        conflicts = setdiff(resis, argis)
+        @assert !isempty(conflicts) "Expected to have some conflicts, but none were found."
+
+        error("""Types do not match between function arguments and results.
+        The following arguments should be traced: $(join(verify_arg_names.args[collect(conflicts)], ", "))
+        """)
+    end
 
     out_tys = if do_transpose
         [transpose_ty(Ops.mlir_type(arg)) for arg in linear_results]