within_tracing, approach taken from Enzyme.jl

EnzymeAD · Dec 17, 2024 · 9b5cd96 · 9b5cd96
1 parent 55d1527
commit 9b5cd96
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 6 deletions.
diff --git a/src/Compiler.jl b/src/Compiler.jl
@@ -15,7 +15,6 @@ import ..Reactant:
     TracedType,
     Cached
 using ScopedValues
-import ReactantCore: enable_tracing
 
 @inline traced_getfield(@nospecialize(obj), field) = Base.getfield(obj, field)
 @inline traced_getfield(
@@ -297,9 +296,7 @@ function compile_mlir!(mod, f, args, callcache; optimize::Union{Bool,Symbol}=tru
     linear_results = MLIR.IR.mmodule!(mod) do
         MLIR.IR.block!(MLIR.IR.body(mod)) do
             callcache!(callcache) do
-                with(enable_tracing=>true) do
-                    return Reactant.TracedUtils.make_mlir_fn(f, args, (), "main", true)
-                end
+                return Reactant.TracedUtils.make_mlir_fn(f, args, (), "main", true)
             end
         end
     end

diff --git a/src/Interpreter.jl b/src/Interpreter.jl
@@ -39,6 +39,30 @@ function set_reactant_abi(
 )
     (; fargs, argtypes) = arginfo
 
+    if f === ReactantCore.within_tracing
+        if length(argtypes) != 1
+            @static if VERSION < v"1.11.0-"
+                return CallMeta(Union{}, Effects(), NoCallInfo())
+            else
+                return CallMeta(Union{}, Union{}, Effects(), NoCallInfo())
+            end
+        end
+        @static if VERSION < v"1.11.0-"
+            return CallMeta(
+                Core.Const(true),
+                Core.Compiler.EFFECTS_TOTAL,
+                MethodResultPure(),
+            )
+        else
+            return CallMeta(
+                Core.Const(true),
+                Union{},
+                Core.Compiler.EFFECTS_TOTAL,
+                MethodResultPure(),
+            )
+        end
+    end
+
     # Improve inference by considering call_with_reactant as having the same results as
     # the original call
     if f === Reactant.call_with_reactant

diff --git a/src/Reactant.jl b/src/Reactant.jl
@@ -1,6 +1,6 @@
 module Reactant
 
-using ReactantCore: ReactantCore, @trace, MissingTracedValue
+using ReactantCore: ReactantCore, @trace, within_tracing, MissingTracedValue
 
 using LinearAlgebra: LinearAlgebra
 using Adapt: Adapt, WrappedArray
@@ -145,7 +145,7 @@ function Enzyme.make_zero(
 end
 
 using .Compiler: @compile, @code_hlo, @jit, traced_getfield, create_result, compile
-export ConcreteRArray, ConcreteRNumber, @compile, @code_hlo, @jit, @trace
+export ConcreteRArray, ConcreteRNumber, @compile, @code_hlo, @jit, @trace, within_tracing
 
 const registry = Ref{MLIR.IR.DialectRegistry}()
 function __init__()