Reimplement getfield to allow customization on external types

EnzymeAD · Jul 23, 2024 · dffa473 · dffa473
1 parent 9eabd69
commit dffa473
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/src/Compiler.jl b/src/Compiler.jl
@@ -269,6 +269,8 @@ macro code_hlo(options, maybe_call=nothing)
     end
 end
 
+traced_getfield(obj, field) = Base.getfield(obj, field)
+
 function compile(f, args; pipeline_options="", client=nothing)
     N = length(args)
     ctx = MLIR.IR.Context()
@@ -333,7 +335,7 @@ function compile(f, args; pipeline_options="", client=nothing)
             end
             res = :(args[$(path[2])])
             for p in path[3:end]
-                res = :(Base.getfield($res, $(Meta.quot(p))))
+                res = :(traced_getfield($res, $(Meta.quot(p))))
             end
             sym = Symbol("sbuf_$i")
             sbuf = :($sym = XLA.synced_buffer($res.data))
@@ -375,7 +377,7 @@ function compile(f, args; pipeline_options="", client=nothing)
                     path = path[3:end]
                 end
                 for p in path
-                    res = :(Base.getfield($res, $(Meta.quot(p))))
+                    res = :(traced_getfield($res, $(Meta.quot(p))))
                 end
                 res = :($res.data = $(Symbol("concrete_res_$(idx)")))
                 push!(delinearized_results, res)
@@ -401,12 +403,12 @@ function compile(f, args; pipeline_options="", client=nothing)
                     path = path[3:end]
                 end
                 for p in path
-                    res = :(Base.getfield($res, $(Meta.quot(p))))
+                    res = :(traced_getfield($res, $(Meta.quot(p))))
                 end
 
                 argres = :(args[argpath[2]])
                 for p in argpath[3:end]
-                    argres = :(Base.getfield($argres, $(Meta.quot(p))))
+                    argres = :(traced_getfield($argres, $(Meta.quot(p))))
                 end
 
                 res = :($res.data = $argres.data)