Introduce transmute, transmute_type functions as future names for make_tracer, traced_type

docs/src/api/api.md

@@ -43,9 +43,11 @@ OptimizeCommunicationOptions
 ShardyPropagationOptions
 ```
 
-## Tracing customization
+## Tracing
 
 ```@docs
+Reactant.transmute
+Reactant.transmute_type
 Reactant.@skip_rewrite_func
 Reactant.@skip_rewrite_type
 ```

ext/ReactantCUDAExt.jl

@@ -1018,7 +1018,7 @@ function Reactant.make_tracer(
 )
     x = Base.unsafe_pointer_to_objref(Base.reinterpret(Ptr{Cvoid}, prev.ptr))
     x = x::TracedRArray
-    Reactant.make_tracer(seen, x, path, mode; kwargs...)
+    Reactant.transmute(seen, x, path, mode; kwargs...)
     return prev
 end
 
@@ -1027,7 +1027,7 @@ function Reactant.make_tracer(
 )
     x = Base.unsafe_pointer_to_objref(Base.reinterpret(Ptr{Cvoid}, prev.ptr))
     x = x::TracedRNumber
-    Reactant.make_tracer(seen, x, path, mode; kwargs...)
+    Reactant.transmute(seen, x, path, mode; kwargs...)
     return prev
 end
 
@@ -1093,7 +1093,7 @@ Reactant.@reactant_overlay @noinline function (func::LLVMFunc{F,tt})(
     kernelargsym = gensym("kernelarg")
 
     for (i, prev) in enumerate(Any[func.f, args...])
-        Reactant.make_tracer(seen, prev, (kernelargsym, i), Reactant.NoStopTracedTrack)
+        Reactant.transmute(seen, prev, (kernelargsym, i), Reactant.NoStopTracedTrack)
     end
     wrapper_tys = MLIR.IR.Type[]
     for arg in values(seen)
@@ -1420,7 +1420,7 @@ function Reactant.make_tracer(
         end
         error("Unsupported runtime $runtime")
     end
-    TT = Reactant.traced_type(eltype(RT), Val(mode), track_numbers, sharding, runtime)
+    TT = Reactant.transmute_type(eltype(RT), Val(mode), track_numbers, sharding, runtime)
     if TT === eltype(RT)
         return prev
     end
@@ -1430,7 +1430,7 @@ function Reactant.make_tracer(
     for I in eachindex(prev)
         if isassigned(prev, I)
             pv = prev[I]
-            nv = Reactant.make_tracer(
+            nv = Reactant.transmute(
                 seen,
                 pv,
                 append_path(path, I),

src/Compiler.jl

@@ -1547,7 +1547,7 @@ function compile_mlir!(
 
     concrete_seen = OrderedIdDict()
 
-    concrete_result = make_tracer(
+    concrete_result = Reactant.transmute(
         concrete_seen, traced_result, ("result",), TracedToConcrete; runtime
     )
 

src/Ops.jl

@@ -1887,7 +1887,7 @@ end
     traced_args = Vector{Any}(undef, N)
 
     for (i, prev) in enumerate(args)
-        @inbounds traced_args[i] = Reactant.make_tracer(
+        @inbounds traced_args[i] = Reactant.transmute(
             seen_args, prev, (), Reactant.NoStopTracedTrack; track_numbers
         )
     end
@@ -1972,14 +1972,14 @@ end
     tb_traced_args = Vector{Any}(undef, N)
     fb_traced_args = Vector{Any}(undef, N)
     for i in 1:N
-        @inbounds tb_traced_args[i] = Reactant.make_tracer(
+        @inbounds tb_traced_args[i] = Reactant.transmute(
             tb_seen_args,
             args[i],
             (true_fn_names[1], i),
             Reactant.TracedSetPath;
             track_numbers,
         )
-        @inbounds fb_traced_args[i] = Reactant.make_tracer(
+        @inbounds fb_traced_args[i] = Reactant.transmute(
             fb_seen_args,
             args[i],
             (false_fn_names[1], i),
@@ -2043,15 +2043,15 @@ end
     end
 
     seen_true_results = Reactant.OrderedIdDict()
-    traced_true_results = Reactant.make_tracer(
+    traced_true_results = Reactant.transmute(
         seen_true_results,
         tb_result,
         (true_fn_names[2],),
         Reactant.NoStopTracedTrack;
         track_numbers,
     )
     for (i, arg) in enumerate(tb_traced_args)
-        Reactant.make_tracer(
+        Reactant.transmute(
             seen_true_results,
             arg,
             (true_fn_names[3], i),
@@ -2108,15 +2108,15 @@ end
     end
 
     seen_false_results = Reactant.OrderedIdDict()
-    traced_false_results = Reactant.make_tracer(
+    traced_false_results = Reactant.transmute(
         seen_false_results,
         fb_result,
         (false_fn_names[2],),
         Reactant.NoStopTracedTrack;
         track_numbers,
     )
     for (i, arg) in enumerate(fb_traced_args)
-        Reactant.make_tracer(
+        Reactant.transmute(
             seen_false_results,
             arg,
             (false_fn_names[3], i),
@@ -2372,7 +2372,7 @@ end
 @noinline function call(f, args...)
     seen = Reactant.OrderedIdDict()
     cache_key = []
-    Reactant.make_tracer(seen, (f, args...), cache_key, Reactant.TracedToTypes)
+    Reactant.transmute(seen, (f, args...), cache_key, Reactant.TracedToTypes)
     cache = Reactant.Compiler.callcache()
     if haskey(cache, cache_key)
         # cache lookup:
@@ -2414,7 +2414,7 @@ end
     end
 
     seen_cache = Reactant.OrderedIdDict()
-    Reactant.make_tracer(
+    Reactant.transmute(
         seen_cache,
         args,
         (), # we have to insert something here, but we remove it immediately below.
@@ -2438,7 +2438,7 @@ end
     )
 
     seen_results = Reactant.OrderedIdDict()
-    traced_result = Reactant.make_tracer(
+    traced_result = Reactant.transmute(
         seen_results,
         traced_result,
         (), # we have to insert something here, but we remove it immediately below.

src/Reactant.jl

@@ -98,6 +98,8 @@ unwrapped_eltype(::TracedRNumber{T}) where {T} = T
 unwrapped_eltype(::Type{<:AbstractArray{T,N}}) where {T,N} = unwrapped_eltype(T)
 unwrapped_eltype(::AbstractArray{T,N}) where {T,N} = unwrapped_eltype(T)
 
+# TODO replace `promote_trace_type` with `promote_transmute_type` on v0.3 release
+promote_transmute_type(a::Type, b::Type) = promote_traced_type(a, b)
 promote_traced_type(a::Type, b::Type) = Base.promote_type(a, b)
 
 aos_to_soa(x::AbstractArray) = x

src/TracedUtils.jl

@@ -431,7 +431,7 @@ function prepare_mlir_fn_args(
     seen_args0 = OrderedIdDict()
     try
         for i in 1:N
-            @inbounds traced_args[i] = Reactant.make_tracer(
+            @inbounds traced_args[i] = Reactant.transmute(
                 seen_args0, args[i], (argprefix, i), inmode; toscalar, runtime
             )
         end
@@ -623,13 +623,11 @@ function finalize_mlir_fn(
     seen_results = OrderedIdDict()
     MLIR.IR.activate!(fnbody)
     traced_result = try
-        traced_result = Reactant.make_tracer(
-            seen_results, result, (resprefix,), outmode; runtime
-        )
+        traced_result = Reactant.transmute(seen_results, result, (resprefix,), outmode; runtime)
 
         # marks buffers to be donated
         for i in 1:N
-            Reactant.make_tracer(
+            Reactant.transmute(
                 seen_results,
                 traced_args[i],
                 (resargprefix, i),
@@ -1162,7 +1160,7 @@ function elem_apply(f, args::Vararg{Any,Nargs}) where {Nargs}
     end
 
     seen_results = OrderedIdDict()
-    traced2_result = Reactant.make_tracer(
+    traced2_result = Reactant.transmute(
         seen_results, result, (), Reactant.TracedSetPath; tobatch=OutShape
     )