[SymMem 3/5] Add MemoryType::Symmetric

[samnordmann]

• [SymMem 1/5] VMM and multicast reference #5516
• [SymMem 2/5] SymmetricTensor runtime type #5517
• Here ==> [SymMem 3/5] Add MemoryType::Symmetric #5518
• [SymMem 4/5] Broadcast and Allgather lowering + runtime support #5519
• [SymMem 5/5] Contiguous View #5520

Full branch for reference: #5515

[greptile-apps]

Greptile Overview

Greptile Summary

This PR adds MemoryType::Symmetric to support symmetric memory allocations, which enable efficient multi-device operations through CUDA VMM and NVLS multicast. The implementation adds the new memory type enum, preserves it during fusion input/output setup and host IR lowering, and implements allocation logic using SymmetricTensor::allocate() with a contiguity requirement.

Key changes:

• Added MemoryType::Symmetric enum value and string representation
• Modified Fusion::addInput, addOutput, and host IR lowering to preserve symmetric memory type instead of forcing global
• Added symmetric allocation path in HostIrEvaluator::handle(kir::Allocate) with contiguity validation
• Updated allocation sites to use tv->getMemoryType() instead of hardcoding MemoryType::Global
• Added test case validating symmetric memory allocation and properties

Issue found:

• Fusion::replaceOutput still forces MemoryType::Global (acknowledged by TODO comment)

Confidence Score: 4/5

• Safe to merge with one known gap in replaceOutput (documented in TODO)
• The implementation is well-structured and consistent across most of the codebase. However, Fusion::replaceOutput doesn't preserve symmetric memory type, which could cause issues if a symmetric tensor is replaced. The TODO comment acknowledges this gap, suggesting it's a known limitation being tracked for future work.
• csrc/fusion.cpp - replaceOutput needs symmetric memory support

Important Files Changed

File Analysis

Filename	Score	Overview
csrc/fusion.cpp	4/5	Preserved symmetric memory type when adding inputs/outputs, but replaceOutput still forces global memory
csrc/host_ir/lowering.cpp	5/5	Changed allocations to respect the tensor's existing memory type rather than forcing global
csrc/host_ir/evaluator.cpp	5/5	Added symmetric memory allocation path using SymmetricTensor::allocate with contiguity check

Sequence Diagram

sequenceDiagram
    participant User
    participant Fusion
    participant TensorView
    participant HostIrLower
    participant HostIrEvaluator
    participant SymmetricTensor
    
    User->>Fusion: addInput(tv) / addOutput(tv)
    Fusion->>TensorView: getMemoryType()
    alt MemoryType != Symmetric
        Fusion->>TensorView: setMemoryType(Global)
    end
    
    User->>HostIrLower: lower(Fusion)
    HostIrLower->>TensorView: getMemoryType()
    alt MemoryType != Symmetric
        HostIrLower->>TensorView: setMemoryType(Global)
    end
    HostIrLower->>HostIrLower: Create Allocate HIR nodes
    Note over HostIrLower: Uses tv->getMemoryType()
    
    User->>HostIrEvaluator: runWithInput()
    HostIrEvaluator->>HostIrEvaluator: handle(kir::Allocate)
    HostIrEvaluator->>TensorView: getMemoryType()
    
    alt MemoryType == Symmetric
        HostIrEvaluator->>HostIrEvaluator: isTvContiguous(tv)
        Note over HostIrEvaluator: Error if not contiguous
        HostIrEvaluator->>SymmetricTensor: allocate(sizes, dtype, device)
        SymmetricTensor->>SymmetricTensor: Check VMM support
        SymmetricTensor->>SymmetricTensor: cuMemCreate with granularity
        SymmetricTensor->>SymmetricTensor: cuMemAddressReserve & cuMemMap
        SymmetricTensor-->>HostIrEvaluator: at::Tensor
    else MemoryType != Symmetric
        HostIrEvaluator->>HostIrEvaluator: empty_strided_cuda()
    end

Loading

[greptile-apps]

_{8 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

logic: The TODO is valid - replaceOutput should preserve symmetric memory like addInput and addOutput do

Suggested change

	MemoryType::Global); // TODO: may need a patch to support symmetric
	// memory type
	if (tv->getMemoryType() != MemoryType::Symmetric) {
	tv->setMemoryType(MemoryType::Global);
	}