TIRx: 진화하는 프론티어 ML 커널을 위한 오픈 컴파일러 스택

오늘 우리는 Apache TVM을 기반으로 구축된 ML 커널을 위한 오픈 소스 하드웨어 네이티브 DSL(Domain-Specific Language)이자 컴파일러인 TIRx를 소개합니다. TIRx는 빠르게 변화하는 커널과 빠르게 변화하는 하드웨어가 만나는 AI 소프트웨어 스택의 영역을 목표로 합니다. TIRx는 현재 GPU 및 특화된 AI 가속기로 컴파일되며, 이후 세대의 하드웨어와 함께 성장하도록 설계되었습니다. 동일한 설계는 전문가가 작성한 커널, 에이전트가 생성한 커널, 그리고 메가커널(megakernel) 시스템 모두에 적용됩니다.

우리는 출시와 함께 다음과 같은 자료를 제공하기 위해 더 넓은 커뮤니티와 협력해 왔습니다:

PyPI wheel 및 Python 프론트엔드. @T.jit / @T.prim_func 스타일의 저작 방식, 파서 유틸리티, 그리고 TIRx 프로그램을 구성하기 위한 Python API를 갖춘 Python 임베디드 하드웨어 네이티브 커널 DSL입니다.

TIRx 커널 라이브러리 및 벤치마크. Blackwell GPU에서의 GEMM, 어텐션(attention) 스타일 커널, 저정밀도 연산자(low-precision operators)를 다루는 엔드 투 엔드(end-to-end) 예제를 포함합니다.

현대적 GPU 프로그래밍에 관한 공개 강의. 이 큐레이션된 온라인 강의는 Carnegie Mellon University의 머신러닝 시스템 과정의 일부로 강의되었으며, 학생들에게 머신러닝 시스템을 위한 현대적 GPU 프로그래밍을 가르치기 위해 TIRx를 사용합니다.

다음 리소스를 확인하실 수 있습니다:

• GitHub: https://github.com/apache/tvm
• Documentation: https://tvm.apache.org/docs/tirx/overview.html
• PyPI wheel: https://pypi.org/project/apache-tvm/0.25.0/

pip install apache-tvm==0.25.0

• Community TIRx 커널 라이브러리: https://github.com/mlc-ai/tirx-kernels
• 머신러닝 시스템을 위한 현대적 GPU 프로그래밍: https://mlc.ai/modern-gpu-programming-for-mlsys/index.html

동기 (Motivation)

커널 DSL (Domain-Specific Languages)은 프로그래머와 머신(machine) 사이의 적절한 경계를 선택할 때 가장 효과적입니다. 성숙한 커널과 성숙한 하드웨어의 경우, 그 경계는 상위 수준(high-level)일 수 있습니다. 즉, 컴파일러가 스레드 할당(thread assignment), 메모리 이동(memory movement), 레이아웃(layout) 세부 사항 및 명령어 선택(instruction selection)을 간결한 텐서(tensor) 또는 타일(tile) 추상화 뒤로 숨기는 것입니다. Triton이 전형적인 사례이며, Triton의 채택은 이러한 방식이 기존의 커널 패턴에 얼마나 잘 작동하는지를 보여줍니다. 하지만 프론티어(frontier) 영역에서는 동일한 경계가 더 큰 압박을 받습니다. 새로운 명령어(instruction), 메모리 공간(memory spaces), 협업 패턴(cooperation patterns) 및 커널 알고리즘이 컴파일러가 이를 잘 자동화할 수 있는 내장 메커니즘을 갖추기도 전에 등장하는 경우가 많기 때문입니다. 이런 상황이 발생하면, 상위 수준의 컴파일러가 통상적으로 숨기는 부분들이 바로 전문가가 여전히 수동으로 제어해야 하는 부분이 됩니다.

TIRx(발음: “tier-ex”)는 세 가지 결정 사항을 중심으로 구성된, 더 낮고 명시적인 경계를 선택함으로써 이에 대응합니다.

