CuTe 张量算法全解copy / copy_if / gemm / axpby / fill / clear 的接口设计与实现剖析【免费下载链接】cutlassCUDA Templates and Python DSLs for High-Performance Linear Algebra项目地址: https://gitcode.com/GitHub_Trending/cu/cutlassCuTe 是 CUTLASS 提供的张量计算模板库其核心抽象Tensor把数据类型、存储位置global memory / shared memory / registers以及形状与步长Layout在编译期一并编码到类型系统中。本文以 media/docs/cpp/cute/04_algorithms.md 为主线系统讲解 CuTe 在Tensor上实现的copy、copy_if、gemm、axpby、fill、clear等数值算法的接口形态、类型驱动分派type-based dispatch机制、并行性与同步语义并结合 include/cute/algorithm/ 目录下的真实实现与 examples/cute/tutorial/ 中的教程示例给出源码级印证。读完本文你将掌握这些算法在何时选择默认实现、何时通过Copy_Atom/MMA_Atom覆盖默认实现以及它们在真实 GEMM kernel 中如何与谓词predication、异步拷贝和流水线配合。概述算法目录与整体设计CuTe 的所有张量算法实现集中在仓库的 include/cute/algorithm/ 目录下每个算法对应一个或一组头文件头文件提供的算法一句话说明copy.hppcopy、copy_if拷贝、带谓词的拷贝覆盖标量拷贝到cp.async、TMA 等硬件指令gemm.hppgemm张量级 GEMM按模式mode数量分派到 FMA/MMAaxpby.hppaxpby广义 AXPYy αx βyfill.hppfill用标量填充张量clear.hppclear将张量清零fill(t, T{})的特例cooperative_copy.hpp协作式拷贝面向多线程/多 CTA 协作的拷贝cooperative_gemm.hpp协作式 GEMM面向多线程/多 CTA 协作的 GEMMtensor_reduce.hppreduce、batch_reduce、logical_reduce归约类算法tensor_algorithms.hppfor_each、transform类似 STL 的遍历与变换functional.hpp函数对象plus、multiplies等二元/一元算子prefer.hpp重载优先级标签用于重载决议的preferN标签prefetch.hpp预取数据预取相关辅助tuple_algorithms.hpp元组算法面向元组/复合张量的算法这些算法的共同设计哲学是由参数类型驱动分派。一个Tensor的 Engine数据存储与元素类型和 Layout形状与步长都是类型的一部分因此算法可以在编译期根据参数类型选择最合适的实现——从最朴素的逐元素循环到同步/异步的硬件拷贝指令再到各类 MMA矩阵乘加指令。copy类型驱动分派到硬件拷贝指令接口与两大重载copy把源Tensor的元素拷贝到目标Tensor中。其各种重载定义在 include/cute/algorithm/copy.hpp。该算法有两大主重载。第一个重载只接收源张量和目标张量实现由两个Tensor参数的类型决定template class SrcEngine, class SrcLayout, class DstEngine, class DstLayout CUTE_HOST_DEVICE void copy(TensorSrcEngine, SrcLayout const src, TensorDstEngine, DstLayout dst);第二个重载在前者基础上多一个Copy_Atom拷贝原子参数让调用者显式覆盖默认实现template class... CopyArgs, class SrcEngine, class SrcLayout, class DstEngine, class DstLayout CUTE_HOST_DEVICE void copy(Copy_AtomCopyArgs... const copy_atom, TensorSrcEngine, SrcLayout const src, TensorDstEngine, DstLayout dst);其中Copy_Atom的定义位于 include/cute/atom/copy_atom.hpp它把一次具体硬件拷贝操作如UniversalCopy、SM80_CP_ASYNC_CACHEALWAYS、SM90_BULK_COPY_G2S等见 include/cute/arch/copy.hpp 与 include/cute/arch/copy_sm80.hpp封装为可携带内部类型与可选参数的原子对象。Copy_Atom之上还有TiledCopy用于把拷贝铺开到多线程/多指令上详见 include/cute/atom/copy_atom.hpp。从实现来看include/cute/algorithm/copy.hpp带Copy_Atom的copy会先检查源/目标 Layout 的 rank 是否一致static_assert(SrcLayout::rank DstLayout::rank, CopyAtom rank-mismatch.)若 rank 为 1 则直接把src/dst交给copy_atom.call(src, dst)否则用group_modes1,R把第一个模式V 模式与其余模式分开对每个切片调用一次copy_atom.call。对于静态形状的情况还会执行AutoFilter利用nullspace/zipped_divide找出目标布局的零空间nullspace自动过滤掉那些不产生实际访问的冗余坐标在保证不产生竞态Ambiguous race-condition的前提下减少无谓的拷贝调用。默认实现如何自动选择不带Copy_Atom的两参数copy走自动向量化 自动过滤路径include/cute/algorithm/copy.hpp分三档静态 Layout例如寄存器张量假定指针满足 128 位对齐使用AutoFilter(AutoVectorizingCopyWithAssumedAlignment128{})静态形状、动态布局可以进行过滤但不假定对齐降级为AutoVectorizingCopyWithAssumedAlignment8完全动态不假定对齐同样使用 8 位基准的自动向量化。自动向量化逻辑include/cute/algorithm/copy.hpp会计算max_common_vector(src, dst)找出源与目标可以合并的最大公共向量宽度再结合max_alignment与MaxVecBits取最大公约数得到vec_bits。当vec_bits是 8 的倍数且大于元素位宽时就通过recastVecType把张量重铸为更宽的向量类型例如把四个 32 位元素合成一个 128 位元素从而把四条ld.global.b32合并成一条ld.global.b128。这正是本文后面通用实现优化清单第 2 条在真实代码中的落地。此外include/cute/algorithm/copy.hpp 还提供了copy_aligned它假定对齐到 128 位并同样执行自动过滤适合调用者确定指针与布局对齐的场景。所有拷贝重载都接受右值目标张量Accept mutable temporaries便于链式编写代码。并行性与同步语义取决于参数类型无论走默认实现还是Copy_Atom选择的自定义实现copy都可能不使用任何并行性也可能使用全部可用的并行度并伴随多种同步语义——具体行为完全取决于参数类型需要使用者基于目标架构的知识自行判断开发者通常为每种 GPU 架构编写定制的优化 kernel。copy可以是每线程串行的也可以是跨一组线程如一个 thread block 或 cluster并行的。若copy是并行的那么在使用其结果之前参与线程组可能需要同步。例如参与线程构成一个 thread block 时必须先调用__syncthreads()或 Cooperative Groups 的等价物。copy可能使用异步拷贝指令如cp.async及其 C 接口memcpy_async。这种情况下使用者必须执行与该底层实现匹配的额外同步才能安全地使用拷贝结果。examples/cute/tutorial/ 中的 CuTe GEMM 教程示例展示了一种同步方法更优化的 GEMM 实现则用流水线pipelining技术把异步copy与其他有用工作重叠起来。从源码角度异步拷贝路径可以看 include/cute/algorithm/copy.hppcopy(AutoCopyAsync const cpy, ...)只是把谓词恒为真的copy_ifcopy_if(cpy, constant_fntrue_type{}, src, dst)转发出去。而AutoCopyAsync内部include/cute/algorithm/copy.hpp在CUTE_ARCH_CP_ASYNC_SM80_ENABLED使能且满足gmem→smem、源/目标元素大小相同时会选择SM80_CP_ASYNC_CACHEGLOBAL或SM80_CP_ASYNC_CACHEALWAYS16 字节元素走 CACHEGLOBAL4/8/16 字节元素走 CACHEALWAYS否则回退到UniversalCopy。Hopper 及以后的架构还支持 TMA 批量拷贝当CUTE_COPY_ATOM_TMA_SM90_ENABLED开启时copy对SM90_BULK_COPY_AUTO有专门特化include/cute/algorithm/copy.hpp它只允许 gmem↔smem 之间搬运、要求至少 128 位宽度并通过max_common_layout找到公共子张量后自动构造SM90_BULK_COPY_G2S/SM90_BULK_COPY_S2G原子。TiledCopy与ThrCopy也都有对应的copy/copy_if重载TiledCopy会退化为其底层Copy_Atominclude/cute/algorithm/copy.hpp而ThrCopy线程切片的copy/copy_if被显式 deleteinclude/cute/algorithm/copy.hpp——也就是说你不应该把已经按线程切分好的ThrCopy直接交给copy而是要先用get_slice(thread_idx)之类的接口得到本线程的Copy_Atom视图。未识别的拷贝策略会触发static_assert(dependent_falseCopyPolicy, Unrecognized CopyPolicy.)编译期报错。一个通用 copy 实现下面这段代码展示了任意两个Tensor之间最朴素的通用copy实现template class TA, class ALayout, class TB, class BLayout CUTE_HOST_DEVICE void copy(TensorTA, ALayout const src, // Any logical shape TensorTB, BLayout dst) // Any logical shape { for (int i 0; i size(dst); i) { dst(i) src(i); } }该实现把两个张量都当作一维逻辑坐标遍历按逻辑的列主序column-major顺序访问。在此基础上一些合理的、与架构无关的优化方向包括若两个张量的内存空间已知且存在优化的访问指令如cp.async则分派到定制指令若两个张量拥有静态布局且可证明元素向量化是合法的——例如四条ld.global.b32可以合并成一条ld.global.b128——则对源/目标张量做向量化若可能校验将要使用的拷贝指令对源/目标张量是否合适。CuTe 的优化实现全部做到了以上三点分别对应AutoCopyAsync、AutoVectorizingCopyWithAssumedAlignment、以及各种static_assert与形状/对齐校验。copy_if带谓词的拷贝copy_if与copy位于同一头文件 include/cute/algorithm/copy.hpp。它与copy一样接收源、目标Tensor但额外接收一个与源/目标同形状的谓词张量predicationTensor源张量的元素只有在谓词张量对应元素非零true时才会被拷贝。朴素的标量版实现非常直白include/cute/algorithm/copy.hpptemplate class PrdTensor, class SrcEngine, class SrcLayout, class DstEngine, class DstLayout CUTE_HOST_DEVICE void copy_if(PrdTensor const pred, TensorSrcEngine, SrcLayout const src, TensorDstEngine, DstLayout dst) { using SrcType typename SrcEngine::value_type; using DstType typename DstEngine::value_type; CUTE_UNROLL for (int i 0; i size(dst); i) { if (pred(i)) { dst(i) static_castDstType(static_castSrcType(src(i))); } } }除标量版外还有三个变体Copy_Atom版include/cute/algorithm/copy.hpp要求源/目标/谓词三者的 rank 一致CopyAtom rank-mismatch.rank 为 1 时直接copy_atom.call(prd, src, dst)否则对 V 模式之外的每个切片循环调用。旧的rank-1 谓词重载pred比src少一维已被标记为 deprecated官方建议改用 bool 张量或 transform 张量作为谓词include/cute/algorithm/copy.hpp。AutoCopyAsync版include/cute/algorithm/copy.hpp在cp.async可用时选择对应的SM80_CP_ASYNC_*原子否则用UniversalCopy逐元素按谓词执行异步拷贝。TiledCopy版退化为底层Copy_Atom的copy_if。关于copy_if的使用动机与详细方法请参考教程的 谓词predication章节。核心思想是当矩阵无法被 tile 整除时例如把 41×55 的矩阵切成 4×8 的 tile必然有剩余CuTe 采用向上取整 谓词屏蔽的策略——logical_divide会把形状圆整到 128×8 这样的完整 tile再构造一个与数据张量同形、记录了原始坐标的恒等/坐标张量将其坐标与原始形状比较生成谓词最后用copy_if只拷贝合法元素。典型用法// Copy a k_tile from global memory to shared memory copy_if(tApA, tAgA(_,_,k_tile), tAsA); copy_if(tBpB, tBgB(_,_,k_tile), tBsB);这种先取整、再屏蔽的方式避免了带 warp 分歧的分支与 CUDA 一维并行中先判越界再做功的惯用法一脉相承。gemm按模式数量分派的张量级 GEMMgemm计算什么gemm算法接收 A、B、C 三个Tensor其具体行为取决于各Tensor的模式mode数量。CuTe 用字母表达这些模式V表示向量即由相互独立元素构成的一个模式M、N分别表示 BLAS GEMM 例程中结果矩阵 C 的行数与列数K表示 GEMM 的归约模式即 GEMM 沿着求和的模式。CuTe 用(...) x (...) (...)记号描述输入 A、B 与输出 C 的模式左侧两个括号依次是 A 和 B右侧是 C。gemm共有五种形态该表同时记录在 include/cute/algorithm/gemm.hpp 的头部注释中#形态语义分派策略1(V) x (V) (V)向量逐元素乘积C_v A_v B_v分派到 FMA 或 MMA2(M) x (N) (M,N)向量外积C_mn A_m B_n以 V1 分派到形态 43(M,K) x (N,K) (M,N)矩阵乘积C_mn A_mk B_nk对每个 K 分派到形态 24(V,M) x (V,N) (V,M,N)向量批量外积C_vmn A_vm B_vn优化寄存器复用对每个 (M,N) 分派到形态 15(V,M,K) x (V,N,K) (V,M,N)矩阵批量乘积C_vmn A_vmk B_vnk对每个 K 分派到形态 4关于模式顺序的约定详见 GEMM 教程K 总是最右最外层V 总是最左最内层。分派到优化实现与copy相同CuTe 的gemm实现也依据Tensor参数的类型分派到经过优化的实现gemm同样接受一个可选的MMA_Atom参数允许调用者覆盖 CuTe 依据Tensor类型默认选择的FMA指令。MMA_Atom的定义在 include/cute/atom/mma_atom.hpp其上层还有TiledMMAinclude/cute/atom/mma_atom.hpp。默认 MMA 为UniversalFMA见 include/cute/arch/mma.hpp它是对标量 FMA 的封装。从实现看include/cute/algorithm/gemm.hpp 把gemm分成若干层次三参到四参的转换include/cute/algorithm/gemm.hppgemm(A, B, C)等价于gemm(C, A, B, C)即默认把 C 同时当作累加输入与输出D A*B C。默认 MMA 注入include/cute/algorithm/gemm.hpp四参gemm(D, A, B, C)依据四个张量的value_type构造MMA_AtomUniversalFMA...并转发。线程本地寄存器 GEMMinclude/cute/algorithm/gemm.hpp通过__CUTE_REQUIRES约束 rank 与内存空间is_rmem为形态 1~5 分别实现形态 1 直接mma.call(D, A, B, C)形态 2 用append2把一维布局补成二维后递归到形态 4形态 3 用prepend3补出 V1 模式后递归到形态 5形态 4 是实现核心之一内部针对元素宽度做了多条serpentine蛇形遍历特化当 A、B 的 V 模式各占 64 位时用行主序蛇形各占 32 位时用kinked serpentine折角蛇形64 位32 位混合时按位宽决定谁在外层循环、谁走蛇形其余情况回退到列主序蛇形。这些路径都围绕一个目标——最大化寄存器复用形态 5 对每个 K 切片递归到形态 4。线程本地共享内存 GEMMinclude/cute/algorithm/gemm.hpp当 A/B 在 shared memoryis_smem而 D/C 在寄存器时先用MMA_AtomMMA::make_fragment_A/B把 smem 片段拷入寄存器片段再逐 K 做线程级寄存器 GEMM。这也印证了教程中先用copy把数据搬进寄存器片段、再调用线程级gemm的写法。gemm的实现中遍布CUTE_STATIC_ASSERT_V形状校验如AM CM、BN CN、AK BK等在编译期保证形态间的形状一致性。关于MMA_Atom与各架构下gemm特化的更多细节请参考教程的 MMA 章节。axpby广义 AXPYaxpby定义在 include/cute/algorithm/axpby.hpp它计算$$y \alpha x \beta y$$其中 α、β 是标量x、y 是Tensor。名字取 Alpha times X Plus Beta times Y是 BLAS 经典 AXPYAlpha times X Plus Y的推广——AXPY 相当于 β1 的特例。实现细节include/cute/algorithm/axpby.hpp值得注意第三个模板参数PrdTensor constant_fntrue_type允许传入可选的谓词张量默认恒真因此axpby天然支持谓词化的 epilogue——在 谓词教程 的 GEMM epilogue 示例中正是用tCgC(i) alpha * tCrC(i) beta * tCgC(i)这种形式完成累加输出。编译期会检查 β 是否为零支持复数Beta分别比较实部虚部若 β0循环体退化为y(i) alpha * x(i)省去一次乘法与加法这是为 GEMM 中第一个 K-tile 直接写结果这类场景做的优化。fill与clearfill定义在 include/cute/algorithm/fill.hpp把目标Tensor的每个元素覆盖为给定标量。其内部通过prefer1/prefer0标签做重载决议include/cute/algorithm/fill.hpp若底层的 Engine 支持直接fill(tensor.data(), value)例如按连续内存批量填充就优先走该路径否则回退到逐元素循环// Default implementation template class Engine, class Layout, class T CUTE_HOST_DEVICE void fill(TensorEngine, Layout tensor, T const value, prefer0) { CUTE_UNROLL for (int i 0; i size(tensor); i) { tensor(i) value; } }clear定义在 include/cute/algorithm/clear.hpp是把fill的应用特化为填零fill(tensor, T{})其中T是张量的value_type。这两个算法通常用于初始化累加器寄存器片段或清空共享内存缓冲是所有 GEMM kernel 的常见开场动作。其他相关算法除上述核心算法外include/cute/algorithm/ 目录还包含归约tensor_reduce.hppreduce用二元算子BinaryOp默认cute::plus把张量归约成标量batch_reduce沿(RedMode, BatchMode)布局的 Red 模式归约并逐 batch 写入dstlogical_reduce则按target_profile指定的模式_表示保留、整数表示归约对布局切块后做归约。遍历与变换tensor_algorithms.hppfor_each、transform与 STL 同名算法语义类似cute::lazy::transform则返回一个惰性视图张量。协作式算法cooperative_copy.hpp、cooperative_gemm.hpp面向多线程/多 CTA 协作场景的高层封装。函数对象与元组算法functional.hpp、tuple_algorithms.hpp提供plus、multiplies等算子以及面向复合/元组张量的操作。结合教程示例理解实战用法examples/cute/tutorial/ 目录下的示例把上述算法串进了完整流程sgemm_1.cu 与 sgemm_2.cu 是入门级 SGEMM演示copygmem→smem、gemm线程级 MMA与axpby/fillepilogue的配合sgemm_sm70.cu 面向 Volta 架构的 wgmma 风格 MMAsgemm_sm80.cu 展示 Ampere 上cp.async异步拷贝与多级流水线tiled_copy.cu 与 tiled_copy_if.cu 分别演示TiledCopy与谓词化copy_if的用法是理解本文copy/copy_if分派机制的最佳配套阅读hopper/ 与 blackwell/ 子目录则包含面向 HopperTMA、集群与 Blackwell 的进阶示例。以 SGEMM 的典型主循环为例其模式固定为// gmem - smem用 copy可能被 TiledCopy cp.async 优化 copy(gA(_, _, k), sA); copy(gB(_, _, k), sB); __syncthreads(); // 若使用 cp.async需要对应的异步同步原语 // smem - 寄存器片段 - 线程级 gemm gemm(mma, tCrC, tCrA, tCrB, tCrC); // 形态 5(V,M,K) x (V,N,K) (V,M,N) // epilogueaxpby可结合谓词 tCcC 屏蔽越界 tCgC(i) alpha * tCrC(i) beta * tCgC(i);这里的每一步——搬运、计算、写回——都对应本文讲解的一个算法而它们之间的同步点正是并行性与同步语义取决于参数类型一节的现实体现。总结CuTe 的算法层是类型驱动分派设计的集中体现copy/copy_if依据源/目标Tensor的类型在标量循环、向量化拷贝、cp.async与 TMA 批量拷贝之间自动选择gemm依据模式数量在五种形态之间递归分派并最终落到 FMA/MMA 原子axpby/fill/clear则提供 epilogue 与初始化所需的标量级操作。理解这些算法的关键是记住参数类型Engine Layout即实现选择而并行性、同步语义与是否异步完全由底层所选的具体指令决定。若需在真实 kernel 中覆盖默认行为则显式传入Copy_Atom或MMA_Atom即可配合 谓词教程 的坐标张量方法可以写出同时兼顾性能与边界正确性的通用 GEMM。【免费下载链接】cutlassCUDA Templates and Python DSLs for High-Performance Linear Algebra项目地址: https://gitcode.com/GitHub_Trending/cu/cutlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考