添加 C++ 封装的算子#

要添加一个 C++ 封装的算子，你需要首先在安装 FlagGems 时启用 C++ 扩展能力特性。 请参阅FlagGems 安装文档， 了解如何在安装 flag_gems 时启用 C++ 扩展的详细步骤。

1. 编写封装层#

按照如下步骤添加一个新的 C++ 封装的算子：

• 在 include/flag_gems/operators.h 文件中为算子添加函数原型声明；

• 在 lib/<op_name>.cpp 文件中添加算子的函数实现；

• 修改 lib/CMakeLists.txt 文件，包含新的算子；

• 在 src/flag_gems/csrc/cstub.cpp 文件中为算子添加 Python 绑定逻辑；

• 在 triton_src/ 下面为算子添加 triton_jit 函数；

  提示

  目前我们使用一个专门的目录来存放 triton_jit 函数。 将来我们会复用 flag_gems 目录下 Python 代码中的 triton_jit 函数。

2. 编写测试用例#

FlagGems 使用 ctest 和 googletest 来执行 C++ 代码的单元测试。 在完成 C++ 封装的算子实现之后，你需要为其添加对应的 C++ 测试用例。 你的测试用例应该添加到 ctests/test_triton_<xxx>.cpp 文件中， 并且在 ctests/CMakeLists.txt 中列出测试用例文件。 最后，构建你的测试代码并使用 C++ 测试文档中所给方法来执行测试。

3. 运行 C++ 测试用例#

如果你在构造 FlagGems 时将 cmake 选项 FLAGGEMS_BUILD_CTESTS 设置为 ON， 进而启用了 C++ 扩展，那么你就可以使用下面的命令在 FlagGems/build/cpython-3xx 目录下运行 ctest。

ctest .

上面这条命令会运行 ctests/ 目录下的所有测试用例文件。 你也可以使用下面的命令来运行一个指定的测试用例，同时启用日志信息输出：

ctest -V —R <regex>

其中

• -R <regex>：运行名字与所给的正则表达式匹配的所有测试用例；
• -V：启用详尽（verbose）输出模式，打印每个测试的详细输出， 测试用例写入到标准输出（stdout）或标准错误输出（stderr）的所有消息。

例如：

TORCH_CPP_LOG_LEVEL=INFO ctest -V -R test_triton_reduction

我们也使用 PyTorch 的 Aten 日志机制，所以你需要设置环境变量 TORCH_CPP_LOG_LEVEL=INFO 才能获得 libtorch_example 中的更多日志信息。

4. 为你的代码提交 PR#

当所有一切都工作正常时，该到提交拉去请求（Pull Request，PR）的时候了。 在提交 PR 时，我们希望你在针对算子所做的工作提供简要描述之外， 能够在 PR 描述中包含端到端的性能测试数据以方便评审。