오케스트레이션(Orchestration)은 하드웨어 네이티브 소스(hardware-native source)에 유지됩니다. 파이프라인 구조(pipeline structure), 동기화(synchronization), 역할 할당(role assignment), 메모리 배치(memory placement) 및 백엔드 인트린직(backend intrinsics)은 프론티어에서 전문가의 제어가 가장 빈번하게 필요한 부분들입니다. 따라서 TIRx는 이러한 요소들을 아직 새로운 기능을 모델링하지 못할 수도 있는 추상화 뒤에 숨기지 않고 소스 코드에 유지합니다.

반복되는 타일 프리미티브(tile primitives)는 컴파일러에 노출됩니다. 실행 범위(execution scope), 텐서 레이아웃(tensor layout) 및 타일 프리미티브 디스패치(tile primitive dispatch)를 통해, 커널 전체를 고정된 컴파일러 파이프라인에 강제로 통과시키지 않고도 일반적인 연산들이 백엔드 전반에 걸쳐 재사용 가능하고, 분석 가능하며, 이식 가능하도록 유지할 수 있습니다. 하드웨어 네이티브 제어의 비용은 엔지니어링 노력입니다. 즉, 각 커널과 백엔드에 대해 모든 연산을 수동으로 작성하는 것은 매우 고된 작업입니다.

반복되는 연산을 타일 프리미티브 (tile primitives)로 노출하면 이러한 문제를 완화할 수 있습니다. 즉, 작성자가 매번 동일한 데이터 이동(data movement)이나 행렬 곱셈(matrix multiply)을 다시 작성하는 대신, 디스패치된 구현체를 재사용할 수 있게 됩니다. 새로운 하드웨어는 처음에는 인트린직 (intrinsics)으로 도입되고, 이후에 타일 프리미티브로 도입됩니다. 새로운 기능은 단일 하드웨어 연산에 대한 얇고 백엔드 특화된 래퍼(wrapper)인 네이티브 인트린직으로서 즉시 사용될 수 있습니다. 커널 전반에 걸쳐 사용 패턴이 안정화되면, 이를 스코프(scope), 피연산자(operand), 백엔드에 걸쳐 디스패치되는 레이아웃 인식 연산인 타일 프리미티브로 승격시킬 수 있습니다. 핵심 추상화는 작게 유지되며, 새로운 기능을 위한 인트린직을 추가하더라도 기존 기능이 깨지는 일은 없습니다.

그 결과, 하드웨어와 함께 성장할 수 있는 DSL 및 컴파일러 스택이 탄생합니다. 이것이 TIRx의 핵심 설계 철학입니다. 즉, 기반은 작고 명시적으로 유지하고, 새로운 가속기 세대가 등장함에 따라 백엔드 라이브러리가 진화하도록 하는 것입니다.

이러한 방식은 TIRx를 TileLang과 같은 시스템보다 하위 계층에 위치시킵니다. TileLang 또한 메모리 스코프와 파이프라이닝 (pipelining)을 노출하여 Triton 대비 경계를 낮추면서도, 레이아웃 추론 (layout inference)과 스레드 바인딩 (thread binding)은 컴파일러에 맡깁니다. TIRx는 의도적으로 이러한 고수준의 관심사들을 핵심 영역 외부에 두고, 그러한 시스템들이 그 위에 구축할 수 있는 최소한의 기반을 제공합니다. 우리는 TIRx를 TileLang 컴파일을 지원하는 새로운 최소 기반으로 가져오기 위해 TileLang 커뮤니티와 협력하고 있습니다.

이와 같이 작고 명시적인 기반 덕분에, 엔지니어링 노력을 최대한 줄이면서도 최고 성능을 추구하는 여러 유형의 사용자들을 지원하는 설계를 할 수 있습니다. 즉, 전문가가 작성한 프로덕션 커널 (production kernels), 에이전트가 생성한 커널 (agent-generated kernels), 그리고 메가커널 (megakernel) 시스템이 이에 해당하며, 이들 각각은 네이티브 수준의 제어와 컴파일러가 인식할 수 있는 반복 연산 모두를 필요로 합니다.

이 포스트의 나머지 부분에서는 프로그래밍 모델을 살펴본 후, 이러한 각 방향을 차례대로 설명합니다.

TIRx 프로그래밍 모델

