[triton] Out-of-tree TTIR/TTGIR 패스 플러그인 시스템
PR 링크: triton-lang/triton#8401 상태: Merged | 변경: +839 / -1
들어가며
Triton의 컴파일 파이프라인을 커스터마이즈하려면 Triton 코어를 직접 수정해야 했습니다. 이 PR은 TRITON_PASS_PLUGIN_PATH 환경변수로 동적 라이브러리를 로딩하고, C API를 통해 커스텀 컴파일 패스를 등록/실행할 수 있는 플러그인 시스템을 도입합니다.
핵심 코드 분석
1. TritonPlugin 핵심 구조
struct TritonPlugin {
llvm::Error loadPlugin();
llvm::Expected<TritonPluginResult>
getPassHandles(std::vector<const char *> &handles);
llvm::Expected<TritonPluginResult>
addPass(mlir::PassManager *pm, const char *passHandle);
llvm::Expected<TritonPluginResult> registerPass(const char *passHandle);
};
2. 플러그인 API 규약
#define TRITON_PLUGIN_API \
extern "C" __attribute__((visibility("default"))) TritonPluginResult
// 플러그인이 구현해야 하는 함수들:
TRITON_PLUGIN_API tritonEnumeratePluginPasses(uint32_t *count, const char **handles);
TRITON_PLUGIN_API tritonAddPluginPass(mlir::PassManager *pm, const char *passHandle);
TRITON_PLUGIN_API tritonRegisterPluginPass(const char *passHandle);
3. 자동 로딩
if (std::string filename =
mlir::triton::tools::getStrEnv("TRITON_PASS_PLUGIN_PATH");
!filename.empty()) {
TritonPlugin TP(filename);
std::vector<const char *> passNames;
TP.getPassHandles(passNames);
for (const char *passName : passNames)
TP.registerPass(passName);
}
왜 이게 좋은가
- 비침습적 확장: Triton 코어를 수정하지 않고 컴파일 패스를 추가할 수 있습니다.
- 동적 로딩:
llvm::sys::DynamicLibrary를 사용하여 런타임에 플러그인을 로드합니다. - LLVM 패턴 준수: LLVM의 pass plugin 구조와 유사한 API 설계입니다.
정리
Triton에 out-of-tree 컴파일 패스 플러그인 시스템을 도입한 대규모 PR입니다. 동적 라이브러리 로딩, C API 기반 enumerate/register/add 패턴, 환경변수 기반 자동 로딩이 핵심입니다.
참고 자료
이 글은 AI의 도움을 받아 작성되었으며, 원본 PR의 코드 변경 사항을 기반으로 분석한 내용입니다.
관련 포스트
PR Analysis 의 다른글
- 이전글 [Triton] Gluon의 to_linear_layout에서 TensorMemory 레이아웃 지원
- 현재글 : [triton] Out-of-tree TTIR/TTGIR 패스 플러그인 시스템
- 다음글 [Ray] DefaultCollateFn 병렬화로 Arrow-to-Tensor 변환 가속
댓글