Expand distributed indexing, match numpy indexing scheme

[ClaudiaComito]

Description

This pull request introduces a significant overhaul of distributed indexing within dndarray.py, specifically targeting the __getitem__ and __setitem__ methods. The primary objective is to achieve full NumPy indexing compliance in a distributed environment while minimizing MPI overhead and memory footprint.

The logic has been refactored to identify zero-communication paths ("early out"), and route heavy unordered advanced indexing through (hopefully?) optimized communication.

The following table shows the distribution semantics of the DNDarray indexing operations. The first column shows the operation, the second column shows the distribution semantics of the key, and the third column shows the distribution semantics of the value. The last column shows the distribution semantics of the result.

Array distributed	Operation	Key distributed	Value distributed	Result distributed
no	array[key]	no	--	no
no	array[key]	yes	--	yes
yes	array[key]	no	--	no
yes	array[key]	yes	--	yes
no	array[key] = value	no	no	no
no	array[key] = value	no	yes (but why?)	no
no	array[key] = value	yes (but why?)	no	no
no	array[key] = value	yes	yes	no
yes	array[key] = value	no	no	yes
yes	array[key] = value	no	yes	yes
yes	array[key] = value	yes	no	yes
yes	array[key] = value	yes	yes	yes

Routing logic

The flowchart (DRAFT) maps out the MPI routing decisions based on the evaluated state of the indexing key.

graph TD
    classDef default fill:#ffffff,stroke:#ced4da,stroke-width:1px,color:#212529,rx:4px,ry:4px;
    classDef terminal fill:#343a40,stroke:#343a40,stroke-width:2px,color:#ffffff,rx:15px,ry:15px;
    classDef decision fill:#e3f2fd,stroke:#4dabf7,stroke-width:2px,color:#000000;
    classDef highlight fill:#e8f5e9,stroke:#69b3a2,stroke-width:2px,color:#212529;

    Start(["Start: arr[key] or arr[key] = value"]):::terminal
    Norm["Normalize Key (e.g., Bool Masks -> Int Indices)"]
    ProcessKey["__process_key() Expands dims, aligns shapes"]
    StateCalc{"Calculate State: split_key_is_ordered"}:::decision
    
    Start --> Norm
    Norm --> ProcessKey
    ProcessKey --> StateCalc
    
    Branch1{"Single item on split axis? (root != None)"}:::decision
    Branch0{"Operation?"}:::decision
    BranchNeg1{"Operation?"}:::decision
    
    StateCalc -- "1: Ordered / Ascending" --> Branch1
    StateCalc -- "0: Unordered / Random" --> Branch0
    StateCalc -- "-1: Descending Slice" --> BranchNeg1
    
    subgraph Ordered ["Fast Path: Ordered Indexing (split_key_is_ordered = 1)"]
        style Ordered fill:#f8f9fa,stroke:#dee2e6,stroke-width:1px,stroke-dasharray: 5 5,color:#495057
        Branch1 -- "Yes: Get" --> RootGet["Root fetches local data"]
        RootGet --> Bcast["MPI.Bcast to all ranks"]
        
        Branch1 -- "Yes: Set" --> RootSet["Root updates local data in-place"]
        
        Branch1 -- "No" --> FastLocal["Pure Basic Slicing: Apply locally, NO MPI needed"]:::highlight
    end
    
    subgraph Descending ["Descending Slices (split_key_is_ordered = -1)"]
        style Descending fill:#f8f9fa,stroke:#dee2e6,stroke-width:1px,stroke-dasharray: 5 5,color:#495057
        BranchNeg1 -- "Set" --> FlipVal["Flip 'value' array"]
        FlipVal --> MatchDist["Align distribution map"]
        MatchDist --> SetLocal["-1 Local Set"]
        
        BranchNeg1 -- "Get" --> UnorderedFallback["Converts to arange -> falls back to unordered"]
    end
    
    subgraph Unordered ["Heavy Path: Unordered Advanced Indexing (split_key_is_ordered = 0)"]
        style Unordered fill:#f8f9fa,stroke:#dee2e6,stroke-width:1px,stroke-dasharray: 5 5,color:#495057
        
        Branch0 -- "Get" --> G_Allgather["MPI.Allgather: Share recv_counts"]
        G_Allgather --> G_SendIdx["MPI.Isend/Recv: Send requested indices to owning ranks"]
        G_SendIdx --> G_Fetch["Owning ranks fetch local data"]
        G_Fetch --> G_SendData["MPI.Isend/Recv: Send requested data back"]
        G_SendData --> G_Reconstruct["Reconstruct recv_buf on original rank"]
        
        Branch0 -- "Set" --> S_CheckVal{"Is 'value' distributed?"}:::decision
        S_CheckVal -- "No / Scalar" --> S_LocalMask["_advanced_setitem_unordered_local (Apply locally)"]
        
        S_CheckVal -- "Yes" --> S_Align["Redistribute 'value' to match 'key' distribution"]
        S_Align --> S_AllToAll["MPI.Alltoallv: Exchange data AND indices"]
        S_AllToAll --> S_ApplyRecv["Apply received data to local elements"]
    end

    Bcast --> End(["Return / Complete"]):::terminal
    RootSet --> End
    FastLocal --> End
    SetLocal --> End
    UnorderedFallback -.-> Branch0
    G_Reconstruct --> End
    S_LocalMask --> End
    S_ApplyRecv --> End

