Start adding sharding

ext/ComradeBaseReactantExt.jl

@@ -6,9 +6,94 @@ using Reactant
 using StaticArrays
 
 import ComradeBase: AbstractSingleDomain, basedim, dims, UnstructuredMap
-using ComradeBase: ReactantEx
+using ComradeBase: ReactantEx, ShardSpec, UnstructuredDomain
 import Reactant: AnyTracedRArray, TracedRArray, unwrapped_eltype
 
+
+# --- Sharding ---------------------------------------------------------------------------------
+# The `ShardSpec` carried on a `ReactantEx` executor is only a *declaration* of how the user wants
+# things laid out. The actual sharding is applied with the public `Reactant.to_rarray` API at the
+# input boundary (`to_sharded` below); ordinary Julia code in the executors / NUFFT then propagates
+# the sharding through tracing. No internal Reactant ops are used.
+
+function ComradeBase.shardmesh(dims::Vararg{Int}; names)
+    return Reactant.Sharding.Mesh(reshape(Reactant.devices(), dims...), names)
+end
+
+# Map each semantic axis name to its array-dimension position for the object being sharded.
+# An unstructured domain/map — and a bare per-visibility vector (a measurement / noise vector with no
+# executor of its own) — is a flat list of points, so the per-point labels :Ti and :Fr both refer to
+# the single data dimension (dim 1); sharding "over Ti/Fr" means sharding that flat axis after the data
+# has been `regroup`ed by the same label. Restricted to `AbstractVector` so a higher-rank payload fails
+# loudly (MethodError) instead of being silently split on dim 1. Rectilinear objects map their named
+# dims (:X, :Y, optionally :Ti, :Fr) to positions in order.
+_axispositions(::Union{UnstructuredDomain, UnstructuredMap, AbstractVector}) = (; Ti = 1, Fr = 1)
+function _axispositions(x::ComradeBase.AbstractRectiGrid)
+    ks = keys(x)
+    return NamedTuple{ks}(ntuple(identity, length(ks)))
+end
+# An IntensityMap is an AbstractArray, so `keys` would give CartesianIndices; take the named dims
+# from its grid instead (the pixel-array dimensions follow the grid's dim order).
+_axispositions(x::IntensityMap) = _axispositions(ComradeBase.axisdims(x))
+
+# Translate a `ShardSpec` + the object's axis layout into a public `Reactant.Sharding.DimsSharding`.
+# `DimsSharding` shards the listed dims and replicates the rest, so one spec works across the
+# different-rank arrays held inside a domain/map. Multiple semantic axes can land on the same array
+# dimension (e.g. :Ti and :Fr both map to the flat dim 1), and a mesh-axis value may itself be a tuple
+# — in either case that dimension is sharded across the tuple of all the mesh axes assigned to it.
+function _dimssharding(spec::ShardSpec, x)
+    pos = _axispositions(x)
+    bydim = Tuple{Int, Symbol}[]
+    for an in keys(spec.axes)
+        haskey(pos, an) || continue
+        v = spec.axes[an]
+        for m in (v isa Tuple ? v : (v,))
+            push!(bydim, (pos[an], m))
+        end
+    end
+    sdims = sort!(unique(first.(bydim)))
+    pspec = map(sdims) do d
+        ax = [m for (dd, m) in bydim if dd == d]
+        length(ax) == 1 ? ax[1] : Tuple(ax)
+    end
+    return Reactant.Sharding.DimsSharding(spec.mesh, Tuple(sdims), Tuple(pspec))
+end
+
+# Resolve the executor's sharding *declaration* into a concrete Reactant sharding for `x`.
+#  - `ShardSpec`: the semantic convenience layer (axis names -> DimsSharding).
+#  - a callable: the escape hatch — full Reactant power; gets `x`, returns any AbstractSharding.
+#  - an AbstractSharding: used as-is.
+_resolve_sharding(spec::ShardSpec, x) = _dimssharding(spec, x)
+_resolve_sharding(sh::Reactant.Sharding.AbstractSharding, x) = sh
+_resolve_sharding(f, x) = f(x)
+
+# Core sharding step: move `x` onto the device for a resolved sharding *declaration* (`nothing` =
+# unsharded). Both public entry points below funnel through here so the one- and two-argument forms,
+# and every object type, can never drift apart.
+_to_sharded(x, ::Nothing) = Reactant.to_rarray(x)
+_to_sharded(x, decl) = Reactant.to_rarray(x; sharding = _resolve_sharding(decl, x))
+
+# An image map's grid axes (`X`, `Y`, `Ti`, `Fr`) are tiny coordinate vectors that the multidomain
+# NUFFT plan builder iterates on the *host*; only the pixel array is large enough to be worth sharding.
+# So shard the value array per the declaration and leave the grid replicated on the host (turning the
+# coordinate arrays into device arrays would make that host iteration a disallowed scalar index). The
+# declaration is resolved against the *image* (for its X/Y/Ti/Fr dim positions) but applied to the bare
+# pixel array; `axisdims(img)` keeps the original host grid.
+_to_sharded(img::IntensityMap, ::Nothing) = IntensityMap(Reactant.to_rarray(baseimage(img)), axisdims(img))
+function _to_sharded(img::IntensityMap, decl)
+    vals = Reactant.to_rarray(baseimage(img); sharding = _resolve_sharding(decl, img))
+    return IntensityMap(vals, axisdims(img))
+end
+
+# One-argument form: shard `x` per its own executor's declaration.
+ComradeBase.to_sharded(x) = _to_sharded(x, ComradeBase.sharding(executor(x)))
+
+# Two-argument form: shard `x` per the declaration carried by `ex` rather than by `x`'s own executor.
+# This is how a likelihood's flat `measurement`/`noise` vectors — which carry no executor — and a
+# visibility domain are placed on the same device blocks from a single `ReactantEx` passed at
+# `prepare_device` time. An `IntensityMap` still keeps its grid on the host via the shared core above.
+ComradeBase.to_sharded(x, ex::ReactantEx) = _to_sharded(x, ComradeBase.sharding(ex))
+
 const RInt = Union{Integer, Reactant.TracedRNumber{<:Integer}}
 
 Base.@propagate_inbounds function ComradeBase.rgetindex(I::Reactant.AnyTracedRArray, i::RInt...)
