From aa93f2c3dfae346f124c77031e9f48ef2b45c346 Mon Sep 17 00:00:00 2001 From: Billy Moses Date: Wed, 3 Jun 2026 16:05:53 -0500 Subject: [PATCH 1/9] Emulated complex --- src/ConcreteRArray.jl | 49 ++++++++--- src/TracedUtils.jl | 150 ++++++++++++++++++++++++++------ src/Types.jl | 181 +++++++++++++++++++++++++++++---------- src/compiler/Compiler.jl | 6 ++ src/utils.jl | 2 + src/xla/Client.jl | 13 +++ src/xla/IFRT/Array.jl | 6 +- src/xla/PJRT/Buffer.jl | 2 +- test/core/complex.jl | 28 ++++++ 9 files changed, 354 insertions(+), 83 deletions(-) diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl index 414ea4222c..d2402559c2 100644 --- a/src/ConcreteRArray.jl +++ b/src/ConcreteRArray.jl @@ -106,19 +106,36 @@ function Base.isempty(x::Union{WrappedConcretePJRTArray,WrappedConcreteIFRTArray end function Base.convert(::Type{<:Array}, X::AbstractConcreteArray{T,N}) where {T,N} - if has_padding(X) - padding = get_padding(X) - data = Array{T,N}(undef, (size(X) .+ padding)...) - write_to_host_buffer!(data, X) - return view(data, [1:size(X, i) for i in 1:ndims(X)]...) + client = XLA.client(X) + if T <: Complex && !XLA.supports_complex(client) + T2 = real(T) + if has_padding(X) + padding = get_padding(X) + data = Array{T,N}(undef, (size(X) .+ padding)...) + data_re = reinterpret(reshape, T2, data) + write_to_host_buffer!(data_re, X) + return view(data, [1:size(X, i) for i in 1:ndims(X)]...) + else + data = Array{T,N}(undef, size(X)...) + data_re = reinterpret(reshape, T2, data) + write_to_host_buffer!(data_re, X) + return data + end else - data = Array{T,N}(undef, size(X)...) - write_to_host_buffer!(data, X) - return data + if has_padding(X) + padding = get_padding(X) + data = Array{T,N}(undef, (size(X) .+ padding)...) + write_to_host_buffer!(data, X) + return view(data, [1:size(X, i) for i in 1:ndims(X)]...) + else + data = Array{T,N}(undef, size(X)...) + write_to_host_buffer!(data, X) + return data + end end end -function write_to_host_buffer!(data::Array, X::ConcretePJRTArray{T,N}) where {T,N} +function write_to_host_buffer!(data::AbstractArray, X::ConcretePJRTArray{T,N}) where {T,N} if Sharding.is_sharded(X) completed = Set{eltype(X.sharding.device_to_array_slices)}() for idx in 1:length(X.data) @@ -135,11 +152,12 @@ function write_to_host_buffer!(data::Array, X::ConcretePJRTArray{T,N}) where {T, return nothing end -function write_to_host_buffer!(data::Array, X::ConcreteIFRTArray{T,N}) where {T,N} +function write_to_host_buffer!(data::AbstractArray, X::ConcreteIFRTArray{T,N}) where {T,N} XLA.to_host(X.data, data, X.sharding) return nothing end + function Base.convert( ::Type{<:Array}, X::Union{WrappedConcretePJRTArray,WrappedConcreteIFRTArray} ) @@ -707,7 +725,16 @@ for aType in (:ConcretePJRTArray, :ConcreteIFRTArray) function Base.copyto!( dest::Array, src::Broadcast.Broadcasted{Broadcast.ArrayStyle{$(aType)}} ) - write_to_host_buffer!(dest, copy(src)) + copied_src = copy(src) + client = XLA.client(copied_src) + T = eltype(copied_src) + if T <: Complex && !XLA.supports_complex(client) + T2 = real(T) + dest_re = reinterpret(reshape, T2, dest) + write_to_host_buffer!(dest_re, copied_src) + else + write_to_host_buffer!(dest, copied_src) + end return dest end diff --git a/src/TracedUtils.jl b/src/TracedUtils.jl index dbf2b6703e..f43524de6f 100644 --- a/src/TracedUtils.jl +++ b/src/TracedUtils.jl @@ -12,6 +12,7 @@ using ..Reactant: MissingTracedValue, OrderedIdDict, Ops, + XLA, promote_to, # keep this to avoid breaking external code broadcast_to_size # keep this to avoid breaking external code using ..Ops: @opcall @@ -317,6 +318,8 @@ function make_mlir_fn( resargprefix::Symbol=:resargs, num_replicas=1, optimize_then_pad::Bool=true, + client=nothing, + emulate_complex::Bool=false, ) if sizeof(typeof(f)) != 0 || f isa Base.BroadcastFunction mlir_fn_res = make_mlir_fn( @@ -338,6 +341,8 @@ function make_mlir_fn( resprefix, resargprefix, num_replicas, + client, + emulate_complex, ) mlir_fn_res.fnwrapped = true return mlir_fn_res @@ -353,7 +358,9 @@ function make_mlir_fn( optimize_then_pad, do_transpose, input_shardings, - verify_arg_names, + verify_arg_names; + client, + emulate_complex, ) Ops.activate_constant_context!(fnbody) @@ -364,7 +371,7 @@ function make_mlir_fn( MLIR.IR.activate(fnbody) result = try - process_linear_args!(linear_args, fnbody, do_transpose, optimize_then_pad, inv_map) + process_linear_args!(linear_args, fnbody, do_transpose, optimize_then_pad, inv_map; client, emulate_complex) if isempty(kwargs) Reactant.call_with_reactant(f, traced_args...) @@ -416,7 +423,9 @@ function make_mlir_fn( args, N, concretein, - toscalar, + toscalar; + client, + emulate_complex, ) return CompiledMlirFnResult( @@ -456,7 +465,9 @@ function prepare_mlir_fn_args( optimize_then_pad, do_transpose, input_shardings, - verify_arg_names, + verify_arg_names; + client=nothing, + emulate_complex=false, ) N = length(args) traced_args = Vector{Any}(undef, N) @@ -499,21 +510,41 @@ function prepare_mlir_fn_args( inv_map[v] = k end + has_complex = !emulate_complex || XLA.supports_complex(client) + in_tys = Vector{MLIR.IR.Type}(undef, length(linear_args)) for (i, arg) in enumerate(linear_args) - elT = MLIR.IR.Type(Reactant.unwrapped_eltype(arg)) - if toscalar - in_tys[i] = MLIR.IR.TensorType(Int[], elT) - else - sz = collect(Int, size(arg)) - if !optimize_then_pad - carg = inv_map[arg] - Reactant.has_padding(carg) && (sz .+= Reactant.get_padding(carg)) + T = Reactant.unwrapped_eltype(arg) + if T <: Complex && !has_complex + elT = MLIR.IR.Type(real(T)) + if toscalar + in_tys[i] = MLIR.IR.TensorType(Int[2], elT) + else + sz = collect(Int, size(arg)) + if !optimize_then_pad + carg = inv_map[arg] + Reactant.has_padding(carg) && (sz .+= Reactant.get_padding(carg)) + end + sz = [2; sz] + typ = MLIR.IR.TensorType(sz, elT) + do_transpose && (typ = transpose_ty(typ)) + in_tys[i] = typ end + else + elT = MLIR.IR.Type(T) + if toscalar + in_tys[i] = MLIR.IR.TensorType(Int[], elT) + else + sz = collect(Int, size(arg)) + if !optimize_then_pad + carg = inv_map[arg] + Reactant.has_padding(carg) && (sz .+= Reactant.get_padding(carg)) + end - typ = MLIR.IR.TensorType(sz, elT) - do_transpose && (typ = transpose_ty(typ)) - in_tys[i] = typ + typ = MLIR.IR.TensorType(sz, elT) + do_transpose && (typ = transpose_ty(typ)) + in_tys[i] = typ + end end end @@ -592,7 +623,8 @@ function prepare_mlir_fn_args( ) end -function process_linear_args!(linear_args, fnbody, do_transpose, optimize_then_pad, inv_map) +function process_linear_args!(linear_args, fnbody, do_transpose, optimize_then_pad, inv_map; client=nothing, emulate_complex=false) + has_complex = !emulate_complex || XLA.supports_complex(client) for (i, arg) in enumerate(linear_args) raw_arg = MLIR.IR.argument(fnbody, i) row_maj_arg = do_transpose ? transpose_val(raw_arg) : raw_arg @@ -608,7 +640,31 @@ function process_linear_args!(linear_args, fnbody, do_transpose, optimize_then_p row_maj_arg = MLIR.IR.result(unpad_val_op(row_maj_arg, padding, sz), 1) end end - set_mlir_data!(arg, row_maj_arg) + + T = Reactant.unwrapped_eltype(arg) + if T <: Complex && !has_complex + T2 = real(T) + sz = size(arg) + N = length(sz) + + start_real = [1; ones(Int, N)] + limit_real = [1; collect(Int, sz)] + start_imag = [2; ones(Int, N)] + limit_imag = [2; collect(Int, sz)] + + traced_raw = TracedRArray{T2, N+1}((), row_maj_arg, (2, sz...)) + + real_slice = Ops.slice(traced_raw, start_real, limit_real) + imag_slice = Ops.slice(traced_raw, start_imag, limit_imag) + + real_part = Ops.reshape(real_slice, sz...) + imag_part = Ops.reshape(imag_slice, sz...) + + complex_val = Ops.complex(real_part, imag_part) + set_mlir_data!(arg, complex_val.mlir_data) + else + set_mlir_data!(arg, row_maj_arg) + end end end @@ -641,8 +697,12 @@ function finalize_mlir_fn( args, N, concretein, - toscalar, + toscalar; + client=nothing, + emulate_complex=false, ) + has_complex = !emulate_complex || XLA.supports_complex(client) + # check which arguments have been mutated mutated_args = Int[] if !construct_function_without_args @@ -856,12 +916,39 @@ function finalize_mlir_fn( col_maj = get_mlir_data(broadcast_to_size(false, ())) out_ty = Ops.mlir_type(TracedRArray{Bool,0}, ()) else - col_maj = get_mlir_data(res) - out_ty = Ops.mlir_type(res) + T = Reactant.unwrapped_eltype(res) + if T <: Complex && !has_complex + res_array = if res isa TracedRNumber + TracedRArray{T,0}((), get_mlir_data(res), ()) + else + res + end + + real_part = Ops.real(res_array) + imag_part = Ops.imag(res_array) + + sz = size(res_array) + sz_reshaped = (1, sz...) + + real_reshaped = Ops.reshape(real_part, sz_reshaped...) + imag_reshaped = Ops.reshape(imag_part, sz_reshaped...) + + concated = Ops.concatenate([real_reshaped, imag_reshaped], 1) + col_maj = get_mlir_data(concated) + out_ty = Ops.mlir_type(concated) + + if do_transpose + col_maj = transpose_val(col_maj) + out_ty = transpose_ty(out_ty) + end + else + col_maj = get_mlir_data(res) + out_ty = Ops.mlir_type(res) - if do_transpose - col_maj = transpose_val(col_maj) - out_ty = transpose_ty(out_ty) + if do_transpose + col_maj = transpose_val(col_maj) + out_ty = transpose_ty(out_ty) + end end end @@ -927,8 +1014,14 @@ function finalize_mlir_fn( sharding.mesh.axis_names, size(sharding.mesh), )] + T = Reactant.unwrapped_eltype(arg) + sz = if T <: Complex && !has_complex + (2, size(arg)...) + else + size(arg) + end attr, dialect = Reactant.Sharding.get_tensor_sharding_attribute( - sharding, ctx, sym_name, mesh_attr, size(arg) + sharding, ctx, sym_name, mesh_attr, sz ) linear_arg_shardings[i] = (attr, dialect) if dialect == :sdy @@ -988,8 +1081,15 @@ function finalize_mlir_fn( ) haskey(mesh_cache, key) || @opcall(mesh(sharding.mesh)) (; sym_name, mesh_attr) = mesh_cache[key] + T = Reactant.unwrapped_eltype(arg) + if T <: Complex && !has_complex + sharding = Reactant.adapt_sharding_for_complex(sharding, ndims(arg)) + sz = (2, size(arg)...) + else + sz = size(arg) + end attr, dialect = Reactant.Sharding.get_tensor_sharding_attribute( - sharding, ctx, sym_name, mesh_attr, size(arg) + sharding, ctx, sym_name, mesh_attr, sz ) if dialect == :sdy MLIR.API.mlirFuncSetResultAttr(func2, i - 1, "sdy.sharding", attr) diff --git a/src/Types.jl b/src/Types.jl index 60590a249e..e80e361a27 100644 --- a/src/Types.jl +++ b/src/Types.jl @@ -244,6 +244,24 @@ function ConcretePJRTArray{T,N}( return ConcretePJRTArray{T,N,D}(data, shape, sharding) end +function adapt_sharding_for_complex(sharding::Sharding.NoSharding, N) + return Sharding.NoSharding() +end +function adapt_sharding_for_complex(sharding::Sharding.Replicated, N) + return Sharding.Replicated(sharding.mesh) +end +function adapt_sharding_for_complex(sharding::Sharding.NamedSharding, N) + return Sharding.NamedSharding( + sharding.mesh, + [[nothing]; sharding.partition_spec], + is_closed=(true, sharding.is_closed...), + priority=(-1, sharding.priority...), + ) +end +function adapt_sharding_for_complex(sharding::Sharding.AbstractSharding, N) + return adapt_sharding_for_complex(Sharding.NamedSharding(sharding, N), N) +end + function ConcretePJRTArray( data::Array{T,N}; client::Union{Nothing,XLA.PJRT.Client}=nothing, @@ -252,12 +270,23 @@ function ConcretePJRTArray( sharding::Sharding.AbstractSharding=Sharding.NoSharding(), ) where {T,N} theclient, thedevice = _select_client_and_device(client, idx, device, sharding) - sharded_data, shardinfo = sharding(theclient, thedevice, data) - shape = size(data) - nsharded = length(sharded_data) - return ConcretePJRTArray{T,N,nsharded}(sharded_data, shape, shardinfo) + if T <: Complex && !XLA.supports_complex(theclient) + T2 = real(T) + data_re = reinterpret(reshape, T2, data) + sharding_re = adapt_sharding_for_complex(sharding, N) + sharded_data, shardinfo = sharding_re(theclient, thedevice, data_re) + shape = size(data) + nsharded = length(sharded_data) + return ConcretePJRTArray{T,N,nsharded}(sharded_data, shape, shardinfo) + else + sharded_data, shardinfo = sharding(theclient, thedevice, data) + shape = size(data) + nsharded = length(sharded_data) + return ConcretePJRTArray{T,N,nsharded}(sharded_data, shape, shardinfo) + end end + Base.wait(x::Union{ConcretePJRTArray,ConcretePJRTNumber}) = foreach(wait, x.data) XLA.client(x::Union{ConcretePJRTArray,ConcretePJRTNumber}) = XLA.client(x.data) function XLA.device(x::Union{ConcretePJRTArray,ConcretePJRTNumber}) @@ -384,9 +413,17 @@ function ConcreteIFRTArray( ) where {T,N} theclient, thedevice = _select_client_and_device(client, idx, device, sharding) shape = size(data) - # ToDo: How to use specified device (non-sharded case)? - sharded_data, shardinfo, padding = sharding(theclient, nothing, data) - return ConcreteIFRTArray{T,N}(sharded_data, shape, shardinfo, padding) + if T <: Complex && !XLA.supports_complex(theclient) + T2 = real(T) + data_re = reinterpret(reshape, T2, data) + sharding_re = adapt_sharding_for_complex(sharding, N) + sharded_data, shardinfo, padding = sharding_re(theclient, nothing, data_re) + return ConcreteIFRTArray{T,N}(sharded_data, shape, shardinfo, padding) + else + # ToDo: How to use specified device (non-sharded case)? + sharded_data, shardinfo, padding = sharding(theclient, nothing, data) + return ConcreteIFRTArray{T,N}(sharded_data, shape, shardinfo, padding) + end end # Assemble data from multiple arrays. Needed in distributed setting where each process wont @@ -405,38 +442,77 @@ function ConcreteIFRTArray( client = client === nothing ? XLA.default_backend() : client - (; hlo_sharding) = Sharding.HloSharding(sharding, array_size) - all_devices = XLA.get_device.((client,), sharding.mesh.device_ids) - ifrt_sharding = XLA.IFRT.Sharding(all_devices, hlo_sharding) + if T <: Complex && !XLA.supports_complex(client) + T2 = real(T) + data_re = [reinterpret(reshape, T2, x) for x in data] + array_size_re = (2, array_size...) + sharding_re = adapt_sharding_for_complex(sharding, N) + + (; hlo_sharding) = Sharding.HloSharding(sharding_re, array_size_re) + all_devices = XLA.get_device.((client,), sharding_re.mesh.device_ids) + ifrt_sharding = XLA.IFRT.Sharding(all_devices, hlo_sharding) + + # Validate that all the slices are as we expected them to be + slices, _ = XLA.sharding_to_concrete_array_indices( + hlo_sharding, array_size_re, 0:(length(all_devices) - 1) + ) + addressable_slices = [ + slice for (slice, device) in zip(slices, all_devices) if XLA.is_addressable(device) + ] + for (i, slice) in enumerate(addressable_slices) + idx = findfirst(Base.Fix1(in, i), data_to_addressable_shard) + @assert idx !== nothing + @assert size(data_re[idx]) == length.(slice) "Expected data_re[$idx] to be at \ + $(slice), but got size \ + $(size(data_re[idx]))" + end - # Validate that all the slices are as we expected them to be - slices, _ = XLA.sharding_to_concrete_array_indices( - hlo_sharding, array_size, 0:(length(all_devices) - 1) - ) - addressable_slices = [ - slice for (slice, device) in zip(slices, all_devices) if XLA.is_addressable(device) - ] - for (i, slice) in enumerate(addressable_slices) - idx = findfirst(Base.Fix1(in, i), data_to_addressable_shard) - @assert idx !== nothing - @assert size(data[idx]) == length.(slice) "Expected data[$idx] to be at \ - $(slice), but got size \ - $(size(data[idx]))" - end + data_to_addressable_shard_copy = [copy(x) for x in data_to_addressable_shard] + @inbounds for shard_idxs in data_to_addressable_shard_copy + shard_idxs .-= 1 + end + ifrt_array = XLA.IFRT.AsyncArray( + XLA.IFRT.Array(client, data_re, data_to_addressable_shard_copy, array_size_re, ifrt_sharding), + nothing, + ) + return ConcreteIFRTArray{T,N}( + ifrt_array, array_size, Sharding.ShardInfo(sharding_re, slices) + ) + else + (; hlo_sharding) = Sharding.HloSharding(sharding, array_size) + all_devices = XLA.get_device.((client,), sharding.mesh.device_ids) + ifrt_sharding = XLA.IFRT.Sharding(all_devices, hlo_sharding) + + # Validate that all the slices are as we expected them to be + slices, _ = XLA.sharding_to_concrete_array_indices( + hlo_sharding, array_size, 0:(length(all_devices) - 1) + ) + addressable_slices = [ + slice for (slice, device) in zip(slices, all_devices) if XLA.is_addressable(device) + ] + for (i, slice) in enumerate(addressable_slices) + idx = findfirst(Base.Fix1(in, i), data_to_addressable_shard) + @assert idx !== nothing + @assert size(data[idx]) == length.(slice) "Expected data[$idx] to be at \ + $(slice), but got size \ + $(size(data[idx]))" + end - # Make the mapping 0-indexed - @inbounds for shard_idxs in data_to_addressable_shard - shard_idxs .-= 1 + # Make the mapping 0-indexed + @inbounds for shard_idxs in data_to_addressable_shard + shard_idxs .-= 1 + end + ifrt_array = XLA.IFRT.AsyncArray( + XLA.IFRT.Array(client, data, data_to_addressable_shard, array_size, ifrt_sharding), + nothing, + ) + return ConcreteIFRTArray{T,N}( + ifrt_array, array_size, Sharding.ShardInfo(sharding, slices) + ) end - ifrt_array = XLA.IFRT.AsyncArray( - XLA.IFRT.Array(client, data, data_to_addressable_shard, array_size, ifrt_sharding), - nothing, - ) - return ConcreteIFRTArray{T,N}( - ifrt_array, array_size, Sharding.ShardInfo(sharding, slices) - ) end + @enumx InterpolationType begin Nearest Linear @@ -731,10 +807,20 @@ function ConcretePJRTArray{T}( sharding::Sharding.AbstractSharding=Sharding.NoSharding(), ) where {T} theclient, thedevice = _select_client_and_device(client, idx, device, sharding) - sharded_data, shardinfo = sharding(theclient, thedevice, T, shape) - N = length(shape) - nsharded = length(sharded_data) - return ConcretePJRTArray{T,N,nsharded}(sharded_data, shape, shardinfo) + if T <: Complex && !XLA.supports_complex(theclient) + T2 = real(T) + shape_re = (2, shape...) + sharding_re = adapt_sharding_for_complex(sharding, length(shape)) + sharded_data, shardinfo = sharding_re(theclient, thedevice, T2, shape_re) + N = length(shape) + nsharded = length(sharded_data) + return ConcretePJRTArray{T,N,nsharded}(sharded_data, shape, shardinfo) + else + sharded_data, shardinfo = sharding(theclient, thedevice, T, shape) + N = length(shape) + nsharded = length(sharded_data) + return ConcretePJRTArray{T,N,nsharded}(sharded_data, shape, shardinfo) + end end function ConcreteIFRTArray{T}( @@ -747,13 +833,22 @@ function ConcreteIFRTArray{T}( ) where {T} theclient, thedevice = _select_client_and_device(client, idx, device, sharding) N = length(shape) - # ToDo: How to avoid allocating dummy array on host? - dummy_array = Array{T}(undef, shape) - # ToDo: How to use specified device (non-sharded case)? - sharded_data, shardinfo, padding = sharding(theclient, nothing, dummy_array) - return ConcreteIFRTArray{T,N}(sharded_data, shape, shardinfo, padding) + if T <: Complex && !XLA.supports_complex(theclient) + T2 = real(T) + dummy_array = Array{T2}(undef, (2, shape...)) + sharding_re = adapt_sharding_for_complex(sharding, N) + sharded_data, shardinfo, padding = sharding_re(theclient, nothing, dummy_array) + return ConcreteIFRTArray{T,N}(sharded_data, shape, shardinfo, padding) + else + # ToDo: How to avoid allocating dummy array on host? + dummy_array = Array{T}(undef, shape) + # ToDo: How to use specified device (non-sharded case)? + sharded_data, shardinfo, padding = sharding(theclient, nothing, dummy_array) + return ConcreteIFRTArray{T,N}(sharded_data, shape, shardinfo, padding) + end end + function _select_client_and_device( client::Union{Nothing,XLA.AbstractClient}, idx::Union{Int,Nothing}, diff --git a/src/compiler/Compiler.jl b/src/compiler/Compiler.jl index f275877eba..9d0113e953 100644 --- a/src/compiler/Compiler.jl +++ b/src/compiler/Compiler.jl @@ -119,6 +119,10 @@ function compile_mlir( mod, f; mlir_fn_res.num_partitions, mlir_fn_res.num_replicas ) + if !XLA.supports_complex(client) + run_pass_pipeline!(mod, "lower-complex-operations", "lower_complex_operations") + end + if drop_unsupported_attributes # Drop some of our attributes run_pass_pipeline!( @@ -283,6 +287,8 @@ function compile_mlir!( true; runtime, compile_options.optimize_then_pad, + client, + emulate_complex=true, kwargs..., ) finally diff --git a/src/utils.jl b/src/utils.jl index e293030cc5..7a32ba73ca 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1204,3 +1204,5 @@ nmantissa(::Type{Float64}) = 52 @inline _unwrap_val(x) = x @inline _unwrap_val(::Val{val}) where {val} = _unwrap_val(val) + + diff --git a/src/xla/Client.jl b/src/xla/Client.jl index ccf715c1ba..7319ad439b 100644 --- a/src/xla/Client.jl +++ b/src/xla/Client.jl @@ -14,3 +14,16 @@ function get_addressable_device end function platform_name end default_device(client::AbstractClient) = first(addressable_devices(client)) + +function supports_complex(client::Union{AbstractClient,Nothing}) + if Reactant.precompiling() + return true + end + client === nothing && return true + pname = lowercase(platform_name(client)) + if pname == "neuron" || pname == "trainium" + return false + end + return true +end + diff --git a/src/xla/IFRT/Array.jl b/src/xla/IFRT/Array.jl index 9ecd8436dc..498fdd57f8 100644 --- a/src/xla/IFRT/Array.jl +++ b/src/xla/IFRT/Array.jl @@ -18,7 +18,7 @@ end function Array( client::Client, - array::Base.Array{T,N}, + array::Base.StridedArray{T,N}, device::Device=XLA.default_device(client), memory_kind::AbstractString=string(convert(MemoryKind, XLA.default_memory(device))), ) where {T<:Reactant.ReactantPrimitive,N} @@ -38,7 +38,7 @@ function Array( end function Array( - client::Client, array::Base.Array{T,N}, sharding::Sharding + client::Client, array::Base.StridedArray{T,N}, sharding::Sharding ) where {T<:Reactant.ReactantPrimitive,N} all_devices = XLA.devices(sharding) all_logical_device_ids = collect(Int64, 0:(length(all_devices) - 1)) @@ -76,7 +76,7 @@ end function Array( client::Client, - host_buffers::Vector{Base.Array{T,N}}, + host_buffers::Vector{<:Base.StridedArray{T,N}}, addressable_shard_indices::Vector{Vector{Int64}}, array_shape, sharding::Sharding, diff --git a/src/xla/PJRT/Buffer.jl b/src/xla/PJRT/Buffer.jl index e4e7e65849..deb14b1b14 100644 --- a/src/xla/PJRT/Buffer.jl +++ b/src/xla/PJRT/Buffer.jl @@ -6,7 +6,7 @@ mutable struct Buffer <: XLA.AbstractBuffer end end -function Buffer(client::Client, array::Array{T,N}, device::Device) where {T,N} +function Buffer(client::Client, array::StridedArray{T,N}, device::Device) where {T,N} sizear = collect(Int64, reverse(size(array))) GC.@preserve client array device begin buffer = MLIR.API.ArrayFromHostBuffer( diff --git a/test/core/complex.jl b/test/core/complex.jl index 041cacfb58..d2acf15f8c 100644 --- a/test/core/complex.jl +++ b/test/core/complex.jl @@ -119,3 +119,31 @@ end @test @jit(isreal.(x_concrete)) == isreal.(x) end end + +@testset "forced complex emulation" begin + @eval Reactant.XLA function supports_complex(client::Union{AbstractClient,Nothing}) + return false + end + + try + @testset "emulated complex addition" begin + a = Reactant.to_rarray(ones(ComplexF32, 2)) + b = Reactant.to_rarray(ones(ComplexF32, 2)) + c = Reactant.compile(+, (a, b))(a, b) + @test c == ones(ComplexF32, 2) + ones(ComplexF32, 2) + end + finally + @eval Reactant.XLA function supports_complex(client::Union{AbstractClient,Nothing}) + if Reactant.precompiling() + return true + end + client === nothing && return true + pname = lowercase(platform_name(client)) + if pname == "neuron" || pname == "trainium" + return false + end + return true + end + end +end + From feecafe1dfcc9ca5c5c54e15e30999ac56deffef Mon Sep 17 00:00:00 2001 From: Billy Moses Date: Wed, 3 Jun 2026 16:28:42 -0500 Subject: [PATCH 2/9] fixup --- src/CompileOptions.jl | 11 ++++++++++- src/Reactant.jl | 6 +++--- src/TracedUtils.jl | 6 +++--- src/compiler/Compiler.jl | 18 ++++++++---------- src/compiler/OptimizationPasses.jl | 2 ++ test/core/complex.jl | 28 +++++----------------------- 6 files changed, 31 insertions(+), 40 deletions(-) diff --git a/src/CompileOptions.jl b/src/CompileOptions.jl index a4806185eb..cc75e419e6 100644 --- a/src/CompileOptions.jl +++ b/src/CompileOptions.jl @@ -220,6 +220,7 @@ struct CompileOptions raise_first::Bool # dialect specific options legalize_chlo_to_stablehlo::Bool + emulate_complex::Bool # backend specific options cudnn_hlo_optimize::Bool # sharding options @@ -263,6 +264,8 @@ function CompileOptions(; raise::Union{Bool,String}=false, raise_first::Bool=false, legalize_chlo_to_stablehlo::Bool=false, + client::Union{Nothing,XLA.AbstractClient}=nothing, + emulate_complex::Union{Nothing,Bool}=nothing, cudnn_hlo_optimize::Bool=false, shardy_passes::Union{Symbol,ShardyPropagationOptions}=:post_sdy_propagation, optimize_then_pad::Bool=true, @@ -316,6 +319,8 @@ function CompileOptions(; @assert shardy_passes in [:none, :to_mhlo_shardings, :post_sdy_propagation] end + emulate_complex_val = emulate_complex === nothing ? (client === nothing ? false : !XLA.supports_complex(client)) : emulate_complex + return CompileOptions( optimization_passes, no_nan, @@ -327,6 +332,7 @@ function CompileOptions(; raise, raise_first, legalize_chlo_to_stablehlo, + emulate_complex_val, cudnn_hlo_optimize, shardy_passes, optimize_then_pad, @@ -357,10 +363,11 @@ end function __compile_options_from_kwargs(; compile_options::Union{Missing,CompileOptions}=missing, optimize::Union{Bool,Symbol,String}=true, + client::Union{Nothing,XLA.AbstractClient}=nothing, kwargs..., ) compile_options isa CompileOptions && return compile_options - return CompileOptions(; optimization_passes=optimize, kwargs...) + return CompileOptions(; optimization_passes=optimize, client, kwargs...) end function __reverse_propagation(sym::Symbol) @@ -382,6 +389,7 @@ function __compile_options_with_reversed_propagation(compile_options::CompileOpt compile_options.raise, compile_options.raise_first, compile_options.legalize_chlo_to_stablehlo, + compile_options.emulate_complex, compile_options.cudnn_hlo_optimize, compile_options.shardy_passes, compile_options.optimize_then_pad, @@ -424,6 +432,7 @@ function __compile_options_with_updated_sync(compile_options::CompileOptions, sy compile_options.raise, compile_options.raise_first, compile_options.legalize_chlo_to_stablehlo, + compile_options.emulate_complex, compile_options.cudnn_hlo_optimize, compile_options.shardy_passes, compile_options.optimize_then_pad, diff --git a/src/Reactant.jl b/src/Reactant.jl index dac3cf1be6..6f0b67a808 100644 --- a/src/Reactant.jl +++ b/src/Reactant.jl @@ -92,14 +92,14 @@ _parent_type(::Type{<:Slices{P}}) where {P} = P include("accelerators/Accelerators.jl") +include("mlir/MLIR.jl") +include("xla/XLA.jl") + include("CompileOptions.jl") export OptimizeCommunicationOptions, ShardyPropagationOptions, CompileOptions, MultiFloatOptions -include("mlir/MLIR.jl") -include("xla/XLA.jl") - include("Configuration.jl") include("Sharding.jl") include("Devices.jl") diff --git a/src/TracedUtils.jl b/src/TracedUtils.jl index f43524de6f..ec455ec1a1 100644 --- a/src/TracedUtils.jl +++ b/src/TracedUtils.jl @@ -510,7 +510,7 @@ function prepare_mlir_fn_args( inv_map[v] = k end - has_complex = !emulate_complex || XLA.supports_complex(client) + has_complex = !emulate_complex in_tys = Vector{MLIR.IR.Type}(undef, length(linear_args)) for (i, arg) in enumerate(linear_args) @@ -624,7 +624,7 @@ function prepare_mlir_fn_args( end function process_linear_args!(linear_args, fnbody, do_transpose, optimize_then_pad, inv_map; client=nothing, emulate_complex=false) - has_complex = !emulate_complex || XLA.supports_complex(client) + has_complex = !emulate_complex for (i, arg) in enumerate(linear_args) raw_arg = MLIR.IR.argument(fnbody, i) row_maj_arg = do_transpose ? transpose_val(raw_arg) : raw_arg @@ -701,7 +701,7 @@ function finalize_mlir_fn( client=nothing, emulate_complex=false, ) - has_complex = !emulate_complex || XLA.supports_complex(client) + has_complex = !emulate_complex # check which arguments have been mutated mutated_args = Int[] diff --git a/src/compiler/Compiler.jl b/src/compiler/Compiler.jl index 9d0113e953..1601428f58 100644 --- a/src/compiler/Compiler.jl +++ b/src/compiler/Compiler.jl @@ -67,7 +67,7 @@ function compile_mlir( try mod = MLIR.IR.Module(MLIR.IR.Location()) - compile_options, kwargs_inner = __get_compile_options_and_kwargs(; kwargs...) + compile_options, kwargs_inner = __get_compile_options_and_kwargs(; client, kwargs...) # Wrap compile_mlir! to catch pass pipeline failures mlir_fn_res = try @@ -119,7 +119,7 @@ function compile_mlir( mod, f; mlir_fn_res.num_partitions, mlir_fn_res.num_replicas ) - if !XLA.supports_complex(client) + if compile_options.emulate_complex run_pass_pipeline!(mod, "lower-complex-operations", "lower_complex_operations") end @@ -288,7 +288,7 @@ function compile_mlir!( runtime, compile_options.optimize_then_pad, client, - emulate_complex=true, + emulate_complex=compile_options.emulate_complex, kwargs..., ) finally @@ -1331,7 +1331,7 @@ function compile_xla( # compile function to MLIR module mod = MLIR.IR.Module(MLIR.IR.Location()) - compile_options, kwargs = __get_compile_options_and_kwargs(; kwargs...) + compile_options, kwargs = __get_compile_options_and_kwargs(; client, kwargs...) mlir_fn_res = compile_mlir!( mod, @@ -1410,7 +1410,7 @@ function compile_xla( finalize(mod) - return exec, hlo_modules, mlir_fn_res, device, client, module_string + return exec, hlo_modules, mlir_fn_res, device, client, module_string, compile_options finally MLIR.IR.deactivate(ctx) end @@ -1422,11 +1422,9 @@ function compile(f, args; kwargs...) end end -function compile(ctx, f, args; kwargs...) - compile_options, kwargs = __get_compile_options_and_kwargs(; kwargs...) - - exec, _, mlir_fn_res, device, client, str = compile_xla( - ctx, f, args; compile_options, kwargs... +function compile(ctx, f, args; client=nothing, kwargs...) + exec, _, mlir_fn_res, device, client, str, compile_options = compile_xla( + ctx, f, args; client, kwargs... ) (; linear_args, seen_args, linear_results, preserved_args, concrete_result) = mlir_fn_res diff --git a/src/compiler/OptimizationPasses.jl b/src/compiler/OptimizationPasses.jl index 9838645885..9c01f380cf 100644 --- a/src/compiler/OptimizationPasses.jl +++ b/src/compiler/OptimizationPasses.jl @@ -232,6 +232,7 @@ function __get_compile_options_and_kwargs(; raise::Union{Bool,String}=false, raise_first::Bool=false, legalize_chlo_to_stablehlo::Bool=false, + emulate_complex::Union{Nothing,Bool}=nothing, cudnn_hlo_optimize::Bool=false, shardy_passes::Union{Symbol,ShardyPropagationOptions}=:post_sdy_propagation, optimize_then_pad::Bool=true, @@ -259,6 +260,7 @@ function __get_compile_options_and_kwargs(; raise, raise_first, legalize_chlo_to_stablehlo, + emulate_complex, cudnn_hlo_optimize, shardy_passes, optimize_then_pad, diff --git a/test/core/complex.jl b/test/core/complex.jl index d2acf15f8c..10581a228c 100644 --- a/test/core/complex.jl +++ b/test/core/complex.jl @@ -121,29 +121,11 @@ end end @testset "forced complex emulation" begin - @eval Reactant.XLA function supports_complex(client::Union{AbstractClient,Nothing}) - return false - end - - try - @testset "emulated complex addition" begin - a = Reactant.to_rarray(ones(ComplexF32, 2)) - b = Reactant.to_rarray(ones(ComplexF32, 2)) - c = Reactant.compile(+, (a, b))(a, b) - @test c == ones(ComplexF32, 2) + ones(ComplexF32, 2) - end - finally - @eval Reactant.XLA function supports_complex(client::Union{AbstractClient,Nothing}) - if Reactant.precompiling() - return true - end - client === nothing && return true - pname = lowercase(platform_name(client)) - if pname == "neuron" || pname == "trainium" - return false - end - return true - end + @testset "emulated complex addition" begin + a = Reactant.to_rarray(ones(ComplexF32, 2)) + b = Reactant.to_rarray(ones(ComplexF32, 2)) + c = Reactant.compile(+, (a, b); emulate_complex=true)(a, b) + @test c == ones(ComplexF32, 2) + ones(ComplexF32, 2) end end From ccca1c47c1ca9820ca007b94d4031c096a640658 Mon Sep 17 00:00:00 2001 From: Billy Moses Date: Wed, 3 Jun 2026 16:31:33 -0500 Subject: [PATCH 3/9] sz 3 --- test/core/complex.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/core/complex.jl b/test/core/complex.jl index 10581a228c..ce222b75bd 100644 --- a/test/core/complex.jl +++ b/test/core/complex.jl @@ -122,10 +122,10 @@ end @testset "forced complex emulation" begin @testset "emulated complex addition" begin - a = Reactant.to_rarray(ones(ComplexF32, 2)) - b = Reactant.to_rarray(ones(ComplexF32, 2)) + a = Reactant.to_rarray(ones(ComplexF32, 3)) + b = Reactant.to_rarray(ones(ComplexF32, 3)) c = Reactant.compile(+, (a, b); emulate_complex=true)(a, b) - @test c == ones(ComplexF32, 2) + ones(ComplexF32, 2) + @test c == ones(ComplexF32, 3) + ones(ComplexF32, 3) end end From b0b8e83a673f574d004bca600965f43e6ca504d9 Mon Sep 17 00:00:00 2001 From: Billy Moses Date: Wed, 3 Jun 2026 16:51:52 -0500 Subject: [PATCH 4/9] fixup --- src/ConcreteRArray.jl | 49 +++-------- src/Types.jl | 181 ++++++++++------------------------------- src/xla/IFRT/Array.jl | 38 +++++++-- src/xla/PJRT/Buffer.jl | 25 +++++- 4 files changed, 106 insertions(+), 187 deletions(-) diff --git a/src/ConcreteRArray.jl b/src/ConcreteRArray.jl index d2402559c2..414ea4222c 100644 --- a/src/ConcreteRArray.jl +++ b/src/ConcreteRArray.jl @@ -106,36 +106,19 @@ function Base.isempty(x::Union{WrappedConcretePJRTArray,WrappedConcreteIFRTArray end function Base.convert(::Type{<:Array}, X::AbstractConcreteArray{T,N}) where {T,N} - client = XLA.client(X) - if T <: Complex && !XLA.supports_complex(client) - T2 = real(T) - if has_padding(X) - padding = get_padding(X) - data = Array{T,N}(undef, (size(X) .+ padding)...) - data_re = reinterpret(reshape, T2, data) - write_to_host_buffer!(data_re, X) - return view(data, [1:size(X, i) for i in 1:ndims(X)]...) - else - data = Array{T,N}(undef, size(X)...) - data_re = reinterpret(reshape, T2, data) - write_to_host_buffer!(data_re, X) - return data - end + if has_padding(X) + padding = get_padding(X) + data = Array{T,N}(undef, (size(X) .+ padding)...) + write_to_host_buffer!(data, X) + return view(data, [1:size(X, i) for i in 1:ndims(X)]...) else - if has_padding(X) - padding = get_padding(X) - data = Array{T,N}(undef, (size(X) .+ padding)...) - write_to_host_buffer!(data, X) - return view(data, [1:size(X, i) for i in 1:ndims(X)]...) - else - data = Array{T,N}(undef, size(X)...) - write_to_host_buffer!(data, X) - return data - end + data = Array{T,N}(undef, size(X)...) + write_to_host_buffer!(data, X) + return data end end -function write_to_host_buffer!(data::AbstractArray, X::ConcretePJRTArray{T,N}) where {T,N} +function write_to_host_buffer!(data::Array, X::ConcretePJRTArray{T,N}) where {T,N} if Sharding.is_sharded(X) completed = Set{eltype(X.sharding.device_to_array_slices)}() for idx in 1:length(X.data) @@ -152,12 +135,11 @@ function write_to_host_buffer!(data::AbstractArray, X::ConcretePJRTArray{T,N}) w return nothing end -function write_to_host_buffer!(data::AbstractArray, X::ConcreteIFRTArray{T,N}) where {T,N} +function write_to_host_buffer!(data::Array, X::ConcreteIFRTArray{T,N}) where {T,N} XLA.to_host(X.data, data, X.sharding) return nothing end - function Base.convert( ::Type{<:Array}, X::Union{WrappedConcretePJRTArray,WrappedConcreteIFRTArray} ) @@ -725,16 +707,7 @@ for aType in (:ConcretePJRTArray, :ConcreteIFRTArray) function Base.copyto!( dest::Array, src::Broadcast.Broadcasted{Broadcast.ArrayStyle{$(aType)}} ) - copied_src = copy(src) - client = XLA.client(copied_src) - T = eltype(copied_src) - if T <: Complex && !XLA.supports_complex(client) - T2 = real(T) - dest_re = reinterpret(reshape, T2, dest) - write_to_host_buffer!(dest_re, copied_src) - else - write_to_host_buffer!(dest, copied_src) - end + write_to_host_buffer!(dest, copy(src)) return dest end diff --git a/src/Types.jl b/src/Types.jl index e80e361a27..60590a249e 100644 --- a/src/Types.jl +++ b/src/Types.jl @@ -244,24 +244,6 @@ function ConcretePJRTArray{T,N}( return ConcretePJRTArray{T,N,D}(data, shape, sharding) end -function adapt_sharding_for_complex(sharding::Sharding.NoSharding, N) - return Sharding.NoSharding() -end -function adapt_sharding_for_complex(sharding::Sharding.Replicated, N) - return Sharding.Replicated(sharding.mesh) -end -function adapt_sharding_for_complex(sharding::Sharding.NamedSharding, N) - return Sharding.NamedSharding( - sharding.mesh, - [[nothing]; sharding.partition_spec], - is_closed=(true, sharding.is_closed...), - priority=(-1, sharding.priority...), - ) -end -function adapt_sharding_for_complex(sharding::Sharding.AbstractSharding, N) - return adapt_sharding_for_complex(Sharding.NamedSharding(sharding, N), N) -end - function ConcretePJRTArray( data::Array{T,N}; client::Union{Nothing,XLA.PJRT.Client}=nothing, @@ -270,23 +252,12 @@ function ConcretePJRTArray( sharding::Sharding.AbstractSharding=Sharding.NoSharding(), ) where {T,N} theclient, thedevice = _select_client_and_device(client, idx, device, sharding) - if T <: Complex && !XLA.supports_complex(theclient) - T2 = real(T) - data_re = reinterpret(reshape, T2, data) - sharding_re = adapt_sharding_for_complex(sharding, N) - sharded_data, shardinfo = sharding_re(theclient, thedevice, data_re) - shape = size(data) - nsharded = length(sharded_data) - return ConcretePJRTArray{T,N,nsharded}(sharded_data, shape, shardinfo) - else - sharded_data, shardinfo = sharding(theclient, thedevice, data) - shape = size(data) - nsharded = length(sharded_data) - return ConcretePJRTArray{T,N,nsharded}(sharded_data, shape, shardinfo) - end + sharded_data, shardinfo = sharding(theclient, thedevice, data) + shape = size(data) + nsharded = length(sharded_data) + return ConcretePJRTArray{T,N,nsharded}(sharded_data, shape, shardinfo) end - Base.wait(x::Union{ConcretePJRTArray,ConcretePJRTNumber}) = foreach(wait, x.data) XLA.client(x::Union{ConcretePJRTArray,ConcretePJRTNumber}) = XLA.client(x.data) function XLA.device(x::Union{ConcretePJRTArray,ConcretePJRTNumber}) @@ -413,17 +384,9 @@ function ConcreteIFRTArray( ) where {T,N} theclient, thedevice = _select_client_and_device(client, idx, device, sharding) shape = size(data) - if T <: Complex && !XLA.supports_complex(theclient) - T2 = real(T) - data_re = reinterpret(reshape, T2, data) - sharding_re = adapt_sharding_for_complex(sharding, N) - sharded_data, shardinfo, padding = sharding_re(theclient, nothing, data_re) - return ConcreteIFRTArray{T,N}(sharded_data, shape, shardinfo, padding) - else - # ToDo: How to use specified device (non-sharded case)? - sharded_data, shardinfo, padding = sharding(theclient, nothing, data) - return ConcreteIFRTArray{T,N}(sharded_data, shape, shardinfo, padding) - end + # ToDo: How to use specified device (non-sharded case)? + sharded_data, shardinfo, padding = sharding(theclient, nothing, data) + return ConcreteIFRTArray{T,N}(sharded_data, shape, shardinfo, padding) end # Assemble data from multiple arrays. Needed in distributed setting where each process wont @@ -442,77 +405,38 @@ function ConcreteIFRTArray( client = client === nothing ? XLA.default_backend() : client - if T <: Complex && !XLA.supports_complex(client) - T2 = real(T) - data_re = [reinterpret(reshape, T2, x) for x in data] - array_size_re = (2, array_size...) - sharding_re = adapt_sharding_for_complex(sharding, N) - - (; hlo_sharding) = Sharding.HloSharding(sharding_re, array_size_re) - all_devices = XLA.get_device.((client,), sharding_re.mesh.device_ids) - ifrt_sharding = XLA.IFRT.Sharding(all_devices, hlo_sharding) - - # Validate that all the slices are as we expected them to be - slices, _ = XLA.sharding_to_concrete_array_indices( - hlo_sharding, array_size_re, 0:(length(all_devices) - 1) - ) - addressable_slices = [ - slice for (slice, device) in zip(slices, all_devices) if XLA.is_addressable(device) - ] - for (i, slice) in enumerate(addressable_slices) - idx = findfirst(Base.Fix1(in, i), data_to_addressable_shard) - @assert idx !== nothing - @assert size(data_re[idx]) == length.(slice) "Expected data_re[$idx] to be at \ - $(slice), but got size \ - $(size(data_re[idx]))" - end + (; hlo_sharding) = Sharding.HloSharding(sharding, array_size) + all_devices = XLA.get_device.((client,), sharding.mesh.device_ids) + ifrt_sharding = XLA.IFRT.Sharding(all_devices, hlo_sharding) - data_to_addressable_shard_copy = [copy(x) for x in data_to_addressable_shard] - @inbounds for shard_idxs in data_to_addressable_shard_copy - shard_idxs .-= 1 - end - ifrt_array = XLA.IFRT.AsyncArray( - XLA.IFRT.Array(client, data_re, data_to_addressable_shard_copy, array_size_re, ifrt_sharding), - nothing, - ) - return ConcreteIFRTArray{T,N}( - ifrt_array, array_size, Sharding.ShardInfo(sharding_re, slices) - ) - else - (; hlo_sharding) = Sharding.HloSharding(sharding, array_size) - all_devices = XLA.get_device.((client,), sharding.mesh.device_ids) - ifrt_sharding = XLA.IFRT.Sharding(all_devices, hlo_sharding) - - # Validate that all the slices are as we expected them to be - slices, _ = XLA.sharding_to_concrete_array_indices( - hlo_sharding, array_size, 0:(length(all_devices) - 1) - ) - addressable_slices = [ - slice for (slice, device) in zip(slices, all_devices) if XLA.is_addressable(device) - ] - for (i, slice) in enumerate(addressable_slices) - idx = findfirst(Base.Fix1(in, i), data_to_addressable_shard) - @assert idx !== nothing - @assert size(data[idx]) == length.(slice) "Expected data[$idx] to be at \ - $(slice), but got size \ - $(size(data[idx]))" - end + # Validate that all the slices are as we expected them to be + slices, _ = XLA.sharding_to_concrete_array_indices( + hlo_sharding, array_size, 0:(length(all_devices) - 1) + ) + addressable_slices = [ + slice for (slice, device) in zip(slices, all_devices) if XLA.is_addressable(device) + ] + for (i, slice) in enumerate(addressable_slices) + idx = findfirst(Base.Fix1(in, i), data_to_addressable_shard) + @assert idx !== nothing + @assert size(data[idx]) == length.(slice) "Expected data[$idx] to be at \ + $(slice), but got size \ + $(size(data[idx]))" + end - # Make the mapping 0-indexed - @inbounds for shard_idxs in data_to_addressable_shard - shard_idxs .-= 1 - end - ifrt_array = XLA.IFRT.AsyncArray( - XLA.IFRT.Array(client, data, data_to_addressable_shard, array_size, ifrt_sharding), - nothing, - ) - return ConcreteIFRTArray{T,N}( - ifrt_array, array_size, Sharding.ShardInfo(sharding, slices) - ) + # Make the mapping 0-indexed + @inbounds for shard_idxs in data_to_addressable_shard + shard_idxs .-= 1 end + ifrt_array = XLA.IFRT.AsyncArray( + XLA.IFRT.Array(client, data, data_to_addressable_shard, array_size, ifrt_sharding), + nothing, + ) + return ConcreteIFRTArray{T,N}( + ifrt_array, array_size, Sharding.ShardInfo(sharding, slices) + ) end - @enumx InterpolationType begin Nearest Linear @@ -807,20 +731,10 @@ function ConcretePJRTArray{T}( sharding::Sharding.AbstractSharding=Sharding.NoSharding(), ) where {T} theclient, thedevice = _select_client_and_device(client, idx, device, sharding) - if T <: Complex && !XLA.supports_complex(theclient) - T2 = real(T) - shape_re = (2, shape...) - sharding_re = adapt_sharding_for_complex(sharding, length(shape)) - sharded_data, shardinfo = sharding_re(theclient, thedevice, T2, shape_re) - N = length(shape) - nsharded = length(sharded_data) - return ConcretePJRTArray{T,N,nsharded}(sharded_data, shape, shardinfo) - else - sharded_data, shardinfo = sharding(theclient, thedevice, T, shape) - N = length(shape) - nsharded = length(sharded_data) - return ConcretePJRTArray{T,N,nsharded}(sharded_data, shape, shardinfo) - end + sharded_data, shardinfo = sharding(theclient, thedevice, T, shape) + N = length(shape) + nsharded = length(sharded_data) + return ConcretePJRTArray{T,N,nsharded}(sharded_data, shape, shardinfo) end function ConcreteIFRTArray{T}( @@ -833,22 +747,13 @@ function ConcreteIFRTArray{T}( ) where {T} theclient, thedevice = _select_client_and_device(client, idx, device, sharding) N = length(shape) - if T <: Complex && !XLA.supports_complex(theclient) - T2 = real(T) - dummy_array = Array{T2}(undef, (2, shape...)) - sharding_re = adapt_sharding_for_complex(sharding, N) - sharded_data, shardinfo, padding = sharding_re(theclient, nothing, dummy_array) - return ConcreteIFRTArray{T,N}(sharded_data, shape, shardinfo, padding) - else - # ToDo: How to avoid allocating dummy array on host? - dummy_array = Array{T}(undef, shape) - # ToDo: How to use specified device (non-sharded case)? - sharded_data, shardinfo, padding = sharding(theclient, nothing, dummy_array) - return ConcreteIFRTArray{T,N}(sharded_data, shape, shardinfo, padding) - end + # ToDo: How to avoid allocating dummy array on host? + dummy_array = Array{T}(undef, shape) + # ToDo: How to use specified device (non-sharded case)? + sharded_data, shardinfo, padding = sharding(theclient, nothing, dummy_array) + return ConcreteIFRTArray{T,N}(sharded_data, shape, shardinfo, padding) end - function _select_client_and_device( client::Union{Nothing,XLA.AbstractClient}, idx::Union{Int,Nothing}, diff --git a/src/xla/IFRT/Array.jl b/src/xla/IFRT/Array.jl index 498fdd57f8..612df5c96a 100644 --- a/src/xla/IFRT/Array.jl +++ b/src/xla/IFRT/Array.jl @@ -18,17 +18,26 @@ end function Array( client::Client, - array::Base.StridedArray{T,N}, + array::Base.Array{T,N}, device::Device=XLA.default_device(client), memory_kind::AbstractString=string(convert(MemoryKind, XLA.default_memory(device))), ) where {T<:Reactant.ReactantPrimitive,N} + ndims = N + sizear = collect(Int64, reverse(size(array))) + ptype = XLA.primitive_type(T) + if T <: Complex && !XLA.supports_complex(client) + ndims += 1 + ptype = XLA.primitive_type(real(T)) + pushfirst!(sizear, 2) + end + GC.@preserve client device begin buffer = MLIR.API.ifrt_client_make_single_shard_array_from_host_buffer( client.client, array, - XLA.primitive_type(T), - N, - collect(Int64, reverse(size(array))), + ptype, + ndims, + sizear, 0, device.device, string(memory_kind), @@ -38,7 +47,7 @@ function Array( end function Array( - client::Client, array::Base.StridedArray{T,N}, sharding::Sharding + client::Client, array::Base.Array{T,N}, sharding::Sharding ) where {T<:Reactant.ReactantPrimitive,N} all_devices = XLA.devices(sharding) all_logical_device_ids = collect(Int64, 0:(length(all_devices) - 1)) @@ -100,11 +109,18 @@ function Array( end array_shape = collect(Int64, reverse(array_shape)) + ndims = length(array_shape) + + if T <: Complex && !XLA.supports_complex(client) + pushfirst!(array_shape, 2) + ndims += 1 + @assert Reactant.Sharding.is_replicated(sharding) "Complex number emulation on sharded arrays is not yet supported." + end GC.@preserve client sharding begin buffer = MLIR.API.ifrt_client_assemble_array_from_single_shards( client.client, - length(array_shape), + ndims, array_shape, sharding.ptr, length(single_device_arrays), @@ -117,8 +133,16 @@ function Array( end function Array( - client::Client, array::Base.Array{T,N}, sharding + client::Client, array::Base.StridedArray{T,N}, sharding ) where {T<:Reactant.ReactantPrimitive,N} + if T <: Complex && !XLA.supports_complex(client) + @assert sharding isa Reactant.Sharding.NoSharding || sharding isa Reactant.Sharding.Replicated + T2 = real(T) + array_re = reinterpret(reshape, T2, array) + sharding_re = Reactant.Sharding.HloSharding(sharding, size(array_re)) + return Array(client, array_re, sharding_re) + end + @assert sharding isa Reactant.Sharding.AbstractSharding if !(sharding isa Reactant.Sharding.HloSharding) sharding = Reactant.Sharding.HloSharding(sharding, size(array)) diff --git a/src/xla/PJRT/Buffer.jl b/src/xla/PJRT/Buffer.jl index deb14b1b14..f342f43d76 100644 --- a/src/xla/PJRT/Buffer.jl +++ b/src/xla/PJRT/Buffer.jl @@ -7,10 +7,17 @@ mutable struct Buffer <: XLA.AbstractBuffer end function Buffer(client::Client, array::StridedArray{T,N}, device::Device) where {T,N} + ndims = N sizear = collect(Int64, reverse(size(array))) + ptype = XLA.primitive_type(T) + if T <: Complex && !XLA.supports_complex(client) + ndims += 1 + ptype = XLA.primitive_type(real(T)) + pushfirst!(sizear, 2) + end GC.@preserve client array device begin buffer = MLIR.API.ArrayFromHostBuffer( - client.client, pointer(array), XLA.primitive_type(T), N, sizear, device.device + client.client, pointer(array), ptype, ndims, sizear, device.device ) end return Buffer(buffer) @@ -85,14 +92,24 @@ end end end + ndims = length(dims) + ptype = XLA.primitive_type(S) + sizear = collect(Int64, reverse(dims)) + + if S <: Complex && !XLA.supports_complex(client) + ndims += 1 + ptype = XLA.primitive_type(real(S)) + pushfirst!(sizear, 2) + end + GC.@preserve client device begin return Buffer( MLIR.API.UninitPJRTBuffer( client.client, device.device, - XLA.primitive_type(S), - length(dims), - collect(Int64, reverse(dims)), + ptype, + ndims, + sizear, ), ) end From 898d474161acc7daeab903e1011cadb91076d998 Mon Sep 17 00:00:00 2001 From: Billy Moses Date: Wed, 3 Jun 2026 17:07:38 -0500 Subject: [PATCH 5/9] more fixup --- src/TracedUtils.jl | 50 ++++++++++++++++++------------------------ src/xla/Client.jl | 3 +-- src/xla/IFRT/Array.jl | 4 ++-- src/xla/PJRT/Buffer.jl | 2 +- 4 files changed, 25 insertions(+), 34 deletions(-) diff --git a/src/TracedUtils.jl b/src/TracedUtils.jl index ec455ec1a1..8a9bb092dc 100644 --- a/src/TracedUtils.jl +++ b/src/TracedUtils.jl @@ -318,7 +318,6 @@ function make_mlir_fn( resargprefix::Symbol=:resargs, num_replicas=1, optimize_then_pad::Bool=true, - client=nothing, emulate_complex::Bool=false, ) if sizeof(typeof(f)) != 0 || f isa Base.BroadcastFunction @@ -341,7 +340,6 @@ function make_mlir_fn( resprefix, resargprefix, num_replicas, - client, emulate_complex, ) mlir_fn_res.fnwrapped = true @@ -358,8 +356,7 @@ function make_mlir_fn( optimize_then_pad, do_transpose, input_shardings, - verify_arg_names; - client, + verify_arg_names, emulate_complex, ) @@ -371,7 +368,7 @@ function make_mlir_fn( MLIR.IR.activate(fnbody) result = try - process_linear_args!(linear_args, fnbody, do_transpose, optimize_then_pad, inv_map; client, emulate_complex) + process_linear_args!(linear_args, fnbody, do_transpose, optimize_then_pad, inv_map, emulate_complex) if isempty(kwargs) Reactant.call_with_reactant(f, traced_args...) @@ -423,8 +420,7 @@ function make_mlir_fn( args, N, concretein, - toscalar; - client, + toscalar, emulate_complex, ) @@ -465,9 +461,8 @@ function prepare_mlir_fn_args( optimize_then_pad, do_transpose, input_shardings, - verify_arg_names; - client=nothing, - emulate_complex=false, + verify_arg_names, + emulate_complex, ) N = length(args) traced_args = Vector{Any}(undef, N) @@ -510,12 +505,10 @@ function prepare_mlir_fn_args( inv_map[v] = k end - has_complex = !emulate_complex - in_tys = Vector{MLIR.IR.Type}(undef, length(linear_args)) for (i, arg) in enumerate(linear_args) T = Reactant.unwrapped_eltype(arg) - if T <: Complex && !has_complex + if T <: Complex && emulate_complex elT = MLIR.IR.Type(real(T)) if toscalar in_tys[i] = MLIR.IR.TensorType(Int[2], elT) @@ -623,8 +616,7 @@ function prepare_mlir_fn_args( ) end -function process_linear_args!(linear_args, fnbody, do_transpose, optimize_then_pad, inv_map; client=nothing, emulate_complex=false) - has_complex = !emulate_complex +function process_linear_args!(linear_args, fnbody, do_transpose, optimize_then_pad, inv_map, emulate_complex) for (i, arg) in enumerate(linear_args) raw_arg = MLIR.IR.argument(fnbody, i) row_maj_arg = do_transpose ? transpose_val(raw_arg) : raw_arg @@ -642,15 +634,18 @@ function process_linear_args!(linear_args, fnbody, do_transpose, optimize_then_p end T = Reactant.unwrapped_eltype(arg) - if T <: Complex && !has_complex + if T <: Complex && emulate_complex T2 = real(T) sz = size(arg) N = length(sz) - start_real = [1; ones(Int, N)] - limit_real = [1; collect(Int, sz)] - start_imag = [2; ones(Int, N)] - limit_imag = [2; collect(Int, sz)] + start_real = ones(Int, N+1) + limit_real = collect(Int, sz) + pushfirst!(limit_real, 1) + start_imag = ones(Int, N+1) + start_imag[1] = 2 + limit_imag = collect(Int, sz) + pushfirst!(limit_real, 2) traced_raw = TracedRArray{T2, N+1}((), row_maj_arg, (2, sz...)) @@ -697,12 +692,9 @@ function finalize_mlir_fn( args, N, concretein, - toscalar; - client=nothing, - emulate_complex=false, + toscalar, + emulate_complex, ) - has_complex = !emulate_complex - # check which arguments have been mutated mutated_args = Int[] if !construct_function_without_args @@ -917,7 +909,7 @@ function finalize_mlir_fn( out_ty = Ops.mlir_type(TracedRArray{Bool,0}, ()) else T = Reactant.unwrapped_eltype(res) - if T <: Complex && !has_complex + if T <: Complex && emulate_complex res_array = if res isa TracedRNumber TracedRArray{T,0}((), get_mlir_data(res), ()) else @@ -1015,7 +1007,7 @@ function finalize_mlir_fn( size(sharding.mesh), )] T = Reactant.unwrapped_eltype(arg) - sz = if T <: Complex && !has_complex + sz = if T <: Complex && emulate_complex (2, size(arg)...) else size(arg) @@ -1082,8 +1074,8 @@ function finalize_mlir_fn( haskey(mesh_cache, key) || @opcall(mesh(sharding.mesh)) (; sym_name, mesh_attr) = mesh_cache[key] T = Reactant.unwrapped_eltype(arg) - if T <: Complex && !has_complex - sharding = Reactant.adapt_sharding_for_complex(sharding, ndims(arg)) + if T <: Complex && emulate_complex + sharding = Reactant.adapt_sharding_for_emulated_complex(sharding, ndims(arg)) sz = (2, size(arg)...) else sz = size(arg) diff --git a/src/xla/Client.jl b/src/xla/Client.jl index 7319ad439b..f51611fe87 100644 --- a/src/xla/Client.jl +++ b/src/xla/Client.jl @@ -15,11 +15,10 @@ function platform_name end default_device(client::AbstractClient) = first(addressable_devices(client)) -function supports_complex(client::Union{AbstractClient,Nothing}) +function supports_complex(client::AbstractClient) if Reactant.precompiling() return true end - client === nothing && return true pname = lowercase(platform_name(client)) if pname == "neuron" || pname == "trainium" return false diff --git a/src/xla/IFRT/Array.jl b/src/xla/IFRT/Array.jl index 612df5c96a..9b7be376b1 100644 --- a/src/xla/IFRT/Array.jl +++ b/src/xla/IFRT/Array.jl @@ -85,7 +85,7 @@ end function Array( client::Client, - host_buffers::Vector{<:Base.StridedArray{T,N}}, + host_buffers::Vector{Base.Array{T,N}}, addressable_shard_indices::Vector{Vector{Int64}}, array_shape, sharding::Sharding, @@ -133,7 +133,7 @@ function Array( end function Array( - client::Client, array::Base.StridedArray{T,N}, sharding + client::Client, array::Base.Array{T,N}, sharding ) where {T<:Reactant.ReactantPrimitive,N} if T <: Complex && !XLA.supports_complex(client) @assert sharding isa Reactant.Sharding.NoSharding || sharding isa Reactant.Sharding.Replicated diff --git a/src/xla/PJRT/Buffer.jl b/src/xla/PJRT/Buffer.jl index f342f43d76..f8dc47efd0 100644 --- a/src/xla/PJRT/Buffer.jl +++ b/src/xla/PJRT/Buffer.jl @@ -6,7 +6,7 @@ mutable struct Buffer <: XLA.AbstractBuffer end end -function Buffer(client::Client, array::StridedArray{T,N}, device::Device) where {T,N} +function Buffer(client::Client, array::Array{T,N}, device::Device) where {T,N} ndims = N sizear = collect(Int64, reverse(size(array))) ptype = XLA.primitive_type(T) From 9e9e9b7da02030c673889547d57c2a6f33bcbdd1 Mon Sep 17 00:00:00 2001 From: Billy Moses Date: Wed, 3 Jun 2026 17:08:38 -0500 Subject: [PATCH 6/9] fixup --- src/utils.jl | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/utils.jl b/src/utils.jl index 7a32ba73ca..e293030cc5 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1204,5 +1204,3 @@ nmantissa(::Type{Float64}) = 52 @inline _unwrap_val(x) = x @inline _unwrap_val(::Val{val}) where {val} = _unwrap_val(val) - - From ed9536ccdfc840515a574e1d3e6cf0849aa164c1 Mon Sep 17 00:00:00 2001 From: Billy Moses Date: Thu, 4 Jun 2026 14:28:35 -0500 Subject: [PATCH 7/9] remove client --- src/compiler/Compiler.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/src/compiler/Compiler.jl b/src/compiler/Compiler.jl index 1601428f58..82a51b2b2e 100644 --- a/src/compiler/Compiler.jl +++ b/src/compiler/Compiler.jl @@ -287,7 +287,6 @@ function compile_mlir!( true; runtime, compile_options.optimize_then_pad, - client, emulate_complex=compile_options.emulate_complex, kwargs..., ) From 4e0a367e0e37a5339c64bb90fa7e575fc41b9620 Mon Sep 17 00:00:00 2001 From: Billy Moses Date: Thu, 4 Jun 2026 17:18:24 -0500 Subject: [PATCH 8/9] fix --- src/TracedUtils.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/TracedUtils.jl b/src/TracedUtils.jl index 8a9bb092dc..bb962803ad 100644 --- a/src/TracedUtils.jl +++ b/src/TracedUtils.jl @@ -645,7 +645,7 @@ function process_linear_args!(linear_args, fnbody, do_transpose, optimize_then_p start_imag = ones(Int, N+1) start_imag[1] = 2 limit_imag = collect(Int, sz) - pushfirst!(limit_real, 2) + pushfirst!(limit_imag, 2) traced_raw = TracedRArray{T2, N+1}((), row_maj_arg, (2, sz...)) From c33a85f86079b48b693f7d478a98f4e9031fddcb Mon Sep 17 00:00:00 2001 From: Billy Moses Date: Mon, 8 Jun 2026 13:14:42 -0500 Subject: [PATCH 9/9] fix --- src/xla/IFRT/Array.jl | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/xla/IFRT/Array.jl b/src/xla/IFRT/Array.jl index 9b7be376b1..7d6368c9d9 100644 --- a/src/xla/IFRT/Array.jl +++ b/src/xla/IFRT/Array.jl @@ -210,8 +210,12 @@ function XLA.to_host(buffer::Array, data, reactant_sharding) if reactant_sharding isa Reactant.Sharding.NoSharding data_buffer = first(single_device_arrays) data_buffer_shape = reverse(size(data_buffer)) - @assert size(data) == data_buffer_shape "Expected data to be of size \ - $(size(data)), got $(data_buffer_shape)" + expected_size = size(data) + if eltype(data) <: Complex && eltype(data_buffer) == real(eltype(data)) + expected_size = (2, expected_size...) + end + @assert expected_size == data_buffer_shape "Expected data to be of size \ + $(expected_size), got $(data_buffer_shape)" GC.@preserve data_buffer begin MLIR.API.ifrt_array_copy_to_host_buffer(data_buffer.buffer, data) end