AwsNeuronLNCShardingConstraint and the SPMD↔LNC Coupling

All addresses on this page apply to neuronx-cc 2.24.5133.0+58f8de22, the hlo2penguin front-end binary, cp310 build, unless tagged otherwise. The two-VA-frame convention holds: .rodata VA = file offset + 0x200000, .text VA = file offset + 0x201000 — VA is not the raw file offset. Other versions will differ. Evidence is static-binary-derived (D-AB08).

Abstract

neuronx-cc runs a stock upstream-XLA SPMD partitioner inside hlo2penguin. The driver, sharding propagation, and the compute handlers are all verbatim XLA — there is no Neuron subclass of SpmdPartitioner, no Neuron ShardingPropagation, no Neuron compute-handler set. Neuron customizes SPMD at exactly two seams, and this page documents the more interesting one: the custom-call targets that ride the stock partitioner. The companion seam — the front-end --distribution-strategy / --spmd / --lnc options that seed the mesh — is owned by distribution-strategy-seeding; this page only touches it where the LNC count couples the two ends.

The namesake op is AwsNeuronLNCShardingConstraint: a side-effect-free HLO custom-call that pins an xla::HloSharding onto an HLO value to bind it to the physical Logical NeuronCore (LNC) device mesh. It is Neuron's analogue of XLA's stock Sharding custom-call / sdy::ShardingConstraintOp, but keyed to the physical LNC topology instead of an abstract logical mesh. The critical structural fact — verified by exhaustive negative xref — is that it has no dedicated CustomCallPartitioner subclass. It flows through the stock target-string-keyed CustomCallPartitioner registry in xla/service/custom_call_sharding_helper.cc, the same registry every other custom-call uses. This is "Neuron data riding stock machinery," not a Neuron-subclassed partitioner.

The page covers, in order: the LNC-constraint op and its penguin-IR printer (§1); the stock registry and the exact dispatch path that reshards an operand to the pinned sharding (§2); the sibling AwsNeuronTransferWithStaticRing collective-transfer custom-call and its parameter-load / rematerialization coupling (§3); and the front↔back contract — how the front-end logical_nc_config and the backend lnc_size are one quantity N, and how the LNC constraint makes the SPMD-tile↔physical-LNC-core mapping deterministic so the backend lnc_splitter produces SPMD-identical replicas (§4).

For reimplementation, the contract is:

• The op shape of AwsNeuronLNCShardingConstraint: an identity-on-data custom-call carrying target_name= + a sharding= HloSharding payload; emitted/validated/printed/cost-modeled only in the front-end.
• The stock CustomCallPartitioner registry mechanism: a process-global flat_hash_map<target, CCP> populated once via absl::CallOnce from ShardingPropagation::Run, consulted by target string at every SPMD touch-point.
• The dispatch-and-reshard path: how an unregistered (or sharding-constraint) target reaches PartitionedHlo::Reshard and forces the operand to the pinned sharding using stock collective machinery.
• The front↔back logical_nc_config == lnc_size == N contract and how the LNC constraint binds an SPMD tile to a specific LNC core.

NOTE — throughout, stock-XLA marks code that is verbatim upstream XLA (xla::, xla::spmd::, xla::sdy::), and Neuron marks Neuron-authored code (xla::hilo::, xla::partition::, AwsNeuron* strings, PenguinizeFunctions, UnpackNestedAWSNTWSR, the mlir::*PythonPrinter family). The whole point of the page is that the mechanism is stock and only the data is Neuron.

1. The LNC Sharding-Constraint Op

Purpose

AwsNeuronLNCShardingConstraint is one of the real AwsNeuron* custom-call targets catalogued in the front-end's verifier whitelist. Semantically it is a sharding-constraint op: a no-op-on-data, side-effect-free custom-call whose output shape equals its input shape, existing solely to carry a required HloSharding into the SPMD partitioner. The sharding it carries is a tile assignment over the LNC device mesh — the N Logical-NeuronCore SPMD devices the module is being partitioned across. It is the one op that ties an abstract SPMD HloSharding to the physical LNC topology.

