본문으로 건너뛰기

[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 의 다른글