[Triton] AMD ConvertToBufferOps에서 i64 offset 지원

2026년 3월 6일

들어가며

AMD GPU의 buffer 명령어(buffer_load/buffer_store)는 32비트 offset만 지원한다. 기존 Triton의 ConvertToBufferOps 패스는 64비트 offset을 사용하는 load/store를 무조건 거부하여, flex attention 같은 64비트 포인터 연산을 사용하는 커널이 느린 일반 load로 fallback되었다. 이 PR은 offset이 안전하다고 증명 가능할 때 자동으로 32비트로 truncate하여 빠른 buffer 명령어를 사용하도록 한다.

핵심 코드 분석

Before

i64 offset이 있는 load는 tt.load로 남아 buffer 변환이 일어나지 않았다:

%ptr = tt.addptr %base, %offset : tensor<1024x!tt.ptr<f32>>, tensor<1024xi64>
%val = tt.load %ptr : tensor<1024x!tt.ptr<f32>>

After

offset이 안전한 범위(non-negative, bounded)라고 증명되면 trunci 후 buffer 연산으로 변환된다:

%truncated = arith.trunci %offset : tensor<1024xi64> to tensor<1024xi32>
%val = amdg.buffer_load %base_ptr[%truncated] : tensor<1024xf32>

안전성 증명은 세 가지 경로로 수행된다:

Range analysis: make_range + extsi + 상수 곱셈이면 값 범위를 추론할 수 있다
tt.pointer_range 속성: 커널 인자에 tt.pointer_range = 32가 설정되어 있으면 32비트 범위 보장
음수 가능성 검사: offset이 음수일 수 있으면 변환하지 않는다

// 안전한 케이스: make_range [0, 256) * 1024 => 항상 non-negative, bounded
%range = tt.make_range {end = 256, start = 0}
%range_ext = arith.extsi %range : tensor<256xi32> to tensor<256xi64>
%offset = arith.muli %range_ext, %stride  // stride=1024

// 위험한 케이스: 임의 i64 값 => 변환 안 함
%splat_off = tt.splat %arg1 : i64 -> tensor<256xi64>