It is a purely front-end construct. The string AwsNeuronLNCShardingConstraint (.rodata 0x360f50, len 30) appears only in hlo2penguin — verified absent (jq count 0) from the hlo-opt mid-level optimizer's string table, where the sibling AwsNeuronTransferWithStaticRing is present (count 1). The constraint is therefore consumed inside the HLO→penguin import path — which is where the stock SPMD partitioner runs — and stripped once partitioning has applied it, exactly as a sharding-constraint op should be. (CONFIRMED.)

QUIRK — the op is "front-end only" not because Neuron hides it, but because it is a transient annotation. SPMD propagation reads it, reshards the operand to its sharding, and the op evaporates before the optimized HLO reaches hlo-opt. A reimplementer who looks for it in the mid-level IR will not find it; it lives only between import and the end of the front-end SPMD pass.

Evidence It Has No Dedicated Pass

The string's complete referencing-function set (from the strings sidecar referenced_by_functions) is exactly four functions, and none is a rewrite/legalize/partition pass:

All four are import / validate / print / cost functions. There is no xla::Legalize* / Partition* / Lower* pass that special-cases this target string in either binary. The op is given no CustomCallPartitioner subclass. This is a confirmed negative result (exhaustive xref enumeration over both binaries) and is what licenses the whole "rides stock machinery" thesis. (CONFIRMED.)

The Penguin-IR Printer

The op's only dedicated emitter is the penguin-IR printer:

MhloToPythonPrinter::print<mhlo::CustomCallOp>   @0x20d1f10 (2762 B)   ── dispatch on call-target name
  └─ std::function<void(Operation*)> thunk                            ── one per AwsNeuron* target
       └─ MhloToPythonPrinter::printLNCShardingConstraint  @0x20e37e0 ── the LNC-constraint emitter
  (StableHLO twin: StableHLOToPythonPrinter::print<...> ─→ printLNCShardingConstraint @0x217fb00)

QUIRK — dispatch from print<CustomCallOp> to printLNCShardingConstraint is not a direct call. The printer template builds a table of std::function<void(mlir::Operation*)> thunks keyed by target name; the only caller of printLNCShardingConstraint in the binary is the std::_Function_handler<...print<CustomCallOp>...>::_M_invoke thunk wrapper, not the template body. A reimplementer chasing a direct callee edge from the template will miss it — the routing goes through a function-pointer table. (CONFIRMED — caller of @0x20e37e0 is the _M_invoke thunk.)

printLNCShardingConstraint @0x20e37e0 emits the penguin op form. Its three referenced string literals and five callees pin the exact emission:

// MhloToPythonPrinter::printLNCShardingConstraint(mlir::Operation* op)   // @0x20e37e0, 1281 B
//
// Emits:   <dsts> = <ctx>.LNCShardingConstraintOp( <srcs>,
//                                                  target_name=<getCallTargetName()>,
//                                                  sharding=<the op's HloSharding> )
void printLNCShardingConstraint(Operation* op):
    dsts = printDsts(op)                       // @callee printDsts
    srcs = printSrcs(op, /*bitset<20>*/)       // @callee printSrcs
    emit(dsts, " = ", ctx, ".LNCShardingConstraintOp(")     // str ".LNCShardingConstraintOp(" @0x26a802
    emit(srcs)
    emit(", target_name=", op.getCallTargetName())          // str "target_name="  @0x27aecd
                                                            //  callee CustomCallOp::getCallTargetName()
    emit(", sharding=",  sharding_attr_of(op))              // str "sharding="     @0x23a742
    // backend_config payload also read:
    bc = op.getBackendConfig()                              // callee CustomCallOp::getBackendConfig()
    printMeta(op)                                            // @callee printMeta
    emit(")")

