[Triton] AMD gfx1250에 buffer_load/buffer_store 연산 노출
들어가며
AMD GPU 아키텍처에서 buffer_load와 buffer_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를 기반으로 합니다.
댓글