[8 GPU and Specialized Hardware - part2]sch.annotate的作用和工作原理

[hbsun2113]

在第八节课程中，提到可以使用sch.annotate(i, "pragma_import_llvm", tmm_kernel())来替代tvm.tir.TensorIntrin.register("tmm16", tmm16_desc, tmm16_impl)。我的理解是对于block_mm内部的计算逻辑使用tmm_kernel而不是tmm16_impl。但是通过sch.mod.show()我发现block_mm的计算逻辑仍然使用了T.call_extern("tmm16"...。请问这里sch.annotate对实际的运行会产生什么影响？

文档没有说明sch.annotate对于运行的影响和原理：

[Hzfengsy]

可能你对这部分有一点误解。sch.annotate 并不能“替代” tvm.tir.TensorIntrin.register("tmm16", tmm16_desc, tmm16_impl)。这两者是需要同时作用才能工作的。
这里的sch.annotate是给定程序一个tmm16的实现（这个函数在是自己实现的，在编译器中找不到所以需要给定程序），配合Tensorize一起使用

[hbsun2113]

明白了。还有一个问题，我没有看到tmm16和tmm_kernel之间的映射，请问sch.annotate是如何工作的，使得tvm可以将tmm16的实现替换为tmm_kernel？
这里没有出现tmm16：sch.annotate(i, "pragma_import_llvm", tmm_kernel())。

[Hzfengsy]

tmm_kernel是一个Python函数，返回的是tmm16的llvm IR代码。

[hbsun2113]

请问哪里有指明tmm_kernel返回的是tmm16的代码呢？

[Hzfengsy]

tmm_kernel的实现 extern "C" int tmm16(...)

[hbsun2113]

明白了，感谢！