실제 환경에서 해당 경계가 어떻게 나타나는지는 다음과 같습니다. TIRx 프로그램은 구조화된 네이티브 커널 (native kernel)로 읽힙니다. 루프 (loops), 분기 (branches), 텐서 (tensors), 동기화 (synchronization), 파이프라인 상태 (pipeline state), 그리고 백엔드 인트린직 (backend intrinsics)이 직접 작성됩니다. 반복되는 하드웨어 연산이 재사용 및 디스패치 (dispatch) 가능해져야 하는 지점에는 타일 프리미티브 (Tile primitives)가 나타납니다. 모델의 대부분은 세 가지 요소로 구성됩니다.

**실행 범위 (Execution scope)**는 누가 어떤 입도 (granularity)로 연산을 실행할지를 결정합니다. 두 가지 요소가 이를 선택합니다: 특정 영역에 진입하는 하드웨어 역할을 선택하는 제어 흐름 (control flow), 그리고 호출의 입도를 설정하는 프리미티브 네임스페이스 (primitive namespace)입니다.

수식어가 없는 Tx.* 호출은 스레드 (thread) 레벨에서 실행되며, Tx.wg.*는 워프 그룹 (warpgroup) 레벨에서 실행됩니다. T.ptx.elect_sync()와 같은 프레디케이트 (predicate)는 스레드 레벨 호출을 더 좁혀 단일 발행 스레드 (issuing thread)까지 제한할 수 있습니다.

**텐서 레이아웃 (Tensor layout)**은 스토리지 우선 (storage-first) 인터페이스를 통해 논리적 텐서가 어디에 존재하는지를 설명합니다. 타일 (tile)은 글로벌 메모리 (global memory), 공유 메모리 (shared memory), 레지스터 (registers), 텐서 메모리 (tensor memory), 또는 가속기 SRAM에 위치할 수 있습니다. 사용자는 각 타일이 어디에 위치하는지, 그리고 그 요소들이 레인 (lanes), 워프 (warps), 레지스터에 어떻게 분산되는지를 선언하며, 해당 선언은 타일에 부착된 상태로 유지됩니다. 프리미티브가 호출되면 컴파일러는 이러한 선언을 읽어 구현체를 선택합니다. 레이아웃은 스토리지 기술 (storage description)이지 루프 변환 유틸리티 (loop-transformation utility)가 아닙니다. 사용자는 타일의 레이아웃을 구성할 수는 있지만, 루프를 변환하기 위해 레이아웃을 사용하지는 않습니다.

**타일 프리미티브 디스패치 (Tile primitive dispatch)**는 하나의 호출을 네이티브 IR로 변환합니다. 피연산자 레이아웃 (operand layouts), 실행 범위 (execution scope), 타겟 (target), 또는 명시적인 dispatch= 힌트로부터 컴파일러는 일치하는 구현체를 선택합니다. 예를 들어, 글로벌에서 공유 메모리로의 복사는 TMA로, 공유 메모리에서 레지스터로의 복사는 ldmatrix로, 텐서 메모리에서 레지스터로의 복사는 tcgen05.ld로 해결됩니다. 행렬 곱셈 (matrix multiply)은 WGMMA, tcgen05, 또는 시스톨릭 어레이 (systolic-array) 명령어로 해결됩니다. 그런 다음 디스패치는 해당 명령어를 전체 타일에 적용하는 데 필요한 루프와 주소 지정 (addressing)을 생성합니다.

이러한 요소들은 스코프(scope)가 중요한 곳이라면 어디든 결합됩니다. 아래의 GEMM 에필로그(epilogue)에서는 워프그룹 스코프(warpgroup-scoped) 및 스레드 스코프(thread-scoped) 프리미티브(primitives)가 동일한 영역인 Tx.wg.*에 위치합니다.

이 호출들은 워프그룹(warpgroup) 전체에 걸쳐 타일(tile)을 이동시키고 캐스트(cast)하며, 명시적인 발행 스레드 술어(issuing-thread predicate)에 의해 보호되는 최종 스레드 스코프 Tx.copy_async가 TMA 스토어(store)를 수행합니다.

