Revert "fix: traced_getfield move to TracedUtils"

This reverts commit f35d911.
EnzymeAD · Dec 16, 2024 · e2d7377 · e2d7377
1 parent d9c167f
commit e2d7377
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 31 deletions.
diff --git a/src/Compiler.jl b/src/Compiler.jl
@@ -13,7 +13,11 @@ import ..Reactant:
     TracedToConcrete,
     append_path,
     TracedType
-import ..TracedUtils: TracedUtils, traced_getfield
+
+@inline traced_getfield(@nospecialize(obj), field) = Base.getfield(obj, field)
+@inline traced_getfield(
+    @nospecialize(obj::AbstractArray{<:Union{ConcreteRNumber,ConcreteRArray}}), field
+) = Base.getindex(obj, field)
 
 function create_result(tocopy::T, path, result_stores) where {T}
     if !isstructtype(typeof(tocopy))

diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl
@@ -1,3 +1,20 @@
+struct XLAArray{T,N} <: RArray{T,N}
+    # size::NTuple{N,Int}
+end
+
+mutable struct ConcreteRArray{T,N} <: RArray{T,N}
+    data::XLA.AsyncBuffer
+    #   data::XLAArray{T, N}
+    shape::NTuple{N,Int}
+end
+
+const WrappedConcreteRArray{T,N} = WrappedArray{T,N,ConcreteRArray,ConcreteRArray{T,N}}
+const AnyConcreteRArray{T,N} = Union{ConcreteRArray{T,N},WrappedConcreteRArray{T,N}}
+
+mutable struct ConcreteRNumber{T} <: RNumber{T}
+    data::XLA.AsyncBuffer
+end
+
 function ConcreteRNumber{T}(
     data::T2; client=XLA.default_backend[], idx=XLA.default_device_idx[], device=nothing
 ) where {T<:Number,T2<:Number}

diff --git a/src/Interpreter.jl b/src/Interpreter.jl
@@ -202,7 +202,7 @@ function set_act!(inp, path, reverse, tostore; emptypath=false)
     end
 
     for p in path
-        x = TracedUtils.traced_getfield(x, p)
+        x = traced_getfield(x, p)
     end
 
     #if inp isa Enzyme.Active || !reverse

diff --git a/src/Reactant.jl b/src/Reactant.jl
@@ -114,25 +114,6 @@ mutable struct TracedRNumber{T} <: RNumber{T}
     end
 end
 
-struct XLAArray{T,N} <: RArray{T,N}
-    # size::NTuple{N,Int}
-end
-
-mutable struct ConcreteRArray{T,N} <: RArray{T,N}
-    data::XLA.AsyncBuffer
-    #   data::XLAArray{T, N}
-    shape::NTuple{N,Int}
-end
-
-const WrappedConcreteRArray{T,N} = WrappedArray{T,N,ConcreteRArray,ConcreteRArray{T,N}}
-const AnyConcreteRArray{T,N} = Union{ConcreteRArray{T,N},WrappedConcreteRArray{T,N}}
-
-mutable struct ConcreteRNumber{T} <: RNumber{T}
-    data::XLA.AsyncBuffer
-end
-
-const TracedType = Union{TracedRArray,TracedRNumber,MissingTracedValue}
-
 include("Ops.jl")
 include("TracedUtils.jl")
 
@@ -143,6 +124,8 @@ include("ConcreteRArray.jl")
 
 include("linear_algebra.jl")
 
+const TracedType = Union{TracedRArray,TracedRNumber,MissingTracedValue}
+
 include("ControlFlow.jl")
 include("Tracing.jl")
 include("Compiler.jl")
@@ -163,7 +146,7 @@ function Enzyme.make_zero(
     return res
 end
 
-using .Compiler: @compile, @code_hlo, @jit, create_result, compile
+using .Compiler: @compile, @code_hlo, @jit, traced_getfield, create_result, compile
 export ConcreteRArray, ConcreteRNumber, @compile, @code_hlo, @jit, @trace
 
 const registry = Ref{MLIR.IR.DialectRegistry}()

diff --git a/src/TracedUtils.jl b/src/TracedUtils.jl
@@ -14,18 +14,12 @@ using ..Reactant:
     AnyTracedRArray,
     MissingTracedValue,
     OrderedIdDict,
-    ConcreteRArray,
-    ConcreteRNumber
+    Compiler
 import ..Reactant
 import ..Reactant.MLIR
 import ..ReactantPrimitive
 import ..Ops
 
-@inline traced_getfield(@nospecialize(obj), field) = Base.getfield(obj, field)
-@inline traced_getfield(
-    @nospecialize(obj::AbstractArray{<:Union{ConcreteRNumber,ConcreteRArray}}), field
-) = Base.getindex(obj, field)
-
 materialize_traced_array(x::TracedRArray) = x
 materialize_traced_array(x::WrappedTracedRArray) = x[axes(x)...]
 function materialize_traced_array(
@@ -330,7 +324,7 @@ end
 
 function push_val!(ad_inputs, x, path)
     for p in path
-        x = traced_getfield(x, p)
+        x = Compiler.traced_getfield(x, p)
     end
     x = x.mlir_data
     return push!(ad_inputs, x)
@@ -350,7 +344,7 @@ end
 
 function set!(x, path, tostore; emptypath=false)
     for p in path
-        x = traced_getfield(x, p)
+        x = Compiler.traced_getfield(x, p)
     end
 
     x.mlir_data = tostore