主要观点:自定义 PTX 指令可利用 NVIDIA GPU 新硬件特性,提升内核性能,本文旨在指导在 Mojo 中调用自定义 PTX 指令并将其纳入 Mojo 的开源std
。
关键信息:
- 展示了在 Mojo 中计算元素级
tanh
的内核程序,通过assert_almost_equal
断言正确性。 - 分析 Mojo
std
中对tanh
的处理,在 NVIDIA GPU 上对于bfloat16
会先转换为float32
再使用 PTX 指令,现倾向使用专门针对半精度浮点数的tanh.approx.bf16
指令。 - 给出了围绕
tanh.approx.bf16
指令的包装函数tanh_bfloat16
。 - 通过
--emit
标志将 Mojo 代码编译为LLVM IR
和asm
代码,对比了使用原始tanh
和自定义tanh_bfloat16
指令编译后的内核代码,前者需先转换再执行指令,后者可直接输入bfloat16
。
重要细节: - 定义了项目并给出相关代码,如
magic init intro-mojo-ptx --format mojoproject
等。 - 详细说明了
tanh
函数的源代码及相关参数。 - 介绍了获取编译信息的函数
_compile_code_asm
及编译后的代码示例,包括原始和自定义指令编译后的输出对比。
**粗体** _斜体_ [链接](http://example.com) `代码` - 列表 > 引用
。你还可以使用@
来通知其他用户。