Insert call instructions (without any caching)

lib/ReactantCore/src/ReactantCore.jl

@@ -9,6 +9,7 @@ export @trace, within_compile, MissingTracedValue, promote_to_traced
 function is_traced((@nospecialize x::T), seen=Base.IdSet()) where {T}
     if !isprimitivetype(x)
         for fn in fieldnames(T)
+            isdefined(x, fn) || continue
             f = getfield(x, fn)
             if !(f in seen)
                 push!(seen, f)
@@ -247,7 +248,10 @@ function trace_while(expr; track_numbers, mincut, checkpointing, first_arg=nothi
             $body_fn_sym = $(arg_syms) -> begin
                 $(to_locals...)
                 $body
+                temp = Reactant.TRACE_CALLS[]
+                Reactant.TRACE_CALLS[] = false
                 $(from_locals...)
+                Reactant.TRACE_CALLS[] = temp
                 nothing
             end
 
@@ -256,7 +260,6 @@ function trace_while(expr; track_numbers, mincut, checkpointing, first_arg=nothi
             else
                 ($(QuoteNode.(args_names.args)...),)
             end
-
             $(ReactantCore).traced_while(
                 $(cond_fn_sym),
                 $(body_fn_sym),

src/TracedRArray.jl

@@ -847,6 +847,8 @@ function maybe_expand_dims(x::AbstractArray{T,N}, dims) where {T,N}
 end
 
 function Base._cat_t(dims, ::Type{T}, X::TracedRArray...) where {T}
+    temp = Reactant.TRACE_CALLS[]
+    Reactant.TRACE_CALLS[] = false
     dims = dispatch_val(dims)
     @assert dims isa Integer "Support for non-integer dimensions is not implemented yet."
 
@@ -860,7 +862,7 @@ function Base._cat_t(dims, ::Type{T}, X::TracedRArray...) where {T}
     # convert to the target eltype
     X = map(Base.Fix1(TracedUtils.promote_to, TracedRArray{RT,length(shape)}), X)
 
-    return TracedRArray{RT,length(shape)}(
+    result = TracedRArray{RT,length(shape)}(
         (),
         MLIR.IR.result(
             # TODO maybe we should do some conversion?
@@ -873,6 +875,8 @@ function Base._cat_t(dims, ::Type{T}, X::TracedRArray...) where {T}
         ),
         shape,
     )
+    Reactant.TRACE_CALLS[] = temp
+    return result
 end
 
 for (minT, maxT) in Iterators.product((Number, TracedRNumber), (Number, TracedRNumber))

src/TracedUtils.jl

@@ -152,6 +152,11 @@ function set_mlir_data!(x::AnyTracedRArray{T}, data) where {T}
     return x
 end
 
+function set_mlir_data!(x::T, data) where {T}
+    @warn "Setting mlir data on a $T is a no-op."
+    return x
+end
+
 get_ancestor_indices(::TracedRArray, indices) = indices
 get_ancestor_indices(::TracedRArray, indices, args...) = (indices, args...)
 
@@ -279,6 +284,10 @@ function make_mlir_fn(
     optimize_then_pad::Bool=true,
 )
     if sizeof(typeof(f)) != 0 || f isa Base.BroadcastFunction
+        if !isnothing(verify_arg_names)
+            verify_arg_names = (nameof(f), verify_arg_names...)
+        end
+
         mlir_fn_res = make_mlir_fn(
             Reactant.apply,
             (f, args...),
@@ -307,6 +316,7 @@ function make_mlir_fn(
         args,
         name,
         concretein,
+        false, # mutate_args
         toscalar,
         argprefix,
         runtime,
@@ -347,7 +357,7 @@ function make_mlir_fn(
         end
     end
 
-    (func2, traced_result, ret, linear_args, in_tys, linear_results, skipped_results, num_partitions, is_sharded, unique_meshes, mutated_args, global_device_ids) = finalize_mlir_fn(
+    (; func2, traced_result, ret, linear_args, in_tys, linear_results, skipped_results, num_partitions, is_sharded, unique_meshes, mutated_args, global_device_ids) = finalize_mlir_fn(
         result,
         traced_args,
         linear_args,
@@ -409,6 +419,7 @@ function prepare_mlir_fn_args(
     args,
     name,
     concretein,
+    mutate_args,
     toscalar,
     argprefix,
     runtime,
@@ -423,7 +434,11 @@ function prepare_mlir_fn_args(
         @assert !toscalar
         Reactant.ConcreteToTraced
     else
-        Reactant.TracedSetPath
+        if mutate_args
+            Reactant.TracedTrack
+        else
+            Reactant.TracedSetPath
+        end
     end
     fnbody = MLIR.IR.Block(MLIR.IR.Type[], MLIR.IR.Location[])
     MLIR.IR.activate!(fnbody)
@@ -836,9 +851,11 @@ function finalize_mlir_fn(
         MLIR.IR.deactivate!(fnbody)
     end
 
+    f_name = __lookup_unique_name_in_module(mod, name)
+
     func2 = MLIR.IR.block!(MLIR.IR.body(mod)) do
         return MLIR.Dialects.func.func_(;
-            sym_name=__lookup_unique_name_in_module(mod, name),
+            sym_name=f_name,
             function_type=MLIR.IR.FunctionType(in_tys, out_tys),
             body=MLIR.IR.Region(),
             arg_attrs=MLIR.IR.attr(func, "arg_attrs"),
@@ -969,8 +986,9 @@ function finalize_mlir_fn(
     MLIR.API.mlirOperationDestroy(func.operation)
     func.operation = MLIR.API.MlirOperation(C_NULL)
 
-    return (
+    return (;
         func2,
+        f_name,
         traced_result,
         ret,
         linear_args,

src/Tracing.jl

@@ -381,7 +381,10 @@ Base.@nospecializeinfer function traced_type_inner(
             }
         end
         error("Unsupported runtime $runtime")
-    elseif mode == TracedTrack || mode == NoStopTracedTrack || mode == TracedSetPath
+    elseif mode == TracedTrack ||
+        mode == NoStopTracedTrack ||
+        mode == TracedSetPath ||
+        mode == TracedToTypes
         return T
     else
         throw("Abstract RArray cannot be made concrete in mode $mode")
@@ -427,7 +430,10 @@ Base.@nospecializeinfer function traced_type_inner(
             }
         end
         error("Unsupported runtime $runtime")
-    elseif mode == TracedTrack || mode == NoStopTracedTrack || mode == TracedSetPath
+    elseif mode == TracedTrack ||
+        mode == NoStopTracedTrack ||
+        mode == TracedSetPath ||
+        mode == TracedToTypes
         return T
     else
         throw("Abstract RNumber cannot be made concrete in mode $mode")
@@ -921,6 +927,7 @@ function make_tracer(
     return prev
 end
 append_path(@nospecialize(path), i) = (path..., i)
+append_path(::Nothing, i) = nothing
 
 Base.@nospecializeinfer function make_tracer_via_immutable_constructor(
     seen,
@@ -1148,6 +1155,23 @@ function make_tracer(
     )
 end
 
+@static if VERSION >= v"1.11.0"
+    Base.@nospecializeinfer function make_tracer(
+        seen,
+        @nospecialize(prev::Memory),
+        @nospecialize(path),
+        mode;
+        @nospecialize(sharding = Sharding.NoSharding()),
+        kwargs...,
+    )
+        if mode == TracedToTypes
+            return nothing
+        end
+        # TODO: does anything more need to be done here?
+        return prev
+    end
+end
+
 Base.@nospecializeinfer function make_tracer(
     seen,
     @nospecialize(prev::ConcretePJRTArray{T,N}),
@@ -1239,31 +1263,41 @@ Base.@nospecializeinfer function make_tracer(
         throw("Cannot trace existing trace type")
     end
     if mode == TracedToTypes
-        push!(path, MLIR.IR.type(prev.mlir_data))
+        # for TracedRArrays, we check for objectid equality because make_mlir_fn gets rid of duplicate TracedRArrays.
+        # i.e. (a, a) should hash differently than (a, b) when a and b are different TracedRArrays.
+        if haskey(seen, objectid(prev))
+            push!(path, seen[objectid(prev)])
+        else
+            push!(path, MLIR.IR.type(prev.mlir_data))
+            seen[objectid(prev)] = VisitedObject(length(seen) + 1)
+        end
         return nothing
     end
     if mode == TracedTrack
-        TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
+        !isnothing(path) &&
+            TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
         if !haskey(seen, prev)
             return seen[prev] = prev
         end
         return prev
     end
     if mode == NoStopTracedTrack
-        TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
+        !isnothing(path) &&
+            TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
         if !haskey(seen, prev)
             seen[prev] = prev # don't return!
         end
         return prev
     end
     if mode == TracedSetPath
-        TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
+        !isnothing(path) &&
+            TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
         if haskey(seen, prev)
             return seen[prev]
         end
         res = if toscalar
             TracedRNumber{T}((path,), nothing)
-        elseif tobatch !== nothing
+        elseif tobatch !== nothing && prev.shape != tobatch
             error("This should not happen...")
         else
             TracedRArray{T,N}((path,), prev.mlir_data, size(prev))
@@ -1317,25 +1351,35 @@ Base.@nospecializeinfer function make_tracer(
         throw("Cannot trace existing trace type")
     end
     if mode == TracedToTypes
-        push!(path, MLIR.IR.type(prev.mlir_data))
+        # for TracedRArrays, we check for objectid equality because make_mlir_fn gets rid of duplicate TracedRArrays.
+        # i.e. (a, a) should hash differently than (a, b) when a and b are different TracedRArrays.
+        if haskey(seen, objectid(prev))
+            push!(path, seen[objectid(prev)])
+        else
+            push!(path, MLIR.IR.type(prev.mlir_data))
+            seen[objectid(prev)] = VisitedObject(length(seen) + 1)
+        end
         return nothing
     end
     if mode == TracedTrack
-        TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
+        !isnothing(path) &&
+            TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
         if !haskey(seen, prev)
             return seen[prev] = prev
         end
         return prev
     end
     if mode == NoStopTracedTrack
-        TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
+        !isnothing(path) &&
+            TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
         if !haskey(seen, prev)
             seen[prev] = prev # don't return!
         end
         return prev
     end
     if mode == TracedSetPath
-        TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
+        !isnothing(path) &&
+            TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
         if haskey(seen, prev)
             return seen[prev]
         end
@@ -1390,21 +1434,24 @@ Base.@nospecializeinfer function make_tracer(
         throw("Cannot have MissingTracedValue as function call argument.")
     end
     if mode == TracedTrack
-        TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
+        !isnothing(path) &&
+            TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
         if !haskey(seen, prev)
             return seen[prev] = prev
         end
         return prev
     end
     if mode == NoStopTracedTrack
-        TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
+        !isnothing(path) &&
+            TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
         if !haskey(seen, prev)
             seen[prev] = prev # don't return!
         end
         return prev
     end
     if mode == TracedSetPath
-        TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
+        !isnothing(path) &&
+            TracedUtils.set_paths!(prev, (TracedUtils.get_paths(prev)..., path))
         haskey(seen, prev) && return seen[prev]
         res = MissingTracedValue((path,))
         seen[res] = res

src/stdlibs/LinearAlgebra.jl

@@ -297,7 +297,7 @@ function LinearAlgebra._diagm(
     m, n = LinearAlgebra.diagm_size(shape, kv...)
 
     # For repeated indices we need to aggregate the values
-    kv_updated = Dict{Integer,AnyTracedRArray{T,1}}()
+    kv_updated = Dict()
     for (k, v) in kv
         if haskey(kv_updated, k)
             kv_updated[k] = kv_updated[k] + v