Avoid mesh size checks

[wujingyue]

Prefer simplicity over unnecessary performance optimizations. Single-GPU performance for a distributed program is largely irrelevant.

[wujingyue]

!test

[github-actions]

Review updated until commit eea10f1

Description

• Remove explicit mesh size checks in communication logic

• Simplify sharding detection by using loop domain pointer checks

• Replace p_sharded and c_sharded boolean variables with direct p_loop_did and c_loop_did checks

• Streamline communication type determination logic

Changes walkthrough

Relevant files

Bug fix

lower_to_communication.cpp Remove mesh size checks and simplify sharding logic            csrc/host_ir/lower_to_communication.cpp Removed mesh size checks (producer_mesh.size() > 1, consumer_mesh.size() > 1) Replaced p_sharded and c_sharded boolean variables with direct loop domain pointer checks Updated all conditional logic to use p_loop_did and c_loop_did directly Simplified sharding detection logic throughout communication analysis +5/-7

PR Reviewer Guide

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review

Potential Logic Error The PR removes mesh size checks (producer_mesh.size() > 1 and consumer_mesh.size() > 1) and replaces them with simple null checks on loop IDs. This could cause unnecessary communication operations when mesh size is 1 (no actual sharding), as the conditions will still trigger communication even when no sharding exists. Need to verify this doesn't break correctness or introduce performance regressions for single-device cases. if (p_loop_did && !c_loop_did) { IterDomain* p_logical_id = getLogicalFromLoopId(producer, p_loop_did); CommunicationType type = same_mesh ? CommunicationType::Allgather : CommunicationType::Gather; fill_communication_info(type, p_logical_id, p2c_map.at(p_logical_id)); } if (!p_loop_did && c_loop_did) { IterDomain* c_logical_id = getLogicalFromLoopId(consumer, c_loop_did); fill_communication_info( CommunicationType::Scatter, c2p_map.at(c_logical_id), c_logical_id); } if (p_loop_did && c_loop_did) {

[wujingyue]

!test

[greptile-apps]

Greptile Overview

Greptile Summary

Removed intermediate boolean variables p_sharded and c_sharded that were checking both pointer nullness and mesh size, replacing them with direct pointer checks (p_loop_did and c_loop_did). This follows the previous commit that removed the mesh size check from the boolean variable definitions, completing the cleanup by eliminating the now-redundant intermediate variables entirely.

• Simplified conditional logic by directly using p_loop_did and c_loop_did pointer checks instead of p_sharded and c_sharded boolean variables
• Removed the now-unused p_sharded and c_sharded variable declarations
• Maintained identical logical behavior since the previous commit already removed mesh size checks from these variables

Confidence Score: 4/5

• Safe to merge - this is a straightforward cleanup refactoring with no logic changes
• The change is a simple code cleanup that removes intermediate boolean variables and uses direct pointer checks instead. The previous commit (7bdff8b) already changed the semantic behavior by removing mesh size checks, and this commit just completes the cleanup by removing the now-redundant variables
• No files require special attention

Important Files Changed

Filename	Overview
csrc/host_ir/lower_to_communication.cpp	Removed intermediate boolean variables p_sharded and c_sharded, directly using pointer checks instead

Sequence Diagram

sequenceDiagram
    participant Caller
    participant getCommunicationInfo
    participant PairwiseLogicalDomainMap
    participant FillCommInfo as fill_communication_info
    
    Caller->>getCommunicationInfo: Call with Expr* e
    getCommunicationInfo->>getCommunicationInfo: Extract producer and consumer TensorViews
    getCommunicationInfo->>PairwiseLogicalDomainMap: Create mapping (producer, consumer)
    PairwiseLogicalDomainMap-->>getCommunicationInfo: Return p2c_map and c2p_map
    
    loop For each ParallelType in kParallelTypeDIDs
        getCommunicationInfo->>getCommunicationInfo: Get p_loop_did and c_loop_did
        
        alt Both nullptr
            getCommunicationInfo->>getCommunicationInfo: Continue to next parallel type
        else p_loop_did exists AND c_loop_did is nullptr (LoadStoreOp)
            getCommunicationInfo->>FillCommInfo: Allgather or Gather
        else p_loop_did is nullptr AND c_loop_did exists (LoadStoreOp)
            getCommunicationInfo->>FillCommInfo: Scatter
        else Both exist (LoadStoreOp)
            getCommunicationInfo->>FillCommInfo: SendRecv or AllToAll
        else ReductionOp/SqueezeOp with p_loop_did
            alt c_loop_did is nullptr
                getCommunicationInfo->>FillCommInfo: Allreduce or Reduce
            else c_loop_did exists and reduction
                getCommunicationInfo->>FillCommInfo: ReduceScatter
            end
        end
    end
    
    alt No communication found
        getCommunicationInfo->>FillCommInfo: Broadcast (nullptr, nullptr)
    end
    
    getCommunicationInfo-->>Caller: Return CommunicationInfo

Loading

[greptile-apps]

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}