feat(triton): add JIT backend with add operator - #800
Conversation
add operator
d2c1db0 to
57b7c00
Compare
57b7c00 to
f615a7f
Compare
voltjia
left a comment
There was a problem hiding this comment.
代码需要遵循 CONTRIBUTING.md,其中 C++ 主要是遵循 Google C++ Style Guide。目前的 naming 之类的好像没有 follow 这些 convention,需要修改一下。
|
|
||
| if(WITH_TRITON) | ||
| target_include_directories(ops PRIVATE | ||
| ${INFINIOPS_TRITON_INCLUDE_DIRS}) |
There was a problem hiding this comment.
这里前缀应该是 INFINI_OPS,目前的原则是参数类的都是需要加下划线这种间隔的,只有和链接库名之类相关的才省去。
| # Ship the JIT compiler and kernel sources so Triton JIT operators | ||
| # can compile kernels at runtime. compile.py uses __file__ to | ||
| # locate ops/ relative to itself; both must live under triton/. |
There was a problem hiding this comment.
可以的话把注释中的代码用 Markdown 语法括一下,比如 `compile.py`、`__file__` 之类的,除了此处以外别的地方也检查一下,包括但不限于注释、error message、assert message 等。目前这个规则没有那么严格,如果有注意不到的也没事,不过能注意到的还是改一下吧。
| #ifndef INFINI_OPS_TRITON_JIT_ADD_H_ | ||
| #define INFINI_OPS_TRITON_JIT_ADD_H_ | ||
|
|
||
| #include <cuda.h> |
There was a problem hiding this comment.
triton 下面引入这种平台相关的头文件不太好,因为这样想跨平台就比较难了,要不然就会变成一堆宏。这个地方可以考虑改成 InfiniRT 调用,这样跨平台就交给了 InfiniRT 了,可以保持一致性,而且我看用到的接口目前的 InfiniRT 应该也都提供了。这里也可以像 native 下面的一些算子,把 Backend 作为模板参数,然后通过特化来提供,但是感觉这样对于 triton 下面的东西反而复杂化了。但是总之这里应该尽量不引入具体平台的东西。
| #ifndef INFINI_OPS_TRITON_JIT_ADD_H_ | ||
| #define INFINI_OPS_TRITON_JIT_ADD_H_ |
There was a problem hiding this comment.
好像不太对,这个得根据路径来,应该是 INFINI_OPS_TRITON_OPS_ADD_JIT_H。
| namespace infini::ops { | ||
|
|
||
| template <> | ||
| class Operator<Add, Device::Type::kNvidia, 7> : public Add { |
There was a problem hiding this comment.
这个 slot 先用 10 吧,目前的设定是从 8 开始是特定后端,比如 8 是 PyTorch,9 是 NineToothed,我记得目前应该是 10 还没用到。
Summary
Addoperator JIT implementation (implementation_index=7) viasrc/triton/ops/add/jit.hscripts/generate_wrappers.py, detects ops usingconfig_tand emitsconfigparameter with inline config-dict parsingsrc/CMakeLists.txt, compiles JIT infra (jit.cc/compiler.cc) and linkspybind11::embedsrc/config.h,Config::set_extension/extension()for attaching opaque extension dataMotivation
Support JIT compilation for Triton operators. Kernel caching, config passing, and autotune orchestration are implemented in C++. The compilation step bridges to
triton.compileviapybind11::embed.Closes N/A
Type of Change
feat— new feature / new operator / new platformfix— bug fixperf— performance improvement (no behavioral change)refactor— code restructuring without behavior changetest— adding or fixing tests onlydocs— documentation onlybuild/ci— build system or CI configurationchore— tooling, formatting, or other non-code changes!in the Conventional Commits prefix or aBREAKING CHANGE:footer)Platforms Affected
WITH_CPU)WITH_NVIDIA)WITH_ILUVATAR)WITH_METAX)WITH_CAMBRICON)WITH_MOORE)WITH_ASCEND)WITH_TORCH)Smoke Test Result
Test Results on Supported Platforms
Full `pytest` output (optional)
Benchmark / Performance Impact
N/A
Notes for Reviewers
feat/triton-backendand depends on them being merged first.generate_wrappers.pyhas additions:_uses_config_extensionfor detecting ops withconfig_t,_generate_triton_jit_config_parserfor emitting inline config-dict parsing, and conditionaluses_configpath in_generate_call.Configclass (src/config.h) now has astd::shared_ptr<Config> extension_member withset_extension/extensionaccessors, used to pass compile configurations from Python bindings to operator implementations.