Recurse through structs to check is_traced (#931)

* recurse through structs to check is_traced

* test

* Apply suggestions from code review

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

* handle recursive objects

* julia 1.10 doesn't have an IdSet

* Apply suggestions from code review

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

* try out not tracking numbers in Ops.while

* bump reactantcore

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Mosè Giordano <[email protected]>

* bump reactant

* nospecialize is_traced arg

* of course IdSet exists in 1.10 🤦‍♂️

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Mosè Giordano <[email protected]>

Author: 3 people

Project.toml

@@ -85,7 +85,7 @@ Preferences = "1.4"
 PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
-ReactantCore = "0.1.5"
+ReactantCore = "0.1.6"
 Reactant_jll = "0.0.92"
 Scratch = "1.2"
 Sockets = "1.10"

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.5"
+version = "0.1.6"
 
 [deps]
 ExpressionExplorer = "21656369-7473-754a-2065-74616d696c43"

lib/ReactantCore/src/ReactantCore.jl

@@ -6,7 +6,18 @@ using MacroTools: MacroTools
 export @trace, within_compile, MissingTracedValue
 
 # Traits
-is_traced(x) = false
+function is_traced((@nospecialize x::T), seen=Base.IdSet()) where {T}
+    if !isprimitivetype(x)
+        for fn in fieldnames(T)
+            f = getfield(x, fn)
+            if !(f in seen)
+                push!(seen, f)
+                is_traced(f, seen) && return true
+            end
+        end
+    end
+    return false
+end
 
 # New Type signifying that a value is missing
 mutable struct MissingTracedValue

src/Ops.jl

@@ -1724,7 +1724,10 @@ end
     traced_args = Vector{Any}(undef, N)
     for i in 1:N
         @inbounds traced_args[i] = Reactant.make_tracer(
-            seen_args, args[i], (), Reactant.NoStopTracedTrack; track_numbers=Number
+            seen_args,
+            args[i],
+            (),
+            Reactant.NoStopTracedTrack, #; track_numbers=Number
         )
     end
 

src/TracedRArray.jl

@@ -22,6 +22,7 @@ using ..TracedUtils: TracedUtils, get_mlir_data, set_mlir_data!, materialize_tra
 using ReactantCore: ReactantCore
 using GPUArraysCore: GPUArraysCore, @allowscalar
 
+ReactantCore.is_traced(::TracedRArray, seen) = true
 ReactantCore.is_traced(::TracedRArray) = true
 
 Base.strides(x::TracedRArray) = Base.size_to_strides(1, size(x)...)

src/TracedRNumber.jl

@@ -4,6 +4,7 @@ using ..Reactant:
     Reactant, TracedRNumber, TracedRArray, TracedUtils, Ops, MLIR, unwrapped_eltype
 using ReactantCore
 
+ReactantCore.is_traced(::TracedRNumber, seen) = true
 ReactantCore.is_traced(::TracedRNumber) = true
 
 Base.getindex(a::TracedRNumber{T}) where {T} = a

test/control_flow.jl

@@ -784,3 +784,30 @@ end
 
     @test (@jit ternary_max(a, b)) == 2
 end
+
+mutable struct MaybeTraced
+    x
+end
+
+@testset "is_traced of struct" begin
+    containstraced = MaybeTraced(
+        MaybeTraced(Reactant.TracedRArray{Float64,1}((), nothing, (3,)))
+    )
+    @test Reactant.ReactantCore.is_traced(containstraced)
+
+    doesnotcontaintraced = MaybeTraced(MaybeTraced(3))
+    @test !Reactant.ReactantCore.is_traced(doesnotcontaintraced)
+
+    recursivetraced = MaybeTraced((
+        1,
+        "string",
+        MaybeTraced(nothing),
+        MaybeTraced(Reactant.TracedRArray{Float64,1}((), nothing, (3,))),
+    ))
+    recursivetraced.x[3].x = recursivetraced
+    @test Reactant.ReactantCore.is_traced(recursivetraced)
+
+    recursivenottraced = MaybeTraced((1, "string", MaybeTraced(nothing)))
+    recursivenottraced.x[3].x = recursivenottraced
+    @test !Reactant.ReactantCore.is_traced(recursivenottraced)
+end