fix: despecialize sharding fields of ConcreteTypes (#1801)

* fix: despecialize sharding fields of ConcreteTypes

* fix: sharding of structs now work

* fix: dispatches

* fix: missing

* test: incorrect type

* fix: preserve sharding

* test: add #1227 as a testcase

Author: avik-pal

docs/src/tutorials/partial-evaluation.md

@@ -36,7 +36,7 @@ addxy(x, y)
 
 # output
 
-ConcretePJRTNumber{Int64, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(7)
+ConcretePJRTNumber{Int64, 1}(7)
 ```
 
 returns a result that depends on both arguments `x` and `y`:
@@ -46,7 +46,7 @@ addxy(ConcreteRNumber(7), ConcreteRNumber(8))
 
 # output
 
-ConcretePJRTNumber{Int64, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(15)
+ConcretePJRTNumber{Int64, 1}(15)
 ```
 
 The StableHLO IR code generated here is:
@@ -76,7 +76,7 @@ addx4(x, 4)
 
 # output
 
-ConcretePJRTNumber{Int64, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(7)
+ConcretePJRTNumber{Int64, 1}(7)
 ```
 
 will only change based on `x`, not on the non-Reactant argument `y`, we get
@@ -87,7 +87,7 @@ addx4(ConcreteRNumber(7), 8)
 
 # output
 
-ConcretePJRTNumber{Int64, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(11)
+ConcretePJRTNumber{Int64, 1}(11)
 ```
 
 The StableHLO code shows that the second argument has been replaced by a

ext/ReactantCUDAExt.jl

@@ -1379,16 +1379,9 @@ Base.@nospecializeinfer function Reactant.traced_type_inner(
     N = ndims(A)
     if mode == Reactant.ArrayToConcrete && T <: Reactant.ReactantPrimitive
         if runtime isa Val{:PJRT}
-            return Reactant.ConcretePJRTArray{
-                T,
-                N,
-                Reactant.Sharding.ndevices(sharding),
-                Reactant.Sharding.shard_type(typeof(sharding), N),
-            }
+            return Reactant.ConcretePJRTArray{T,N,Reactant.Sharding.ndevices(sharding)}
         elseif runtime isa Val{:IFRT}
-            return Reactant.ConcreteIFRTArray{
-                T,N,Reactant.Sharding.shard_type(typeof(sharding), N)
-            }
+            return Reactant.ConcreteIFRTArray{T,N}
         end
         error("Unsupported runtime $runtime")
     else

src/Compiler.jl

@@ -3,6 +3,7 @@ module Compiler
 using Reactant_jll
 using Libdl: dlsym
 using LinearAlgebra: BlasInt
+using Functors: Functors
 
 import ..Reactant:
     Reactant,
@@ -272,7 +273,7 @@ function create_result(
 end
 
 function create_result(
-    tocopy::ConcretePJRTNumber{T,D,S},
+    tocopy::ConcretePJRTNumber{T,D},
     path,
     result_stores,
     path_to_shard_info,
@@ -283,7 +284,7 @@ function create_result(
     result_cache,
     var_idx,
     resultgen_code,
-) where {T,D,S}
+) where {T,D}
     if !haskey(result_cache, tocopy)
         sym = Symbol("result", var_idx[])
         var_idx[] += 1
@@ -314,7 +315,7 @@ function create_result(
 end
 
 function create_result(
-    tocopy::ConcreteIFRTNumber{T,S},
+    tocopy::ConcreteIFRTNumber{T},
     path,
     result_stores,
     path_to_shard_info,
@@ -325,7 +326,7 @@ function create_result(
     result_cache,
     var_idx,
     resultgen_code,
-) where {T,S}
+) where {T}
     if !haskey(result_cache, tocopy)
         sym = Symbol("result", var_idx[])
         var_idx[] += 1
@@ -356,7 +357,7 @@ function create_result(
 end
 
 function create_result(
-    tocopy::ConcretePJRTArray{T,N,D,S},
+    tocopy::ConcretePJRTArray{T,N,D},
     path,
     result_stores,
     path_to_shard_info,
@@ -367,7 +368,7 @@ function create_result(
     result_cache,
     var_idx,
     resultgen_code,
-) where {T,N,D,S}
+) where {T,N,D}
     if !haskey(result_cache, tocopy)
         sym = Symbol("result", var_idx[])
         var_idx[] += 1
@@ -399,7 +400,7 @@ function create_result(
 end
 
 function create_result(
-    tocopy::ConcreteIFRTArray{T,N,S},
+    tocopy::ConcreteIFRTArray{T,N},
     path,
     result_stores,
     path_to_shard_info,
@@ -410,7 +411,7 @@ function create_result(
     result_cache,
     var_idx,
     resultgen_code,
-) where {T,N,S}
+) where {T,N}
     if !haskey(result_cache, tocopy)
         sym = Symbol("result", var_idx[])
         var_idx[] += 1
@@ -1647,12 +1648,6 @@ function compile_mlir!(
         raise isa Bool && (raise = true)
     end
 
-    concrete_seen = OrderedIdDict()
-
-    concrete_result = make_tracer(
-        concrete_seen, traced_result, ("result",), TracedToConcrete; runtime
-    )
-
     toolkit = XLA.CUDA_DATA_DIR[]
 
     if backend == "cpu" || backend == "tpu"
@@ -2307,6 +2302,17 @@ function compile_mlir!(
         ]
     end
 
+    if result_shardings !== missing
+        result_shardings_after_masking = eltype(result_shardings)[]
+        for (i, present) in enumerate(results_mask)
+            if present
+                push!(result_shardings_after_masking, result_shardings[i])
+            end
+        end
+    else
+        result_shardings_after_masking = missing
+    end
+
     func3 = MLIR.Dialects.func.func_(;
         sym_name="main",
         function_type=MLIR.IR.FunctionType(in_tys, out_tys2),
@@ -2384,6 +2390,10 @@ function compile_mlir!(
         end
     end
 
+    concrete_result = make_tracer(
+        OrderedIdDict(), traced_result, ("result",), TracedToConcrete; runtime
+    )
+
     return Reactant.TracedUtils.CompiledMlirFnResult(
         fnwrapped,
         func3,
@@ -2404,7 +2414,7 @@ function compile_mlir!(
         mlir_fn_res.unique_meshes,
         mlir_fn_res.mutated_args,
         use_shardy_partitioner,
-        result_shardings,
+        result_shardings_after_masking,
         mlir_fn_res.global_device_ids,
         donated_args_mask,
         Reactant.TracedUtils.is_pure(func3),
@@ -3151,7 +3161,6 @@ function codegen_unflatten!(
                 p in Reactant.TracedUtils.get_paths(arg) if length(p) > 0 && p[1] == :args
             ))
 
-            res = :result
             path = path[2:end]
 
             if in(path, keys(result_stores))
@@ -3246,10 +3255,15 @@ function codegen_unflatten!(
     end
 
     # generate return object which stores the concrete results in some arbitrary way
-    return Expr[
-        unresharded_code..., resultgen_code..., :(result = $result_code), unflatten_code...
-    ],
-    used_shardinfo
+    return (
+        Expr[
+            unresharded_code...,
+            resultgen_code...,
+            :(result = $result_code),
+            unflatten_code...,
+        ],
+        used_shardinfo,
+    )
 end
 
 """

src/ConcreteRArray.jl

@@ -23,8 +23,8 @@ function Base.copy(x::Union{AbstractConcreteArray,AbstractConcreteNumber})
     return fn(x)
 end
 
-function Base.copy(X::ConcreteIFRTArray{T,D,S,P}) where {T,D,S,P}
-    return ConcreteIFRTArray{T,D,S}(Base.copy(X.data), X.shape, X.sharding, X.padding)
+function Base.copy(X::ConcreteIFRTArray{T,D,P}) where {T,D,P}
+    return ConcreteIFRTArray{T,D}(Base.copy(X.data), X.shape, X.sharding, X.padding)
 end
 
 function Base.copy(X::ConcretePJRTArray)
@@ -422,8 +422,8 @@ function Base.similar(
 end
 
 function Base.similar(
-    a::ConcretePJRTArray{T,N,D,Sh}, ::Type{S}=T, dims::Dims=size(a)
-) where {S,T,Sh,N,D}
+    a::ConcretePJRTArray{T,N,D}, ::Type{S}=T, dims::Dims=size(a)
+) where {S,T,N,D}
     device_to_array_slices, sharding = Sharding.sharding_to_array_slices(
         a.sharding, dims; return_updated_sharding=Val(true), client=XLA.client(a)
     )
@@ -432,7 +432,7 @@ function Base.similar(
         Base.@_inline_meta
         similar(a.data[i], S, Dims(length.(device_to_array_slices[i])))
     end
-    return ConcretePJRTArray{S,length(dims),D,Sh}(sdata, dims, a.sharding)
+    return ConcretePJRTArray{S,length(dims),D}(sdata, dims, a.sharding)
 end
 
 Base.similar(a::ConcretePJRTArray, dims::Dims) = similar(a, eltype(a), dims)

src/TracedUtils.jl

@@ -17,7 +17,6 @@ using ..Reactant:
 using ..Ops: @opcall
 using ReactantCore: ReactantCore
 using ReactantCore: MissingTracedValue, is_traced, materialize_traced_array
-using Functors: Functors
 
 ReactantCore.materialize_traced_array(x::AbstractArray) = x