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 dbf2b6703e..bb962803ad 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,7 @@ function make_mlir_fn( resargprefix::Symbol=:resargs, num_replicas=1, optimize_then_pad::Bool=true, + emulate_complex::Bool=false, ) if sizeof(typeof(f)) != 0 || f isa Base.BroadcastFunction mlir_fn_res = make_mlir_fn( @@ -338,6 +340,7 @@ function make_mlir_fn( resprefix, resargprefix, num_replicas, + emulate_complex, ) mlir_fn_res.fnwrapped = true return mlir_fn_res @@ -354,6 +357,7 @@ function make_mlir_fn( do_transpose, input_shardings, verify_arg_names, + emulate_complex, ) Ops.activate_constant_context!(fnbody) @@ -364,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) + 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...) @@ -417,6 +421,7 @@ function make_mlir_fn( N, concretein, toscalar, + emulate_complex, ) return CompiledMlirFnResult( @@ -457,6 +462,7 @@ function prepare_mlir_fn_args( do_transpose, input_shardings, verify_arg_names, + emulate_complex, ) N = length(args) traced_args = Vector{Any}(undef, N) @@ -501,19 +507,37 @@ function prepare_mlir_fn_args( 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 && emulate_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 +616,7 @@ 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, 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 @@ -608,7 +632,34 @@ 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 && emulate_complex + T2 = real(T) + sz = size(arg) + N = length(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_imag, 2) + + 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 @@ -642,6 +693,7 @@ function finalize_mlir_fn( N, concretein, toscalar, + emulate_complex, ) # check which arguments have been mutated mutated_args = Int[] @@ -856,12 +908,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 && emulate_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) - if do_transpose - col_maj = transpose_val(col_maj) - out_ty = transpose_ty(out_ty) + 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) + end end end @@ -927,8 +1006,14 @@ function finalize_mlir_fn( sharding.mesh.axis_names, size(sharding.mesh), )] + T = Reactant.unwrapped_eltype(arg) + sz = if T <: Complex && emulate_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 +1073,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 && emulate_complex + sharding = Reactant.adapt_sharding_for_emulated_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/compiler/Compiler.jl b/src/compiler/Compiler.jl index f275877eba..82a51b2b2e 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,6 +119,10 @@ function compile_mlir( mod, f; mlir_fn_res.num_partitions, mlir_fn_res.num_replicas ) + if compile_options.emulate_complex + 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,7 @@ function compile_mlir!( true; runtime, compile_options.optimize_then_pad, + emulate_complex=compile_options.emulate_complex, kwargs..., ) finally @@ -1325,7 +1330,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, @@ -1404,7 +1409,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 @@ -1416,11 +1421,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/src/xla/Client.jl b/src/xla/Client.jl index ccf715c1ba..f51611fe87 100644 --- a/src/xla/Client.jl +++ b/src/xla/Client.jl @@ -14,3 +14,15 @@ function get_addressable_device end function platform_name end default_device(client::AbstractClient) = first(addressable_devices(client)) + +function supports_complex(client::AbstractClient) + if Reactant.precompiling() + return true + end + 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..7d6368c9d9 100644 --- a/src/xla/IFRT/Array.jl +++ b/src/xla/IFRT/Array.jl @@ -22,13 +22,22 @@ function Array( 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), @@ -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), @@ -119,6 +135,14 @@ end function Array( 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 + 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)) @@ -186,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 diff --git a/src/xla/PJRT/Buffer.jl b/src/xla/PJRT/Buffer.jl index e4e7e65849..f8dc47efd0 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::Array{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 diff --git a/test/core/complex.jl b/test/core/complex.jl index 041cacfb58..ce222b75bd 100644 --- a/test/core/complex.jl +++ b/test/core/complex.jl @@ -119,3 +119,13 @@ end @test @jit(isreal.(x_concrete)) == isreal.(x) end end + +@testset "forced complex emulation" begin + @testset "emulated complex addition" begin + 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, 3) + ones(ComplexF32, 3) + end +end +