@@ -31,7 +116,10 @@ Base.@nospecializeinfer function Reactant.make_tracer(
         @nospecialize(runtime),
         kwargs...
     )
-    return Reactant.traced_type(typeof(prev), Val(mode), track_numbers, sharding, runtime)()
+    # Serial/ThreadsEx carry no sharding intent, so they convert to an unsharded ReactantEx.
+    # A domain that was explicitly given `ReactantEx(spec)` keeps its spec via the generic
+    # struct tracer (the Mesh is a leaf type), so no rule is needed for `ReactantEx` itself.
+    return ReactantEx()
 end
 
 Base.@nospecializeinfer function Reactant.traced_type_inner(
@@ -42,7 +130,7 @@ Base.@nospecializeinfer function Reactant.traced_type_inner(
         @nospecialize(ndevices),
         @nospecialize(runtime)
     )
-    return ReactantEx
+    return ReactantEx{Nothing}
 end
 
 
@@ -59,9 +147,6 @@ end
 end
 
 
-# Copied from ComradeBaseKernelAbstractionsExt, these
-# probably will need to be modified still:
-
 function ComradeBase.allocate_map(
         ::Type{<:StructArray{T}},
         g::ComradeBase.AbstractRectiGrid{D, <:ReactantEx}

src/domains/executors.jl

@@ -1,4 +1,5 @@
-export Serial, ThreadsEx
+export Serial, ThreadsEx, ReactantEx, ShardSpec, shard_image, shard_frequency, shard_time, shardmesh,
+    to_sharded
 
 """
     Serial()
@@ -20,12 +21,146 @@ ThreadsEx() = ThreadsEx(:dynamic)
 ThreadsEx(s) = ThreadsEx{s}()
 
 """
-    ReactantEx
+    ShardSpec(mesh, axes::NamedTuple)
+
+A backend-agnostic description of how to shard image or visibility evaluation across a device
+`mesh`. The `mesh` is treated opaquely by ComradeBase (at use time it is a `Reactant.Sharding.Mesh`)
+so that the core package carries no Reactant dependency. `axes` maps *semantic* ComradeBase axis
+names — the same names Comrade uses, `:X`, `:Y`, `:Ti`, `:Fr` — to mesh-axis names (`Symbol`s):
+
+  - Image (`RectiGrid`/`IntensityMap`): `:X`, `:Y`, `:Ti`, `:Fr` are genuine array dimensions and map
+    straight onto a sharding of those dimensions.
+  - Visibility (`UnstructuredDomain`): the data is a flat vector, so `:Ti`/`:Fr` are per-point labels.
+    Sharding over them means [`regroup`](@ref)ing the domain by that label so each `(Ti, Fr)` block is
+    contiguous, then sharding the flat data axis. For ALMA-like data (constant baselines per stamp) the
+    even split lands exactly on block boundaries; for ragged VLBI data a few blocks straddle device
+    boundaries (a bounded cost — a dense pad+mask layout can remove it later).
+
+To co-locate each `(Ti, Fr)` block's image slice with its visibilities, declare the **same** `:Ti`/`:Fr`
+mesh axes on both the image and the visibility domain.
+
+Prefer the convenience constructors on [`ReactantEx`](@ref) instead of building this directly, e.g.
+`ReactantEx(mesh; X = :dx, Y = :dy)` or `ReactantEx(mesh; Fr = :dev)`.
+"""
+struct ShardSpec{M, A <: NamedTuple}
+    mesh::M
+    axes::A
+end
+
+"""
+    ReactantEx()
+    ReactantEx(sharding::ShardSpec)
+    ReactantEx(mesh; axisname = meshaxis...)
+    ReactantEx(f)                            # f(x) -> a Reactant.Sharding.AbstractSharding
+
 Uses Reactant.jl for execution when computing the intensitymap or visibilitymap. Note that specifying
 this should be unnecessary as ComradeBase will automatically switch to this backend when
 it detects that it is inside a Reactant tracing context.
+
+When a [`ShardSpec`](@ref) is attached (typically via the keyword constructor) the image or
+visibility arrays are sharded across the supplied device mesh during evaluation. The keyword
+constructor maps semantic axis names to mesh-axis names, e.g.
+
+```julia
+mesh = ComradeBase.shardmesh(2, 2; names = (:dx, :dy))  # requires Reactant
+g = imagepixels(μas2rad(100.0), μas2rad(100.0), 256, 256; executor = ReactantEx(mesh; X = :dx, Y = :dy))
+```
+
+A mesh-axis value may itself be a tuple to shard one image/data dimension across several mesh axes,
+e.g. `ReactantEx(mesh; X = (:dx1, :dx2))`.
+
+**Escape hatch.** The semantic [`ShardSpec`](@ref) only covers the common "shard dimension *D* across
+mesh axis *M*" case. For anything Reactant supports that it does not express — `NamedSharding` with
+sub-axes, priorities, open/closed dims, partial replication, layouts that depend on the array — pass a
+**function** instead. [`to_sharded`](@ref) calls it with the object being sharded and uses whatever
+`Reactant.Sharding.AbstractSharding` it returns:
+
+```julia
+executor = ReactantEx(x -> Reactant.Sharding.NamedSharding(mesh, (:dx, nothing, :dev)))
+```
+
+See also [`shard_image`](@ref), [`shard_frequency`](@ref), and [`regroup`](@ref).
+"""
+struct ReactantEx{S}
+    sharding::S
+    # Explicit inner constructor so Julia does not auto-generate the `ReactantEx(::Any)` outer
+    # constructor, which would collide with the keyword constructor below.
+    ReactantEx{S}(sharding) where {S} = new{S}(sharding)
+end
+ReactantEx() = ReactantEx{Nothing}(nothing)
+ReactantEx(s::Union{Nothing, ShardSpec, Function}) = ReactantEx{typeof(s)}(s)
+function ReactantEx(mesh; axes...)
+    spec = ShardSpec(mesh, NamedTuple(axes))
+    return ReactantEx{typeof(spec)}(spec)
+end
+
+"""
+    sharding(ex::ReactantEx)
+
+Returns the sharding declaration attached to an executor, or `nothing` when the executor carries no
+sharding (the default for [`Serial`](@ref)/[`ThreadsEx`](@ref) and for an unsharded `ReactantEx`).
+"""
+sharding(@nospecialize(::Any)) = nothing
+sharding(ex::ReactantEx) = getfield(ex, :sharding)
+
+"""
+    shard_image(mesh; x = :dx, y = :dy)
+
+Convenience helper returning a [`ReactantEx`](@ref) that shards the image `:X`/`:Y` axes across the
+`x`/`y` axes of `mesh`.
+"""
+shard_image(mesh; x = :dx, y = :dy) = ReactantEx(mesh; X = x, Y = y)
+
+"""
+    shard_frequency(mesh; axis = :dev)
+
+Convenience helper returning a [`ReactantEx`](@ref) that shards the image `:Fr` axis across the
+`axis` axis of `mesh`. Combine with [`regroup`](@ref)`(visdomain, :Fr)` so each frequency block
+maps to a single device.
+"""
+shard_frequency(mesh; axis = :dev) = ReactantEx(mesh; Fr = axis)
+
+"""
+    shard_time(mesh; axis = :dev)
+
+Convenience helper returning a [`ReactantEx`](@ref) that shards the image `:Ti` axis across the
+`axis` axis of `mesh`. Combine with [`regroup`](@ref)`(visdomain, :Ti)` so each time block maps to a
+single device.
+"""
+shard_time(mesh; axis = :dev) = ReactantEx(mesh; Ti = axis)
+
+"""
+    shardmesh(dims...; names)
+
+Convenience constructor for a device mesh over the available Reactant devices, laid out with the
+given shape `dims` and mesh-axis `names`. Equivalent to
+`Reactant.Sharding.Mesh(reshape(Reactant.devices(), dims...), names)`. Requires Reactant to be loaded.
+"""
+function shardmesh end
+
+"""
+    to_sharded(x)
+    to_sharded(x, ex::ReactantEx)
+
+Move a domain or map `x` onto the device mesh declared by its `ReactantEx` executor's [`ShardSpec`](@ref),
+returning a sharded copy via the public `Reactant.to_rarray` API. The sharding rides on the input
+arrays; ordinary evaluation code then propagates it. If the executor carries no `ShardSpec`, this is
+just `Reactant.to_rarray(x)` (unsharded). Requires Reactant to be loaded.
+
+The two-argument form shards `x` using the declaration carried by `ex` rather than by `x`'s own
+executor. Use it to place a visibility domain and a likelihood's flat `measurement`/`noise` vectors —
+which carry no executor of their own — onto the same blocks from a single `ReactantEx`. A bare vector
+is treated as a flat list of points, so `:Ti`/`:Fr` select its first dimension.
+
+```julia
+mesh = shardmesh(length(Reactant.devices()); names = (:dev,))
+dvis = domain(arrayconfig; executor = ReactantEx(mesh; Fr = :dev))
+rdvis, perm = regroup(dvis, :Fr)   # sort by frequency so each Fr block is contiguous
+dvis_sharded = to_sharded(rdvis)   # flat coordinate arrays now sharded across the mesh
+# apply `perm` to the data/noise vectors so they stay aligned with the model visibilities
+```
 """
-struct ReactantEx end
+function to_sharded end
 
 
 #TODO can this be made nicer?

src/domains/unstructured/domain.jl

@@ -1,4 +1,4 @@
-export UnstructuredDomain
+export UnstructuredDomain, regroup
 
 const DataNames = Union{
     <:NamedTuple{(:X, :Y, :T, :F)}, <:NamedTuple{(:X, :Y, :F, :T)},
@@ -76,6 +76,77 @@ function Base.getindex(domain::UnstructuredDomain; Ti = nothing, Fr = nothing)
     return UnstructuredDomain(points[indices], executor(domain), header(domain))
 end
 
+"""
+    regroup(domain::UnstructuredDomain, axes::Symbol...)
+
+Return `(rdomain, perm)` where `rdomain` is `domain` reordered by a stable **lexicographic** sort over
+the points' `axes` properties so that all points sharing the same `axes` values are contiguous, and
+`perm` is the permutation such that `domainpoints(rdomain) == domainpoints(domain)[perm]`.
+
+With a single axis (`regroup(d, :Fr)`) this groups by frequency; with several
+(`regroup(d, :Fr, :Ti)`) the first axis is the major key, so the data nests as `Fr`-blocks each split
+into `Ti`-blocks — matching a sharding declared as `ReactantEx(mesh; Fr = :devf, Ti = :devt)` (the
+declaration's axis order sets the same major/minor order). This is the companion to sharding the
+corresponding image dimensions across a device mesh (see [`shard_frequency`](@ref)/[`shard_time`](@ref)):
+contiguous-per-group ordering lets the per-block visibilities lay out cleanly across devices. Apply
+`perm` to any data/noise vectors that must stay aligned with the model visibilities, and use
+`invperm(perm)` to map results back to the original order.
+
+!!! note
+    Even per-device placement of whole groups requires the groups to be (close to) equal sized. With
+    ragged group sizes the grouping is still contiguous but a uniform mesh split will not fall exactly
+    on group boundaries.
+"""
+# Shared core for every `regroup` method: build one stable **lexicographic** permutation from the
+# per-point key columns (first column is the major key) and return the reordered domain plus `perm`.
+# All flavours of grouping (by raw value, or by grid-plane position) differ only in how they build
+# `keycols`, so the ordering lives in exactly one place.
+function _regroup(domain::UnstructuredDomain, keycols)
+    points = domainpoints(domain)
+    keyvecs = collect(zip(keycols...))   # vector of tuples, compared lexicographically
+    perm = sortperm(keyvecs; alg = Base.Sort.DEFAULT_STABLE)
+    rdomain = UnstructuredDomain(points[perm], executor(domain), header(domain))
+    return rdomain, perm
+end
+
+function regroup(domain::UnstructuredDomain, axes::Symbol...)
+    isempty(axes) && throw(ArgumentError("regroup requires at least one axis"))
+    points = domainpoints(domain)
+    return _regroup(domain, map(a -> getproperty(points, a), axes))
+end
+
+# Rank a coordinate `v` by its position within the grid axis values `order`, matched with `isapprox`
+# so a data frequency/time that differs from the grid's by floating-point round-off still lands on the
+# right plane. A coordinate absent from the axis sorts to the end (`lastindex + 1`).
+_gridrank(order, v) = something(findfirst(x -> isapprox(x, v), order), lastindex(order) + 1)
+
+"""
+    regroup(domain::UnstructuredDomain, grid::AbstractRectiGrid)
+
+Reorder `domain` into the plane order the multidomain Fourier transform enumerates the image planes
+of `grid` — `DimPoints(dims(grid)[3:end])`, i.e. column-major over the grid's *stored* non-spatial
+(`Ti`/`Fr`) coordinate order. Unlike the value-sorting [`regroup`](@ref)`(domain, axes...)`, points are
+ranked by their *position within* each grid axis (matched with `isapprox`), so a descending or
+non-monotonic grid still lays each block out to co-locate with its image plane — the precondition for
+sharding the blocks across a device mesh. Points whose coordinate is absent from a grid axis sort to
+the end. Returns `(rdomain, perm)`.
+
+This is the single source of the data-side plane ordering; it derives its axes from the same
+`dims(grid)[3:end]` the NUFFT planner iterates, so the two cannot silently drift apart.
+"""
+function regroup(domain::UnstructuredDomain, grid::AbstractRectiGrid)
+    points = domainpoints(domain)
+    # non-spatial image axes in stored (column-major) order, keeping only those the data carries
+    gaxes = filter(a -> hasproperty(points, a), map(name, dims(grid)[3:end]))
+    isempty(gaxes) && return domain, Base.OneTo(length(points))
+    ranks = map(gaxes) do a
+        order = getproperty(grid, a)
+        map(v -> _gridrank(order, v), getproperty(points, a))
+    end
+    # `DimPoints` varies the first non-spatial dim fastest, so the last grid dim is the major sort key.
+    return _regroup(domain, reverse(ranks))
+end
+
 function Base.summary(io::IO, g::UnstructuredDomain)
     n = propertynames(domainpoints(g))
     printstyled(io, "│ "; color = :light_black)