feat: pass in new options to optimize-comm (#1172)

* feat: pass in new options to optimize-comm

[skip ci]

* feat: more configurable options

* fix: ambiguity

* wip

* rm print

---------

Co-authored-by: William Moses <[email protected]>
Co-authored-by: William S. Moses <[email protected]>

Author: avik-pal

src/CompileOptions.jl

@@ -0,0 +1,22 @@
+# TODO: make the other optimize options into a struct as well
+@kwdef struct OptimizeCommunicationOptions
+    periodic_concat::Int = 0
+    rotate_comm::Int = 0
+    rotate_to_pad_comm::Int = 1
+    wrap_comm::Int = 0
+    extend_comm::Int = 0
+    dus_to_pad_manual_comp_comm::Int = 0
+    dus_to_pad_comm::Int = 1
+    concat_two_operands_comm::Int = 0
+    concat_to_pad_comm::Int = 1
+    extend_to_pad_comm::Int = 1
+    wrap_to_pad_comm::Int = 1
+end
+
+function Base.String(options::OptimizeCommunicationOptions)
+    return (
+        "optimize-communication{" *
+        join(["$(f)=$(getfield(options, f))" for f in fieldnames(typeof(options))], " ") *
+        "}"
+    )
+end

src/Compiler.jl

@@ -1036,21 +1036,25 @@ function raising!(f, is_raising::Bool)
     end
 end
 
-# TODO investigate which options need enable/disable
-# const comm_pass = "optimize-communication{periodic_concat=1 rotate_comm=1 wrap_comm=1 dus_to_pad_manual_comp_comm=1 dus_to_pad_comm=0 concat_two_operands_comm=0 concat_to_pad_comm=0 extend_to_pad_comm=0 wrap_to_pad_comm=0 concat_two_dus_like=1 extend_dus_like=1}"
-const comm_pass = "optimize-communication"
-
-const optimize_comms_passes = (
-    # rotate handler presently broken (and handled okay presently), disabling for now
-    "enzyme-hlo-generate-td{patterns=lower_rotate;concat_to_onedim_dus;concat_to_onedim_dusslice}",
-    "transform-interpreter",
-    "enzyme-hlo-remove-transform",
-    comm_pass,
-    "enzyme-hlo-generate-td{patterns=lower_rotate;lower_wrap;lower_extend}",
-    "transform-interpreter",
-    "enzyme-hlo-remove-transform",
-    comm_pass,
-)
+function get_optimize_comms_passes(options::Bool)
+    options || return String[]
+    return get_optimize_comms_passes(Reactant.OptimizeCommunicationOptions())
+end
+
+function get_optimize_comms_passes(options::Reactant.OptimizeCommunicationOptions)
+    options_str = String(options)
+    res = [
+        "enzyme-hlo-generate-td{patterns=lower_rotate;concat_to_onedim_dus;concat_to_onedim_dusslice;concatreshape_to_onedim_dus}",
+        "transform-interpreter",
+        "enzyme-hlo-remove-transform",
+        options_str,
+        "enzyme-hlo-generate-td{patterns=lower_rotate;lower_wrap;lower_extend}",
+        "transform-interpreter",
+        "enzyme-hlo-remove-transform",
+        options_str,
+    ]
+    return res
+end
 
 function compile_mlir!(
     mod,
@@ -1065,6 +1069,7 @@ function compile_mlir!(
     no_nan::Bool=false,
     transpose_propagate::Symbol=:up,
     reshape_propagate::Symbol=:up,
+    optimize_communications::Bool=true,
     assert_nonallocating::Bool=false,
     backend="gpu",
     raise::Union{Bool,String}=false,
@@ -1459,7 +1464,7 @@ function compile_mlir!(
                 sym_visibility=MLIR.IR.attr(compiled_f, "private"),
             )
             fnbody = MLIR.IR.Block(
-                in_tys_padded, [MLIR.IR.Location() for _ in in_tys_padded]
+		   in_tys_padded, [MLIR.IR.Location(MLIR.API.mlirValueGetLocation(MLIR.IR.argument(MLIR.IR.first_block(MLIR.IR.region(compiled_f, 1)), i))) for i in 1:length(linear_args)]
             )
             push!(MLIR.IR.region(func_with_padding, 1), fnbody)
             MLIR.IR.activate!(fnbody)
@@ -1592,7 +1597,7 @@ function compile_mlir!(
                 join(
                     [
                         "sdy-close-shardings",
-                        optimize_comms_passes...,
+                        get_optimize_comms_passes(optimize_communications)...,
                         "xla-sdy-stablehlo-export-pipeline",
                     ],
                     ",",
@@ -1606,7 +1611,7 @@ function compile_mlir!(
                     [
                         "sdy-propagation-pipeline",
                         "sdy-close-shardings",
-                        optimize_comms_passes...,
+                        get_optimize_comms_passes(optimize_communications)...,
                         "xla-sdy-stablehlo-export-pipeline",
                     ],
                     ",",
@@ -1764,6 +1769,7 @@ macro code_hlo(args...)
         :transpose_propagate => :(:up),
         :reshape_propagate => :(:up),
         :optimize_then_pad => true,
+        :optimize_communications => true,
     )
     compile_expr, (; compiled) = compile_call_expr(
         __module__, compile_mlir, default_options, args...
@@ -1797,6 +1803,7 @@ macro code_mhlo(args...)
         :transpose_propagate => :(:up),
         :reshape_propagate => :(:up),
         :optimize_then_pad => true,
+        :optimize_communications => true,
     )
     compile_expr, (; compiled) = compile_call_expr(
         __module__, compile_xla, default_options, args...
@@ -1830,6 +1837,7 @@ macro code_xla(args...)
         :transpose_propagate => :(:up),
         :reshape_propagate => :(:up),
         :optimize_then_pad => true,
+        :optimize_communications => true,
     )
     compile_expr, (; compiled) = compile_call_expr(
         __module__, compile_xla, default_options, args...
@@ -1863,6 +1871,7 @@ macro compile(args...)
         :transpose_propagate => :(:up),
         :reshape_propagate => :(:up),
         :optimize_then_pad => true,
+        :optimize_communications => true,
     )
     return esc(first(compile_call_expr(__module__, compile, default_options, args...)))
 end
@@ -1885,6 +1894,7 @@ macro jit(args...)
         :transpose_propagate => :(:up),
         :reshape_propagate => :(:up),
         :optimize_then_pad => true,
+        :optimize_communications => true,
     )
     compile_expr, (; compiled, args) = compile_call_expr(
         __module__, compile, default_options, args...

src/Reactant.jl

@@ -179,6 +179,10 @@ const TracedType = Union{TracedRArray,TracedRNumber,MissingTracedValue}
 
 include("ControlFlow.jl")
 include("Tracing.jl")
+
+include("CompileOptions.jl")
+export OptimizeCommunicationOptions
+
 include("Compiler.jl")
 
 include("Overlay.jl")

src/Sharding.jl

@@ -732,7 +732,9 @@ function NamedSharding(sharding::Replicated, ndims::Int)
     return NamedSharding(sharding.mesh, ntuple(Returns(nothing), ndims))
 end
 
-function (sharding::Replicated)(client::XLA.AbstractClient, device, x)
+function (sharding::Replicated)(
+    client::XLA.AbstractClient, device, x::Union{AbstractArray,Number}
+)
     return (NamedSharding(sharding, ndims(x)))(client, device, x)
 end