Handle promoted traced numbers (#1247)

* Handle promoted traced numbers

* fix mlir ownership

* fix

* Apply suggestions from code review

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

---------

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

Author: wsmoses

ext/ReactantCUDAExt.jl

@@ -991,6 +991,7 @@ Reactant.@reactant_overlay @noinline function (func::LLVMFunc{F,tt})(
 ) where {F,tt}
     blockdim = CUDA.CuDim3(blocks)
     threaddim = CUDA.CuDim3(threads)
+    mod = MLIR.IR.mmodule()
 
     if convert == Val(true)
         args = recudaconvert.(args)
@@ -1021,7 +1022,6 @@ Reactant.@reactant_overlay @noinline function (func::LLVMFunc{F,tt})(
     end
 
     sym_name = String(gensym("call_$fname"))
-    mod = MLIR.IR.mmodule()
     CConv = MLIR.IR.Attribute(
         MLIR.API.mlirLLVMCConvAttrGet(ctx, MLIR.API.MlirLLVMCConvPTX_Kernel)
     )

src/Ops.jl

@@ -104,7 +104,9 @@ end
     end
 
     value = MLIR.IR.DenseElementsAttribute(x)
-    constants = constant_context()[2]
+    constant_blk, constants = constant_context()
+    parent = MLIR.IR.parent_op(constant_blk)
+    @assert MLIR.IR.name(parent) != "builtin.module"
     if haskey(constants, value)
         return constants[value]
     else
@@ -137,6 +139,7 @@ end
 @noinline function constant(
     x::AbstractArray{T,N}; location=mlir_stacktrace("constant", @__FILE__, @__LINE__)
 ) where {T,N}
+    @assert !(x isa TracedRArray)
     return constant(collect(x); location)
 end
 

src/TracedUtils.jl

@@ -280,12 +280,9 @@ function make_mlir_fn(
         return mlir_fn_res
     end
 
-    seen_args = OrderedIdDict()
-
-    (; N, traced_args, linear_args, inv_map, in_tys, sym_visibility, mod, traced_args_to_shardings, func, fnbody) = prepare_mlir_fn_args(
+    (; N, traced_args, linear_args, inv_map, in_tys, sym_visibility, mod, traced_args_to_shardings, func, fnbody, seen_args, skipped_args) = prepare_mlir_fn_args(
         args,
         name,
-        seen_args,
         concretein,
         toscalar,
         argprefix,
@@ -327,14 +324,12 @@ function make_mlir_fn(
         end
     end
 
-    seen_results = OrderedIdDict()
-
     (func2, traced_result, ret, linear_args, in_tys, linear_results, num_partitions, is_sharded, unique_meshes, mutated_args, global_device_ids) = finalize_mlir_fn(
         result,
         traced_args,
         linear_args,
+        skipped_args,
         seen_args,
-        seen_results,
         fnbody,
         func,
         mod,
@@ -388,7 +383,6 @@ end
 function prepare_mlir_fn_args(
     args,
     name,
-    seen_args,
     concretein,
     toscalar,
     argprefix,
@@ -406,16 +400,35 @@ function prepare_mlir_fn_args(
     else
         Reactant.TracedSetPath
     end
-    for i in 1:N
-        @inbounds traced_args[i] = Reactant.make_tracer(
-            seen_args, args[i], (argprefix, i), inmode; toscalar, runtime
-        )
+    fnbody = MLIR.IR.Block(MLIR.IR.Type[], MLIR.IR.Location[])
+    MLIR.IR.activate!(fnbody)
+    Ops.activate_constant_context!(fnbody)
+    seen_args0 = OrderedIdDict()
+    try
+        for i in 1:N
+            @inbounds traced_args[i] = Reactant.make_tracer(
+                seen_args0, args[i], (argprefix, i), inmode; toscalar, runtime
+            )
+        end
+    finally
+        MLIR.IR.deactivate!(fnbody)
+        Ops.deactivate_constant_context!(fnbody)
     end
 
+    seen_args = OrderedIdDict()
     linear_args = Reactant.TracedType[]
+    skipped_args = Reactant.TracedType[]
     inv_map = IdDict()
-    for (k, v) in seen_args
+    for (k, v) in seen_args0
         v isa Reactant.TracedType || continue
+        arg = get_mlir_data(v)
+        if (arg isa MLIR.IR.Value) &&
+            MLIR.IR.is_op_res(arg) &&
+            MLIR.IR.block(MLIR.IR.op_owner(arg)) == fnbody
+            push!(skipped_args, v)
+            continue
+        end
+        seen_args[k] = v
         push!(linear_args, v)
         inv_map[v] = k
     end
@@ -468,7 +481,7 @@ function prepare_mlir_fn_args(
     end
 
     arglocs = MLIR.IR.Location[]
-    for arg in linear_args
+    for (i, arg) in enumerate(linear_args)
         path = get_idx(arg, argprefix)
         stridx = if verify_arg_names isa Nothing
             "arg" * string(path[2])
@@ -490,9 +503,12 @@ function prepare_mlir_fn_args(
                 aval = getfield(aval, idx)
             end
         end
-        push!(arglocs, MLIR.IR.Location(stridx * " (path=$path)", MLIR.IR.Location()))
+        MLIR.IR.push_argument!(
+            fnbody,
+            in_tys[i];
+            location=MLIR.IR.Location(stridx * " (path=$path)", MLIR.IR.Location()),
+        )
     end
-    fnbody = MLIR.IR.Block(in_tys, arglocs)
     push!(MLIR.IR.region(func, 1), fnbody)
 
     return (;
@@ -506,6 +522,8 @@ function prepare_mlir_fn_args(
         traced_args_to_shardings,
         func,
         fnbody,
+        seen_args,
+        skipped_args,
     )
 end
 
@@ -533,8 +551,8 @@ function finalize_mlir_fn(
     result,
     traced_args,
     linear_args,
+    skipped_args,
     seen_args,
-    seen_results,
     fnbody,
     func,
     mod,
@@ -578,6 +596,7 @@ function finalize_mlir_fn(
         Reactant.TracedTrack
     end
 
+    seen_results = OrderedIdDict()
     MLIR.IR.activate!(fnbody)
     traced_result = try
         traced_result = Reactant.make_tracer(
@@ -602,6 +621,9 @@ function finalize_mlir_fn(
     linear_results = Reactant.TracedType[]
     for (k, v) in seen_results
         v isa Reactant.TracedType || continue
+        if any(Base.Fix1(===, k), skipped_args)
+            continue
+        end
         if args_in_result != :all
             if has_idx(v, argprefix)
                 if !(

src/Tracing.jl

@@ -1128,12 +1128,19 @@ Base.@nospecializeinfer function make_tracer_unknown(
             if xi !== xi2
                 changed = true
             end
-            if mode != TracedToTypes && !(Core.Typeof(xi2) <: fieldtype(TT, i))
-                throw(
-                    AssertionError(
-                        "Could not recursively make tracer of object of type $RT into $TT at field $i (named $(fieldname(TT, i))), need object of type $(fieldtype(TT, i)) found object of type $(Core.Typeof(xi2)) ",
-                    ),
-                )
+            FT = fieldtype(TT, i)
+            if mode != TracedToTypes && !(Core.Typeof(xi2) <: FT)
+                if FT <: TracedRNumber && xi2 isa unwrapped_eltype(FT)
+                    xi2 = FT(xi2)
+                    xi2 = Core.Typeof(xi2)((newpath,), xi2.mlir_data)
+                    seen[xi2] = xi2
+                else
+                    throw(
+                        AssertionError(
+                            "Could not recursively make tracer of object of type $RT into $TT at field $i (named $(fieldname(TT, i))), need object of type $(fieldtype(TT, i)) found object of type $(Core.Typeof(xi2)) ",
+                        ),
+                    )
+                end
             end
             flds[i] = xi2
         else

src/mlir/IR/Block.jl

@@ -38,14 +38,14 @@ Base.unsafe_convert(::Core.Type{API.MlirBlock}, block::Block) = block.block
 
 Returns the closest surrounding operation that contains this block.
 """
-parent_op(block::Block) = Operation(API.mlirBlockGetParentOperation(block))
+parent_op(block::Block) = Operation(API.mlirBlockGetParentOperation(block), false)
 
 """
     parent_region(block)
 
 Returns the region that contains this block.
 """
-parent_region(block::Block) = Region(API.mlirBlockGetParentRegion(block))
+parent_region(block::Block) = Region(API.mlirBlockGetParentRegion(block), false)
 
 Base.parent(block::Block) = parent_region(block)