[triton] triton-ext Plugin API에 문자열 인자 지원 추가

PR 링크: triton-lang/triton#9691 상태: Merged | 변경: +73 / -25

들어가며

Triton은 플러그인 시스템을 통해 커스텀 컴파일러 패스를 추가할 수 있습니다. 그러나 기존 addPass API는 패스 이름만 받을 뿐 추가 인자를 전달할 수 없었습니다. 실제 Triton 내부에서는 num_warps, capability 등의 파라미터를 패스에 전달하는 경우가 많은데, 플러그인에서는 이것이 불가능했습니다.

핵심 코드 분석

Before:

using AddPassCType = TritonPluginResult (*)(mlir::PassManager *, const char *);

TRITON_PLUGIN_API
tritonAddPluginPass(mlir::PassManager *pm, const char *passName) {
  passMap[passNameStr](pm);
  return TP_SUCCESS;
}

After:

using AddPassCType = TritonPluginResult (*)(mlir::PassManager *, const char *,
                                            const std::vector<std::string> &);

TRITON_PLUGIN_API
tritonAddPluginPass(mlir::PassManager *pm, const char *passName,
                    const std::vector<std::string> &args) {
  passMap[passNameStr](pm, args);
  return TP_SUCCESS;
}

플러그인 측에서는 문자열 인자를 파싱하여 LLVM cl::opt 스타일의 옵션으로 변환합니다:

static void addTritonPluginPass(mlir::PassManager *pm,
                                const std::vector<std::string> &args) {
  if (args.empty()) {
    pm->addPass(createTritonGPUMLIRPlugin());
    return;
  }
  TritonGPUMLIRPluginOptions opts;
  opts.num_warps = std::atoi(args[0].c_str());
  pm->addPass(createTritonGPUMLIRPlugin(opts));
}

왜 이게 좋은가

이 변경은 Triton 플러그인 생태계의 표현력을 크게 확장합니다. registerPasses API에는 cl::opt를 사용하라고 안내하면서, addPass API에는 문자열 리스트라는 유연한 인터페이스를 제공하여 Python 측에서 간단하게 인자를 전달할 수 있습니다. Python 바인딩도 py::arg("args") = std::vector<std::string>()로 기본값을 설정하여 기존 코드와의 하위 호환성을 유지합니다.

정리

• addPass API에 std::vector<std::string> 인자 추가
• Python 바인딩에서 기본값으로 하위 호환성 유지
• 예제 플러그인에 num_warps 파라미터 지원 데모 포함
• 내부 패스의 convert_to_ttgpuir, f32_dot_tc 등과 유사한 패턴 지원

참고 자료

• triton-lang/triton#9691

이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.