[Triton] Pipeliner에서 cp_async의 alignment 정보 손실 수정

PR 링크: triton-lang/triton#8752 상태: Merged | 변경: +68 / -23

들어가며

Triton의 pipeliner는 loop 내의 load를 async copy로 변환한다. 이 과정에서 컴파일러 변환(예: loop peeling)이 alignment 정보를 잃어버릴 수 있다. Lowering 단계에서 alignment를 다시 분석하려 해도, 분석이 원래와 다른 결과를 낼 수 있어 vectorization 검증에 실패하고 크래시가 발생한다. 이 PR은 async_copy Op 자체에 contiguity 정보를 저장하여 이 문제를 해결한다.

핵심 코드 분석

Op 정의에 contiguity 추가

def TTG_AsyncCopyGlobalToLocalOp : TTG_Op<"async_copy_global_to_local", [...]> {
  // ...
  let arguments = (ins
    // 기존 인자들...
    TT_Tensor:$src,
    TT_MemDescType:$result,
    Optional<TT_Tensor>:$mask,
    Optional<TT_Tensor>:$other,
    DefaultValuedAttr<I32Attr, "0">:$contiguity  // 새로 추가
  );
}

Pipeliner에서 contiguity 저장

// async_copy 생성 시 분석된 contiguity를 Op에 저장
unsigned contiguity = axisInfoAnalysis.getContiguity(src);
if (mask)
  contiguity = std::min<unsigned>(contiguity,
                                  axisInfoAnalysis.getMaskAlignment(mask));
copyOp.setContiguity(contiguity);

Lowering에서 저장된 contiguity 사용

// 이전: 다시 분석 -> 결과가 다를 수 있음
// unsigned contiguity = axisInfoAnalysis.getContiguity(src);

// 이후: Op에 저장된 값 사용 -> 일관성 보장
unsigned contiguity = copyOp.getContiguity();

왜 이게 좋은가

1. 정보 손실 방지: 컴파일러 변환 후에도 원래의 alignment 정보가 유지된다.
2. 크래시 방지: vectorization 검증 실패로 인한 lowering 크래시를 방지한다.
3. 분석 의존성 제거: lowering 단계에서 axis info 분석에 의존하지 않아 더 robust하다.
4. 기본값 0: contiguity가 설정되지 않은 경우 기존 동작(분석 기반)으로 fallback한다.

정리

이 PR은 async_copy_global_to_local Op에 contiguity attribute를 추가하여, pipeliner에서 분석한 alignment 정보를 lowering까지 안전하게 전달한다. 컴파일러 변환이 analysis 결과를 변경하더라도 올바른 vectorization이 보장된다.

참고 자료

• triton-lang/triton#8752

---

이 글은 AI를 활용하여 PR의 핵심 변경사항을 분석하고 정리한 것입니다. 실제 코드의 맥락은 원본 PR을 참고해 주세요.