Emulated complex

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,

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")

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)

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

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,