从 __ldcs 谈起一个编译器层面的平台解耦方案1. 从一个 CUDA Intrinsic 说起最近在将llm.c的部分代码移植到 AET 时我遇到了一个与 CUDA Intrinsic 有关的问题。例如layernorm_forward_kernel3中使用了__ldcs(x c)。这是 CUDA 提供的一个设备端函数可以使用 PTX 指令ld.global.cs从全局内存加载数据。CUDA 的相关实现大致如下long__ldcs(constlong*ptr){unsignedlongret;asmvolatile(ld.global.cs.s64 %0, [%1];:l(ret):__LDG_PTR(ptr));return(long)ret;}直接使用这个函数没有问题但用户需要依赖 CUDA 提供的声明和实现。对于只面向 CUDA 的程序这种方式很自然。但 AET 的设计目标不只是支持 CUDA而是通过 MTCS 支持多个计算平台。如果平台无关的代码直接调用 CUDA 专有的 Intrinsic那么代码本身就会与特定平台绑定。我在想能不能把 CUDA 的__ldcs换成 AET 的Target::ldcs让用户只写统一接口具体实现交给编译器去拼装这样移植代码时就无需四处改动切换平台也只需要重新编译而不必改源码。2. AET 与 MTCS 的关系在介绍具体设计之前需要先说明 AET 与 MTCS 的关系。AET 是基于 GCC 的语言扩展与编译器体系。它在 C 的基础上扩展了面向对象、泛型等编程能力同时也为异构计算提供了相应的编译机制。MTCS 是 AET 编译器体系中的多平台编译基础设施职责是支持多个计算平台。它负责处理不同目标平台相关的编译工作例如面向不同芯片的代码生成。因此面向对象和泛型属于语言层面的编程能力而 MTCS 属于编译器层面的多平台支持机制。它们处于不同的层次但共同构成 AET 的整体设计。这也决定了 AET Intrinsic 机制的设计方向既要让编译器能够识别和处理 Intrinsic也要让平台相关的实现与平台无关的调用接口分离。3. 在 AET 中增加第四类函数Intrinsic在此前的设计中我将函数分为三类普通函数由用户或库提供实现。Builtin由编译器识别的内建函数。Internal Function编译器内部使用的函数形式参与中间表示及相关优化和代码生成。现在我为 AET 增加了第四类Intrinsic。这里的重点不只是让编译器识别一种特殊函数而是建立一套平台无关接口与平台相关实现之间的关联机制。在 AET 中我使用两个标记micMTCS Intrinsic Class用于标记平台无关的 Intrinsic 接口类。miiMTCS Intrinsic Implement用于标记平台相关的 Intrinsic 实现类。这两个标记承担不同的职责。mic描述用户调用的接口mii描述特定平台上的具体实现。在继续之前需要解释一下 AET 语法中的一个记号$。AET 是基于 GCC 的 C 语言扩展。为了引入public$、private$、class$、final$等新关键字同时避免与 C 语言中已有的合法标识符产生冲突我在这些关键字后面加上了$。$是关键字的一部分用于将 AET 扩展语法与原有 C 语法明确区分开来让编译器能够准确识别这些新增的语言特性。这与 Objective-C 使用来标记interface、implementation等扩展语法的思路类似通过一个特殊字符明确标识语言扩展既便于编译器识别也尽量少地影响原有语言的写法。因此在下面的示例中public$、class$、static$都要当作整体来看待不能去掉$。4. 平台相关的实现mii以 CUDA 为例可以定义一个实现类public$final$__attribute__((mii(cuda,sm_35)))class$ CudaMemIntrinsic{public$:public$static$ __device__floatldcs(constfloat*ptr){floatret;asmvolatile(ld.global.cs.f32 %0, [%1];:f(ret):l(ptr));returnret;}};这里的mii(cuda, sm_35)表示该类提供面向 CUDA 平台及指定架构条件的实现。示例中的内联汇编直接使用 PTX 指令。需要说明的是这里的sm_35只是一个示例参数并不代表该机制要求所有 CUDA 实现都必须精确到特定架构。实际匹配时MTCS 会优先精确匹配目标平台和架构如果找不到对应架构的实现可能回退到更通用的平台实现。如果完全没有匹配项编译器会报错提示缺少当前目标平台所需的 Intrinsic 实现。同样示例中的内联汇编省略了部分地址空间限定细节。真正的 PTX 可能还需要显式添加.global等限定这里只为展示机制的组织方式。实现类由 AET 编译器识别和关联不要求普通用户直接调用这个类。将来如果其他平台需要提供对应实现也可以按照相同的设计思路增加相应的实现类。这里要强调的是不同平台不一定有完全相同的底层指令。统一的是上层接口而不是要求所有平台使用相同的汇编指令。5. 平台无关的接口mic接下来定义一个供用户调用的接口类public$__attribute__((mic()))class$ Target{#pragmamtcs intrinsicpublic$static$ __device__floatldcs(constfloat*ptr);};这里的Target只是示例名称并不是 AET 要求用户必须使用的类名。mic()标记这个类是平台无关的 Intrinsic 接口类。用户通过它声明希望使用的操作而不需要在接口中写入 CUDA 专用的 PTX 汇编。这里还有一个重要的细节#pragma mtcs intrinsic。这个 pragma 的作用是声明其后相应的函数属于 Intrinsic 函数。编译器识别到这个标记后才会针对这些函数执行平台实现查找并将接口与当前目标平台适用的实现关联起来。它不是简单的代码注释也不是用来声明某个目标平台的指令。平台和架构的信息由实现侧的mii标记提供#pragma mtcs intrinsic则告诉编译器接口中的函数需要走 Intrinsic 的处理机制。在当前的实现中该 pragma 必须紧邻函数声明放置作用于它后面的单个函数声明。如果放在类定义之前则当前行为是对该类第一个匹配的静态方法生效后续版本可能会扩展为整个类作用域。这样用户代码就可以直接写floatvalueTarget::ldcs(ptr);具体调用语法以 AET 当前的语言实现为准。关键在于用户代码只依赖统一接口不必直接依赖CudaMemIntrinsic也不必在调用处写出 CUDA 专用的汇编。6. 编译器在其中承担什么职责整个过程可以概括为用户代码 | v 平台无关的 Intrinsic 接口mic | v AET 编译器识别 Intrinsic 函数 | v 根据当前目标平台查找匹配的实现mii | v 生成目标平台对应的代码例如当前编译目标为 CUDA 时编译器查找适用的 CUDA 实现如果未来增加了其他目标平台的实现编译器则可以根据当前目标选择对应实现。这并不意味着编译器必须依次遍历所有平台。查找范围由当前编译目标以及实现所声明的平台和架构条件决定。如果找不到匹配的实现编译器应当报告错误而不是静默地选择不适用的实现。因此这套机制并不是简单地给一个普通函数加上特殊名称而是让编译器能够在接口声明、目标平台和具体实现之间建立明确的联系。7. 相关的已有工作在设计这套机制的过程中我注意到 HIP、SYCL、LLVM Intrinsic 等项目也在不同方向上处理类似问题。由于我对这些项目的了解有限这里不做深入对比。感兴趣的读者可以自行查阅官方文档。8. 总结与下一步这次设计来自一个具体问题移植llm.c时直接调用 CUDA 的__ldcs会使代码依赖 CUDA 专用接口。通过mic定义平台无关的接口通过mii标记平台相关的实现再由#pragma mtcs intrinsic告诉编译器哪些函数需要执行 Intrinsic 实现查找AET 可以把用户调用与目标平台实现分离。这套机制仍在随着 AET 和 MTCS 的开发逐步完善。它的核心思路是让用户面向统一接口编程让编译器负责关联目标平台的实现。接下来的计划包括将llm.c完整移植到 AET作为多目标编译的端到端测试案例。这也是 AET 多平台编译方向的一部分麻烦留给编译器方便留给用户。真正实现**“编写一次多芯运行”**