본문으로 건너뛰기

[Triton] AMD gfx1250에 buffer_load/buffer_store 연산 노출

들어가며

AMD GPU 아키텍처에서 buffer_loadbuffer_store는 메모리 접근을 최적화하는 핵심 명령어다. CDNA3에서 이미 지원되던 이 연산들을 새로운 gfx1250 아키텍처에서도 사용할 수 있도록 노출하는 간결한 PR이다.

핵심 코드 분석

Before

# gfx1250/__init__.py
from .._ops import _wmma, _verify_wmma
from .._ops import _mma_scaled
from . import tdm

__all__ = ["tdm", "wmma", "wmma_scaled", "get_wmma_scale_layout"]

gfx1250 모듈에는 WMMA와 TDM 연산만 노출되어 있었고, buffer 연산은 없었다.

After

# gfx1250/__init__.py
from .._ops import _wmma, _verify_wmma, _mma_scaled
from .._layouts import AMDWMMALayout
from ..cdna3 import buffer_load, buffer_store
from . import tdm

__all__ = ["tdm", "wmma", "wmma_scaled", "buffer_load", "buffer_store", "get_wmma_scale_layout"]

CDNA3에서 구현된 buffer_load, buffer_store를 import하여 gfx1250에서도 사용할 수 있게 했다. 테스트에서는 불필요한 parametrize를 제거하여 테스트 대상을 CDNA3으로 고정했다.

왜 이게 좋은가

  • 코드 재사용: 새 아키텍처를 위해 buffer 연산을 재구현하지 않고 CDNA3의 구현을 그대로 상속받았다.
  • +7/-9의 최소 변경: 아키텍처 확장을 import 한 줄로 해결한 깔끔한 사례다.

정리

아키텍처 계층 구조를 잘 설계하면, 새 하드웨어 지원을 최소한의 코드로 추가할 수 있다는 것을 보여주는 PR이다.

참고 자료


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

댓글