Loading

Main changes

• abstracts key parsing and alignment into a centralized private method that handles dimension expansion, shape broadcasting, and classifies the state of the indexing operation to determine network routing.
• enforces standard last-assignment-wins semantics for advanced indexing duplicates on cuda tensors by generating linear indices and mapping local occurrence priorities (thanks @Hakdag97 ).
• intercepts multidimensional and single-dimensional boolean masks early in the pipeline, converting them to explicit integer configurations locally to prevent unnecessary cross-rank broadcasting.
• maps and isolates zero-communication assignments during slice operations, executing completely local pytorch tensor modifications when the requested indices and data already reside on the active rank.
• structures unordered read requests by compiling global communication matrices, enabling the dispatch of non-blocking Isend and Recv calls strictly between nodes that own the requested indices and those requesting them.
• forces distribution alignment during set operations if the right-hand side assignment value is also distributed, utilizing an Alltoallv operation to shuffle payload data and target indices concurrently.
• introduces a value broadcasting helper function to natively squeeze or expand the dimensions of scalar or tensor payloads to match the specific dimensional footprint of the target slice before assignment occurs.

To Be Continued...

Memory footprint

Scaling behaviour

Issue/s resolved: #914 #918

Changes proposed:

• feature extension in __process_key, getitem, and setitem methods
• edge case handling
• extensive comparison to numpy API in unittests

Type of change

Memory requirements

Performance

Due Diligence

• All split configurations tested
• Multiple dtypes tested in relevant functions
• Documentation updated (if needed)
• Updated changelog.md under the title "Pending Additions"

Does this change modify the behaviour of other functions? If so, which?

yes / no

skip ci

[ClaudiaComito]

• pass key to process_key immediately, separate by view/copy output

[JuanPedroGHM]

Create different functions for scalar:

• Scalar
• Mask
• Slice
• Advanced indexing
• Distributed key

Process key should return the tuple with the key information, and send to the separate functions.

Additionally, unnecessary function definitions in __getitem/__setitem.

Make table and diagram available in a markdown file in the docs folder for now.

[brownbaerchen]

I haven't looked at the actual advanced indexing stuff. But I think we would do well to clean up a bit before getting into the details of this. There have been some changes mixed in that, to me, seem unrelated to advanced indexing in heat. They should be moved to separate PRs if we want to keep them. Other changes are cosmetic or temporary and should be removed entirely.

[brownbaerchen]

Can we remove this change from this PR?

[brownbaerchen]

Let's remove this change from this PR (see #2216 for a similar case)

[brownbaerchen]

Not sure about these changes, but some them can be removed from this PR, right?

[brownbaerchen]

Let's remove this change from this PR.

[brownbaerchen]

Let's remove these changes from this PR.

[brownbaerchen]

These changes seem unrelated to advanced indexing. Maybe put in a separate PR.

[brownbaerchen]

Is this needed here?

[brownbaerchen]

Let's remove this from the PR

[brownbaerchen]

This seems unrelated to advanced indexing. Maybe move to a separate PR?

[brownbaerchen]

This seems unrelated to advanced indexing. Maybe move to separate PR?