在 Mojo 中使用 PTX 指令

主要观点:自定义 PTX 指令可利用 NVIDIA GPU 新硬件特性,提升内核性能,本文旨在指导在 Mojo 中调用自定义 PTX 指令并将其纳入 Mojo 的开源std
关键信息

  • 展示了在 Mojo 中计算元素级tanh的内核程序,通过assert_almost_equal断言正确性。
  • 分析 Mojostd中对tanh的处理,在 NVIDIA GPU 上对于bfloat16会先转换为float32再使用 PTX 指令,现倾向使用专门针对半精度浮点数的tanh.approx.bf16指令。
  • 给出了围绕tanh.approx.bf16指令的包装函数tanh_bfloat16
  • 通过--emit标志将 Mojo 代码编译为LLVM IRasm代码,对比了使用原始tanh和自定义tanh_bfloat16指令编译后的内核代码,前者需先转换再执行指令,后者可直接输入bfloat16
    重要细节
  • 定义了项目并给出相关代码,如magic init intro-mojo-ptx --format mojoproject等。
  • 详细说明了tanh函数的源代码及相关参数。
  • 介绍了获取编译信息的函数_compile_code_asm及编译后的代码示例,包括原始和自定义指令编译后的输出对比。
阅读 12
0 条评论