The three literals (.LNCShardingConstraintOp(, target_name=, sharding=) and the five callees (getCallTargetName, getBackendConfig, printSrcs, printDsts, printMeta) are all CONFIRMED in the function's strings_referenced and callees. The sharding= value is the SPMD HloSharding attached to the custom-call op — the constraint payload, i.e. the tile-over-LNC assignment. The output shape equals the input shape; the op is identity on data and exists only to carry the sharding requirement into SPMD. The sharding= payload being the constraint's HloSharding is STRONG (the attribute is the op's defining content, which is exactly what a sharding-constraint op stores).

Validation and Cost

Two of the four referencing functions consume the op outside the printer:

• CustomCallOpChecker::CheckMisc @0x203e890 holds a std::_Hashtable<string> whitelist of legal AwsNeuron* target names (including AwsNeuronLNCShardingConstraint) and rejects unknown targets via hilo::formatErrorMessage(ErrorCode, ...); it logs the "[ERROR] [" prefix on failure. The front-end verifier therefore accepts the LNC constraint as a known op. (CONFIRMED.)
• NeuronHloCostAnalysis::HandleCustomCall @0x21b8b10 routes the same target set through the Neuron cost model. Being an identity annotation, AwsNeuronLNCShardingConstraint contributes ~zero compute — it is a metadata op — unlike AwsNeuronCollectiveMatmul, which carries dot dimensions (lhs/rhs contracting + batch dims) read from backend_config. (CONFIRMED string xref; STRONG zero-cost rationale.)

2. The Stock CustomCallPartitioner Seam

Purpose

Because the LNC constraint has no dedicated partitioner subclass, it must ride the stock CustomCallPartitioner registry. This section is the heart of the page: it shows the exact stock machinery that a Neuron custom-call target flows through, and how a sharding-constraint custom-call ends up reshard­ing its operand to the pinned sharding. Everything in this section is upstream XLA — the source-file string xla/service/custom_call_sharding_helper.cc (@0x3cf7d0) confirms it. (stock-XLA.)

The Registry

// xla::(anon)::GetPartitioners()                          // @0x2d24f30, 143 B   (stock-XLA)
// Process-global, lazily-constructed, mutex-guarded:
absl::flat_hash_map<std::string, std::unique_ptr<CustomCallPartitioner>>&
GetPartitioners():
    static map  guarded by static absl::Mutex
    return map

// xla::RegisterCustomCallPartitioner(string_view target,                  // @0x2d25830, 440 B
//                                    unique_ptr<CustomCallPartitioner> p)
void RegisterCustomCallPartitioner(target, p):
    lock GetPartitioners() mutex
    if !GetPartitioners().insert({target, move(p)}):
        LOG "Failed to register custom call partitioner for "    // str present in fn
        //   source: "xla/service/custom_call_sharding_helper.cc"

// xla::GetCustomCallPartitioner(string const& target)         // @0x2d25140, 353 B
const CustomCallPartitioner* GetCustomCallPartitioner(target):
    lock GetPartitioners() mutex
    it = GetPartitioners().find(target)        // hash the string, mutex-locked find
    return it == end ? nullptr : it->second.get()

The registry is keyed by custom_call_target. It is populated once via absl::CallOnce — its sole caller is the .constprop CallOnceImpl<...ShardingPropagation::Run...UlvE_> lambda, i.e. registration is run lazily from ShardingPropagation::Run. (CONFIRMED: the only caller of RegisterCustomCallPartitioner @0x2d25830 is that CallOnceImpl thunk.)

Who Consults the Registry

GetCustomCallPartitioner @0x2d25140 has exactly seven callers — the dispatch key (the target string) is consulted at every SPMD touch-point:

All seven are stock XLA. (CONFIRMED — the full caller list of @0x2d25140.)

Dispatch and Reshard — the Algorithm

// SpmdPartitioningVisitor::HandleCustomCall(HloInstruction* hlo)   // @0x2c15e30, 4748 B  (stock-XLA)
Status HandleCustomCall(hlo):
    target = hlo->custom_call_target()
    ccp = GetCustomCallPartitioner(target)              // @0x2d25140 — registry dispatch (CONFIRMED edge)
    if ccp != nullptr:
        return ccp->Partition(this, hlo)                // bespoke partitioner, if one was registered

    // No registered partitioner -> fall to the stock pass-through / reshard paths:
    if is_elementwise_like(hlo):
        return HandleElementwise(hlo)                   // @callee HandleElementwise
    if forced_single_device(hlo):
        return HandleSingleDevice(hlo)                  // @callee HandleSingleDevice

    // A sharding-constraint custom-call resolves to: reshard the operand to the
    // requested HloSharding (the "sharding=" payload from §1):
    operand   = GetPartitionedHlo(hlo->operand(0))
    requested = hlo->sharding()                         // the pinned LNC tiling
    return operand.Reshard(requested, /*literal=*/nullopt)   // @0x2cad850 (CONFIRMED edge)

HandleCustomCall @0x2c15e30 confirmedly calls both GetCustomCallPartitioner and PartitionedHlo::Reshard @0x2cad850, plus HandleElementwise / HandleSingleDevice as the pass-through paths. The base CustomCallPartitioner::Partition @0x2d250a0 is an unimplemented stub (it logs "Implement sharding for %s"), and the base CanSideEffectingHaveReplicatedSharding @0x2bfc590 is a 3-byte return false (xor eax,eax; ret). So an unregistered target falls to the default reshard/elementwise/replicate handling — which is precisely the reshard a sharding-constraint forces. (CONFIRMED stubs and edges.)

Reshard @0x2cad850 is the stock collective machinery: all-to-all / collective-permute / all-gather / dynamic-slice, selected per the stock GetReshardAllToAllSourceTargetDims / CanReshardWithCollectivePermute predicates (all present in hlo2penguin). The same function also handles the stock SPMDFullToShardShape / SPMDShardToFullShape custom-calls. (CONFIRMED.)

                       Neuron data  ──────────────────────────────────►  stock-XLA machinery
  AwsNeuronLNCShardingConstraint
     (target string + sharding=)
            │
            ▼
   HandleCustomCall ── GetCustomCallPartitioner(target) ── registry miss (no Neuron subclass)
            │                                                        │
            │                                                        ▼
            └──────────────────────────────────────────────►  PartitionedHlo::Reshard(sharding)
                                                                     │
                                          all-to-all / collective-permute / all-gather / dynamic-slice

GOTCHA — the honest gap. No symbol literally named e.g. LNCShardingConstraintPartitioner is registered in GetPartitioners(); the registrant set is populated indirectly (CallOnce from ShardingPropagation::Run) and is not enumerable from the skipped-decompile import. The dispatch mechanism (target → GetCustomCallPartitioner → Partition / Propagate* → Reshard) is CONFIRMED stock XLA. That AwsNeuronLNCShardingConstraint flows through this mechanism rather than a bespoke pass is STRONG — proven by the exhaustive negative xref of §1 (no Neuron pass references the string), not by observing a registry entry for it. The reshard-to-pinned-sharding behavior is the stock sharding-constraint behavior; whether a thin Neuron CCP is registered for it or it falls to the default reshard path is not separable from the skipped import, but the result (operand resharded to the pinned LNC tiling) is identical either way.

NOTE — the LNC constraint is side-effect-free (a pure annotation), so it is never gated by the side-effecting-sharding RET_CHECK path that fires only for side-effecting custom-calls. This is why it can be freely resharded and then stripped without a replication-legality check. (STRONG.)

3. AwsNeuronTransferWithStaticRing

Purpose

AwsNeuronTransferWithStaticRing (.rodata 0x2dfd28, len 31) is the Neuron collective-transfer custom-call: it moves a tensor between LNC cores along a static ring — a fixed, compile-time-determined ring of Logical NeuronCores — rather than a dynamically-routed collective. "Static" means the cross-core transfer path and order are baked at compile time (fixed by the LNC mesh), so the backend can emit a deterministic DMA schedule instead of a runtime-negotiated route. Unlike the LNC constraint, this op is present in both hlo2penguin and hlo-opt (count 1 in the optimizer's string table) — it survives into the mid-level optimizer because it is a real data-movement op, not a transient annotation. (CONFIRMED.)

Who Touches It

The string's full referencing set (seven functions) splits into validate / cost / cleanup / partition-boundary / penguinize:

Flattening Nested Transfers

// xla::UnpackNestedAWSNTWSR::Run(HloModule* m, flat_hash_set<string_view> exec)   // @0x20537f0
//   1897 B, 79 basic blocks, 68 callees.  ("AWSNTWSR" = AWS-Neuron-Transfer-With-Static-Ring)
Status UnpackNestedAWSNTWSR::Run(m, exec):
    RET_CHECK(m->entry_computation() != nullptr)                  // "nullptr != entry_computation_"
    for comp in m->computations():
        for inst in comp->MakeInstructionPostOrder():            // @callee MakeInstructionPostOrder
            if inst->custom_call_target() != "AwsNeuronTransferWithStaticRing":
                continue
            source = inst->mutable_operand(0)                    // @callee mutable_operand
            if source->custom_call_target() == "AwsNeuronTransferWithStaticRing":   // nested!
                NeuronLogger << "found illegally nested AwsNeuronTransferWithStaticRing, "
                                "replacing <tail> with <source>"  // str @0x354968, len 66 (CONFIRMED)
                comp->ReplaceInstruction(/*tail=*/inst, source)   // @callee HloComputation::ReplaceInstruction
    return OK

The pass collapses transfer(transfer(x)) → transfer(x): two consecutive static-ring transfers on the same value are illegal (a value already in transit on the ring must not be re-wrapped), so the inner is removed and the outer's operand is rebound to the source, keeping the static-ring schedule single-hop. The callee set (MakeInstructionPostOrder, mutable_operand, ReplaceInstruction) and the NeuronLogger diagnostic chain (getInstance / setSourceFile / setSourceLine / shouldLogToFile / shouldLogToConsole) are all CONFIRMED in the function. (CONFIRMED.)

Static-Ring Transfer as a Parameter-Load

Both isCustomCallLoadingParam predicates are tiny (@0x1f3d6b0, 221 B; @0x1f16550, 237 B) and return true iff custom_call_target() == "AwsNeuronTransferWithStaticRing". Their callers locate them in the Neuron graph-partition layer (xla::partition — the module/LNC splitter, distinct from xla::spmd):

partition::isCustomCallLoadingParam @0x1f3d6b0
   ├─ partition::SplitFinder::findInitSplits    ── seeds the init-split set (input boundary)
   └─ partition::replicateTrivial               ── trivial replicate
Rematter::isCustomCallLoadingParam @0x1f16550
   └─ partition::Rematter::trivialRemats        ── trivial rematerialization

A static-ring transfer that brings a parameter onto a core is treated like a parameter load: (i) SplitFinder::findInitSplits seeds it as an input boundary of a partition — it is where data enters the per-core subgraph — and (ii) Rematter::trivialRemats marks it trivially rematerializable, so the loaded value can be re-fetched via the ring rather than spilled/recomputed, and replicateTrivial may replicate it cheaply. (CONFIRMED callers; STRONG "why" from caller names + the LNC split model.)

Penguinization

PenguinizeFunctions (MHLO) @0x2087ee0 and StableHLOPenguinizeFunctions @0x2127230 each carry a per-CustomCallOp lambda that memcmp-matches the static-ring target name and lowers the op to the penguin transfer/collective representation. The exact penguin op kind was not decompilable (the lambda body is a skipped-decompile export); that it is a transfer/collective op rather than a compute op is STRONG (the cost model accounts its bytes-moved as a transfer cost, and the target-name match in the penguinize lambda is CONFIRMED). (CONFIRMED match; STRONG op-kind.)

4. The Front↔Back LNC Contract

The LNC Count Is the Device-Mesh Cardinality

A single number — the LNC count, logical_nc_config = NeuronCores-per-Logical-NeuronCore — is what the SPMD mesh, the front-end sharding, and the backend split all key off:

• Front-end (CompileCommand, owned by distribution-strategy-seeding): self.logical_nc_config (--lnc / --logical-nc-config), default 2 on Trn2 (codename sunda), forced 1 on other archs (constraint string args.arch != "sunda" or args.logical_nc_config == 1). It gates enable_internal_spmd_opt and is surfaced to hlo2penguin as the HloPassOptions core-count knob. (CONFIRMED; see the seeding page.)
• Backend (lnc_splitter, walrus/lnc-splitter): lnc_size = PassOptions+0x1A4 (dword index 420), cached to LncSplitter+0x60, default 1 — the per-module replication factor. LncSplitter::run @0x16d5ca0 (pipeline order 8) clones the one symbolic BIR module lnc_size times and concretizes per-core shard-ids; the cross-core collective lowering computes getRemoteCores={0..lnc-1}\{self} and getAllCores={0..lnc-1} from the same field. (CONFIRMED.)

GOTCHA — the raw backend field lnc_size (PassOptions+0x1A4) defaults to 1, not 2. The "2 on Trn2" is an arch-conditioned front-end default that plumbs down into that field; do not state "lnc_size defaults to 2." A single-core compile leaves the field at 1 and lnc_splitter's if (lnc != 1) fast path skips the clone loop entirely. (CONFIRMED — cross-checked against arch/lnc-memory-model and walrus/lnc-splitter.)

Front-end logical_nc_config and backend lnc_size are the same logical quantity N (the LNC mesh size): the front-end shards an N-way mesh, the backend realizes it as N per-core BIR modules. This equality is STRONG, not CONFIRMED — there are two independent confirmations of N with the same semantic, but the exact plumbing from the CLI attribute to PassOptions+0x1A4 spans the driver→penguin boundary and is not traced through a single symbol.

LNCShardingConstraint Is the Contract

(1) front-end SPMD (stock XLA, inside hlo2penguin)
        propagates an HloSharding over the N-way LNC mesh;
        AwsNeuronLNCShardingConstraint PINS specific values to specific LNC tiles  (§1 "sharding=")
            │
(2) after SPMD partitioning, each HLO value carries its per-LNC tile assignment
        → these become BIR shard-id symbols (QuasiAffineExpr shard-id) in ONE symbolic bir::Module
            │
(3) ExpandReplication::run  @0x1553500  (order 7)  ── makes replication explicit on the symbolic module
        LncSplitter::run    @0x16d5ca0  (order 8)  ── clone ×lnc_size, concretize each clone's
                                                       shard-id to its core index, PRUNE off-core
                                                       (evalMask()==0) instructions
            │
(4) per-core BIR modules, SPMD-identical replicas (the LncVerifier invariant)

The sharding the front-end pins via AwsNeuronLNCShardingConstraint is exactly the shard-id assignment the backend splitter concretizes. The constraint is the contract that makes the SPMD-tile↔physical-LNC-core mapping deterministic, so lnc_splitter's per-core prune produces SPMD-identical replicas. The ordering 7-before-8 is mandatory: replication is made explicit on the single module first; running it after the split would break SPMD identity. (CONFIRMED endpoints; STRONG linkage.)

Static-Ring Transfer at the Split Boundary

AwsNeuronTransferWithStaticRing (§3) is the cross-LNC data movement the split materializes: getRemoteCores / getAllCores feed the cross-core DMA / remote-target lowering. Because partition::isCustomCallLoadingParam treats a static-ring transfer as a parameter load (§3), the splitter sees it as a per-core input boundary — where the ring delivers this core's shard — rather than recomputable compute, so each cloned core reads its slice off the static ring. (STRONG.)

Confidence Ledger

Related Components

Cross-References

• SPMD Partitioner Driver — the stock driver/visitor the custom-call targets flow through (13.1)
• Distribution-Strategy Seeding — the companion seam: --distribution-strategy / --spmd / --lnc that seed the mesh before SPMD (13.10)
• Sharding Propagation — registers and consults the CustomCallPartitioner table (13.2)
• SPMD Compute Handlers — the stock per-op handlers HandleCustomCall shares the visitor with (13.5)
• Shardy ↔ HloSharding Bridge — the sdy mesh / HloSharding conversion that produces the sharding pinned here (13.4)
• LNC Splitter — backend lnc_size clone-×N + concretize-shard-id that realizes the front-end pin (Part 8)
• LNC Memory Model — what a Logical NeuronCore is; lnc_size = PassOptions+0x1A4, default 1 / 2 on Trn2 (1.07)