[Triton] translator_helpers에서 builtin 사용 정리

2026년 3월 18일

들어가며

Triton-to-Gluon translator는 Triton 커널을 Gluon API로 자동 변환하는 도구다. 기존에는 tl.atomic_add, tl.make_tensor_descriptor 같은 builtin 함수들이 translator_helpers에서 수동으로 재작성(rewrite)되었다. 이 PR은 gluon.constexpr_function 헬퍼를 활용하여 불필요한 rewrite를 제거하고, Gluon API의 실제 버그도 수정한다.

핵심 코드 분석

Before

# 불필요한 수동 rewrite 목록
(tl.atomic_add, "tl_atomic_add"),
(tl.make_tensor_descriptor, "tl_make_tensor_descriptor"),
(tl.load_tensor_descriptor, "tl_load_tensor_descriptor"),
(tl.store_tensor_descriptor, "tl_store_tensor_descriptor"),

After

# atomic_add rewrite 제거 - Gluon이 직접 처리
# make_tensor_descriptor도 Gluon이 직접 처리

Gluon의 atomic_add에서 scalar mask 지원이 추가되었다:

# Before: mask가 tensor일 때만 작동
ttgl.atomic_add(offset + ptr, val, mask=mask)

# After: scalar mask도 지원
scalar_mask = True
ttgl.atomic_add(offset + ptr, val, mask=scalar_mask, sem="acquire", scope="cta")

semantic 쪽에서는 _broadcast_ptr_val_mask 헬퍼로 중복 코드를 제거했다:

def _broadcast_ptr_val_mask(self, ptr, val, mask):
    ptr_shape = ptr.shape
    if mask is None:
        ptr, val = self.broadcast_tensors(ptr, val)
    else:
        ptr, val, mask = self.broadcast_tensors(ptr, val, mask)
    if ptr_shape != ptr.shape:
        raise ValueError(...)
    return ptr, val, mask