위의 발췌본들은 단순화된 것입니다. 전체적인 그림을 이해하기 위해, 완전한 FP16/BF16 GEMM 커널의 두 가지 역할인 TMA 프로듀서(producer)와 텐서 메모리 라이트백(tensor-memory writeback)을 살펴보겠습니다. 이를 한 줄씩 읽을 필요는 없습니다. 핵심은 오케스트레이션(orchestration)과 관련된 모든 것(파이프라인 상태, 배리어 프로토콜(barrier protocol), 역할 선택, tcgen05.wait 및 cp_async.bulk와 같은 저수준 동기화 인트린직(synchronization intrinsics))은 일반 소스 코드에 머무는 반면, 반복되는 데이터 이동은 스코프, 레이아웃(layout), 디스패치(dispatch) 설정으로부터 로워링(lowering)이 선택되는 타일 프리미티브(tile primitives)로 나타난다는 점입니다.

세 가지 요소 중 레이아웃은 가장 많은 설계 결정을 수반하므로 자세히 살펴볼 가치가 있습니다.

텐서 레이아웃을 위한 스토리지 우선 인터페이스 (A Storage-First Interface for Tensor Layouts)

TIRx는 레이아웃을 텐서 스토리지의 일급 표현(first-class representation)으로 취급합니다. CuTe에 익숙한 독자라면 이 영역이 낯설지 않을 것입니다. 두 시스템 모두 텐서 데이터가 하드웨어 리소스에 어떻게 매핑되는지 설명하기 위해 레이아웃을 사용하지만, CuTe는 타일 작업이 스레드 간에 어떻게 분할되는지 도출하기 위한 프로그래밍 가능한 인터페이스로 레이아웃을 노출하는 반면, TIRx는 프리미티브 디스패치에 의해 소비되는 스토리지 계약(storage contract)으로 레이아웃을 사용합니다.

TIRx 레이아웃은 논리적 텐서 인덱스를 명명된 축(axes) 상의 물리적 좌표로 매핑합니다. 이 모델은 스트라이드(strides)를 의미론적 하드웨어 축에 부착하고 명시적인 샤드(shard), 레플리카(replica), 오프셋(offset) 구성 요소를 추가함으로써 셰이프-스트라이드(shape-stride) 레이아웃을 일반화합니다. 샤드는 논리적 요소가 물리적 축을 따라 어떻게 분할되는지를 설명합니다. 레플리카는 동일한 논리적 요소가 어디에 복제되는지를 설명합니다. 오프셋은 물리적 배치가 어디서 시작되는지를 설명합니다. 구체적으로,

D (Shard, 샤드). 하나 이상의 이터레이터(iterator) 목록으로, 각 이터레이터는 특정 축(axis)에 대한 범위(extent)와 스트라이드(stride)를 가집니다. D는 이러한 이터레이터들에 걸쳐 논리적 인덱스(logical index)를 분할(partition)하고 기본 좌표(base coordinate)를 생성합니다. 이는 shape-stride 개념을 다중 축으로 일반화한 것입니다. R (Replica, 레플리카). 논리적 인덱스와 독립적으로 하드웨어 공간에서의 오프셋(offset)을 열거하는 복제 이터레이터(replication iterator)들의 집합입니다. 이 집합의 각 요소를 D의 결과에 더하면 복제(replication) 또는 브로드캐스팅(broadcasting)이 수행됩니다. O (Offset, 오프셋). 모든 결과에 고정된 좌표 오프셋(축당 하나의 정수)이 더해집니다. 이를 통해 데이터를 특정 기본 위치에 배치하거나 독점적인 리소스를 예약할 수 있습니다.

TIRx 레이아웃 Python API의 구체적인 예시는 다음과 같습니다:

이는 레인(lane)과 워프(warp)에 걸쳐 분산되고, 다른 워프 그룹(warpgroup)에 걸쳐 복제되며, 워프 축(warp axis) 상의 오프셋에 배치된 논리적 타일(logical tile)을 나타냅니다. (8, 16) 형상(shape) 공간의 논리적 좌표 (i, j)가 주어지면, 다음과 같이 계산하여 각각 워프(warp), 레인(lane), 레지스터(reg) 축으로 매핑합니다.

