[triton] AMD: PartitionedSharedEncodingAttr의 LLVM lowering 지원으로 공유 메모리 파티셔닝 구현

PR 링크: triton-lang/triton#9362 상태: Merged | 변경: +669 / -138

들어가며

GPU 공유 메모리는 여러 파티션으로 나뉘어 있으며, 같은 파티션에 대한 동시 접근은 충돌(conflict)을 유발합니다. 이번 PR은 텐서를 여러 물리적 공유 메모리 파티션에 분할 저장하는 PartitionedSharedEncodingAttr의 LLVM IR 변환을 구현하여 파티션 충돌을 줄입니다.

핵심 코드 분석

Before: 단일 base pointer만 지원

// Utility.h - 기존 SharedMemoryObject
class SharedMemoryObject {
public:
  SharedMemoryObject(Value base, Type baseElemType, ArrayRef<Value> offsets);
  Value getBase() const { return base; }
private:
  Value base;  // 단일 base pointer
};

After: 다중 base pointer 지원 (파티션별)

// Utility.h - 개선된 SharedMemoryObject
class SharedMemoryObject {
public:
  SharedMemoryObject(Value base, Type baseElemType, ArrayRef<Value> offsets);
  SharedMemoryObject(ArrayRef<Value> bases, Type baseElemType,
                     ArrayRef<Value> offsets);  // 다중 base

  Value getBase() const {
    assert(bases.size() == 1 &&
           "getBase() called on partitioned tensor with multiple bases.");
    return bases[0];
  }
  ArrayRef<Value> getBases() const { return bases; }
  size_t getNumBases() const { return bases.size(); }

private:
  SmallVector<Value> bases;  // 파티션별 base pointer
};

또한 lowerLdStShared와 lowerLdSt 함수도 ArrayRef<Value> smemBases를 받도록 변경되어, linear layout의 partition 차원에 따라 올바른 base를 선택합니다.

왜 이게 좋은가

1. 파티션 충돌 감소: 텐서 데이터를 여러 물리적 파티션에 분산하여 동시 접근 시 충돌을 줄입니다.
2. 기존 코드 호환: getBase() 호출 시 assert로 단일 base인지 확인하여 기존 코드의 잘못된 사용을 조기에 발견합니다.
3. PaddedShared와 조합 가능: getPaddedEncoding(), isPaddedEncoding() 헬퍼 함수로 padding과 partitioning을 함께 처리합니다.

정리

이 PR은 공유 메모리 접근 패턴을 물리적 파티션 수준에서 최적화하는 기반을 제공합니다. 단일 base에서 다중 base로의 추상화 확장은 호환성을 유지하면서도 새로운 메모리 레이아웃을 지원할 수 있는 좋은 설계입니다.

참고 자료

• triton-lang/triton#9362

---

이 글은 AI(Claude)의 도움을 받아 작성되었으며, PR의 실제 diff를 기반으로 분석한 내용입니다.