[\begin{aligned} L(i,j)_{(8,16)} &= L(i\cdot 16 + j) && \text{(flatten)} \ &= L\bigl(i,\ \lfloor j/8\rfloor,\ \lfloor j/2\rfloor,%,4,\ j,%,2\bigr) && \text{(unflatten)} \end{aligned}] [\begin{cases} @\mathrm{warp}:\ {,\lfloor j/8\rfloor + 5 + 4r \mid r \in [0,2),} \ @\mathrm{lane}:\ 4i + \lfloor j/2\rfloor,%,4 \ @\mathrm{reg}:\ \ j,%,2 \end{cases}]예를 들어, 논리적 (3, 9) 위치의 57번 요소는 다음과 같이 매핑됩니다:

• 기본 위치(base location): 6@warpid, 12@laneid, 1@m
• 소유자(owners) (레플리카를 통해 2배): { warpid=6 laneid=12 }, { warpid=10 laneid=12 }

(대화형 데모를 열고 57번 요소를 클릭하여 정확히 이러한 소유자들을 확인해 보세요.)

TIRx의 레이아웃 인터페이스는 네 가지 설계 선택을 중심으로 구축되었습니다.

1. 레이아웃은 작업 분할(work-partitioning) 인터페이스가 아니라 저장 계약(storage contract)입니다.

CuTe에서 레이아웃(layout)은 단순히 데이터 배치(data placement)를 나타내는 표현일 뿐만 아니라, 타일 연산(tile operations)이 스레드(threads)에 어떻게 분산되는지를 도출하기 위한 프로그래밍 인터페이스(programming interface)의 일부이기도 합니다. 사용자들은 복사(copy) 및 연산(compute) 작업을 위한 데이터와 작업 분배를 표현하기 위해 레이아웃을 구성(compose), 타일링(tile), 분할(partition)합니다. TIRx는 이 경계를 다르게 설정합니다. 사용자는 각 타일의 저장 레이아웃(storage layout)을 기술하고, 해당 타일들에 대해 타일 프리미티브(tile primitives)를 호출합니다. 레이아웃은 샤딩(sharding), 복제(replication), 오프셋(offset)을 포함하여 논리적 텐서 좌표(logical tensor coordinates)가 물리적 하드웨어 좌표(physical hardware coordinates)에 어떻게 매핑되는지를 기록하며, 실행 분할(execution partitioning)을 구축하는 데 사용되는 표면(surface)이 아닙니다. 프리미티브(primitive)가 로워링(lowered)될 때, 디스패치(dispatch)는 피연산자 레이아웃(operand layouts), 실행 범위(execution scope), 백엔드 타겟(backend target)을 사용하여 스레드 분할(thread partitioning), 루프 중첩(loop nest), 어드레싱(addressing), 명령 시퀀스(instruction sequence)를 생성합니다. 이러한 의미에서 TIRx의 레이아웃은 저장(storage)을 정확하게 표현하기만 하면 됩니다. 변환 로직(transformation logic)은 사용자가 작성한 레이아웃 구성(layout composition)이 아니라 프리미티브 디스패치(primitive dispatch) 내부에 존재합니다.

2. 레이아웃은 논리적 텐서 좌표를 물리적 하드웨어 좌표에 매핑합니다.

명시적인 복제(replica) 및 오프셋(offset) 구조는 지정된 논리적-물리적 공식화(logical-to-physical formulation)에서 비롯됩니다. 레이아웃을 공식화하는 한 가지 대안은 물리적 위치를 논리적 좌표에 매핑하는 것인데, 이렇게 하면 여러 물리적 위치에 저장되는 하나의 논리적 요소인 복제(replication)를 여전히 점값 함수(point-valued function)로 정의할 수 있습니다. 그러나 스트라이드 패턴(strided pattern)으로 물리적 위치에 걸쳐 있는 텐서의 경우, 일부 물리적 위치는 잘 정의된 매핑을 갖지 못할 수도 있습니다.

3. 레이아웃은 일반적인 형상(general shapes)을 지원합니다.