模型编译深度学习推理引擎【免费下载链接】tvmOpen Machine Learning Compiler Framework项目地址https://gitcode.com/gh_mirrors/tv/tvm点击查看免费下载本文以 docs/tirx/tile_primitives/elementwise.rst 为主线系统讲解 TVM Tirxtile IR extension中 CUDA 后端元素级 tile 原语的完整设计从cast、fill、一元zero、reciprocal、sqrt、exp、exp2、log2、silu、二元add、sub、mul、fdiv、maximum到fma的操作集合到每个操作注册reg/smem两个变体、以操作数存储作用域作为互斥判别器的派发模型。读完本文你将掌握何时走寄存器路径、何时走共享内存路径、两种路径各自的分区生成算法induced vs synthesized、生成的 TIRx IR 与 CUDA 代码形态以及如何通过源码和测试用例验证这些行为。一、总览元素级 tile 原语的操作集与双变体注册模型在 Tirx 的Tx.tile目录中元素级elementwise原语覆盖了深度学习算子中最基础的一类对 tile 中每个元素独立施加同样的计算。当前 CUDA 元素级集合包含类型转换cast填充fill一元操作zero、reciprocal、sqrt、exp、exp2、log2、silu二元操作add、sub、mul、fdiv、maximum融合乘加fma。每个操作在 CUDA 后端注册两个变体——reg与smem优先级均为10。二者的判别完全由**缓冲区操作数的存储作用域storage scope**决定且是互斥的reg变体所有缓冲区操作数都在local作用域寄存器smem变体所有缓冲区操作数都在shared作用域共享内存。标量输入scalar inputs不拥有存储作用域因此在这两个变体中都被允许出现具体哪些操作允许标量取决于该操作的 authoring API见下文。每个操作都由一个OpSpec描述。OpSpec的核心是一套parse逻辑它负责构建目的区destination 源列表source list并携带可选的 dtype 检查check_extras以及逐元素应用的标量表达式compute_scalar。从源码看这一模型在 python/tvm/backend/cuda/tile_primitive/elementwise/register.py 中实现得十分简洁——遍历ALL_OPS中的每个OpSpec分别调用_register_reg与_register_smem注册两个派发变体例如register_dispatch( spec.name, cuda, variantreg, priority10, when[predicate(f{spec.name}_reg, is_reg_ewise(spec))], ) def _dispatch(op: TilePrimitiveCall, sctx: DispatchContext, _specspec) - PrimFunc: return emit_reg(op, _spec, sctx)smem变体与之对称只是variantsmem、谓词换成is_smem_ewise(spec)、发射函数换成emit_smem。该文件注释明确指出这种双变体模型模仿了 copy 原语的 PR-640 设计——scope-pair作用域组合决定派发选择底层算法induced 诱导 vs synthesized 合成随之而来。两个变体的差异可先由下表概览VariantOperandsLoweringreg所有缓冲区操作数均为local分区由本地缓冲区的布局诱导产生操作逐线程应用smem所有缓冲区操作数均为shared合成[outer, threads, vec]分区操作逐向量化元素应用值得注意的限制是全局Tx.tile目录中的其他构造器包括minimum、memset和select目前没有这些 CUDAreg/smem变体。也就是说一次Tx.tile.*调用必须由其所选目标后端注册的变体来承接否则无法派发。二、OpSpec解析、dtype 检查与逐元素标量表达式OpSpec是理解整个元素级派发机制的关键抽象。它把一个 tile 原语调用拆成三个层次parse(op_call) - (Plan, msg)把T.op(dst, src, ...)形式的调用解析为一个PlanPlan记录dst目的 TensorRegion与srcs源列表每个源要么是一个 buffer region要么是一个标量。check_extras(extras, compute_dtype)可选的 dtype 规则校验。例如一元操作的scale/bias标量必须与 compute dtype 一致。compute_scalar(src_vals, extras, dst_dtype)逐元素应用的标量表达式即真正施加在每个元素上的数学运算。以 python/tvm/backend/cuda/tile_primitive/elementwise/ops/unary.py 为例一元操作统一采用T.unary(dst, src[, bias, scale])签名bias/scale可选silu为保持历史行为忽略bias/scale。其UNARY_OPS字典清晰展示了 OpSpec 的四元组结构UNARY_OPS: dict[str, OpSpec] { zero: OpSpec(zero, _parse_unary, _compute_zero, _check_unary_extras), fill: OpSpec(fill, _parse_unary, _compute_fill, _check_unary_extras), reciprocal: OpSpec(reciprocal, _parse_unary, _compute_reciprocal, _check_unary_extras), sqrt: OpSpec(sqrt, _parse_unary, _with_bias_scale(T.sqrt), _check_unary_extras), exp: OpSpec(exp, _parse_unary, _with_bias_scale(T.exp), _check_unary_extras), exp2: OpSpec(exp2, _parse_unary, _with_bias_scale(T.exp2), _check_unary_extras), log2: OpSpec(log2, _parse_unary, _with_bias_scale(T.log2), _check_unary_extras), silu: OpSpec(silu, _parse_unary, _compute_silu, _check_unary_extras), }其中_with_bias_scale是一个高阶包装先施加scale乘法、再施加biasbuffer 或常量最后调用裸运算如T.exp_compute_zero恒返回0.0_compute_reciprocal计算1.0 / x_compute_silu计算x / (1 exp(0 - x))。_parse_unary支持源既可以是TensorRegion也可以是 primitive 表达式标量并处理bias为 buffer 与常量两种形态。二元操作在 python/tvm/backend/cuda/tile_primitive/elementwise/ops/binary.py 中实现还包含常量左操作数的交换律commute逻辑_COMMUTATIVE {add, mul, maximum}——当左源是常量时对可交换操作将其移到右侧不可交换操作则拒绝。两个源都是常量时直接报错both inputs are constants。广播extent1 的维度在派发的_broadcast_lift布局层面处理解析器只是按原样记录每个源。三、reg变体local 作用域的 induced 分区路径当一次元素级操作的所有缓冲区操作数都位于local寄存器作用域时派发落入reg变体。其实现位于 python/tvm/backend/cuda/tile_primitive/elementwise/reg.py设计上镜像了 copy 的 reg 路径vec_auto的寄存器路径。3.1 接受条件is_reg_ewise谓词is_reg_ewise(spec)返回一个检查函数它依次验证def check(op_call, sctx): if not sctx.is_target(cuda): return False, non-cuda target if sctx.scope_kind not in (thread, warp, warpgroup, cta): ... ok, reason _all_threads_active(sctx) plan, msg spec.parse(op_call) for br in buffer_regions(plan): if br.source.scope() ! local: # 每个缓冲区操作数都必须是 local return False, foperand scope {br.source.scope()} ! local if br.source.layout is None: ... # spec.check_extrasdtype 规则 # pick_anchor _validate_anchor_layout锚点必须是 TileLayout、非 swizzle # _validate_scope_level_anchorwarp/warpgroup/cta 作用域下锚点的线程轴 # 必须覆盖完整作用域线程链 # NumPy 风格右对齐形状广播检查 # _check_layout_operands_agree各操作数 replica 签名一致将接受条件整理为属性表PropertyRequirementtarget / scope / prioritycudathread/warp/warpgroup/cta全部线程活跃优先级10operands每个缓冲区操作数都在localfill、二元操作与fma允许标量源op父页面列出的任意 CUDA 元素级OpSpec一元 / 二元 /fmaspec.check_extras校验 dtype 组合local-buffer layout锚点布局必须通过校验且其线程轴必须匹配完整的作用域链它诱导分区所有操作数在 thread、per-thread storage、replica 签名上必须一致shapes每个输入区都与目的区 NumPy 风格右对齐广播兼容几个关键校验点的含义值得展开**锚点anchor**由pick_anchor选出通常是目的操作数。_validate_anchor_layout要求锚点布局不是 swizzle、且必须是TileLayout。**作用域级锚点_validate_scope_level_anchor**针对warp/warpgroup/cta作用域经过目标规范化后锚点切片的线程轴只能是该作用域对应的轴laneid/tid_in_wg/tx且按 stride 排序后必须完整铺满该作用域的所有线程stride 依次为 1、e0、e0×e1…直至线程总数不留空隙、不重叠。thread作用域豁免此检查。布局签名一致_check_layout_operands_agree对齐align_operands_to_anchor只能规范 thread 与 local 部分replica 部分不被对齐触碰因此各操作数的 replica 签名轴 key、extent、stride必须逐项一致否则发射结果会是错的。3.2 演示程序文档给出的演示是一个 warp 对32×8的float32局部 tile 求sqrt局部布局为S[(32,8):(1laneid,1)]——即 lanei拥有第i行from tvm.tirx.layout import S, TileLayout, laneid r_layout TileLayout(S[(32, 8) : (1 laneid, 1)]); fs (slice(0, 32), slice(0, 8)) Tx.prim_func def k(A_ptr: Tx.handle, B_ptr: Tx.handle): A Tx.match_buffer(A_ptr, (32, 8), float32); B Tx.match_buffer(B_ptr, (32, 8), float32) Tx.device_entry(); Tx.cta_id([1]); Tx.lane_id([32]); tid Tx.thread_id([32]) A_smem Tx.alloc_buffer((32, 8), float32, scopeshared, layoutTileLayout(S[(32, 8)])) Tx.tile.warp.copy(A_smem[fs], A[fs]); Tx.cuda.cta_sync() R Tx.alloc_buffer((32, 8), float32, scopelocal, layoutr_layout) Tx.tile.warp.copy(R[fs], A_smem[fs]) Tx.tile.warp.sqrt(R[fs], R[fs]) # 元素级 reg 派发 Tx.tile.warp.copy(A_smem[fs], R[fs]); Tx.cuda.cta_sync() Tx.tile.warp.copy(B[fs], A_smem[fs])程序先通过Tx.tile.warp.copy把全局数据搬进共享内存、再搬进按 lane 分行的局部 tileR随后Tx.tile.warp.sqrt(R[fs], R[fs])原地求平方根此处即触发reg变体的元素级派发最后写回全局。3.3 算法步骤emit_reg进入_emit_inducedinduced 路径锚点存在时无锚点的 trivial 情况则退化为每个线程在自己的私有存储上跑完整循环的扁平路径核心流程解析与检查spec.parse构建操作 plan谓词确认每个缓冲区操作数都是 local、校验 NumPy 风格广播、把锚点与完整作用域链对照检查并要求各操作数 thread/local/replica 布局签名兼容。诱导分区从锚点的线程轴此处即laneid出发丢弃线程迭代器给每个线程留下私有 bundle——本例中每个 lane 分到 8 个元素。对每线程 bundle 施加操作_pick_vec_and_carve按最宽优先widest-first枚举spec.vec_impls中注册的打包实现packed implementation如f32x2。对每个候选vec_len先尝试_carve_tail切出每个操作数仅含内存迭代器的布局尾部验证切出的尾部物理连续stride-1 链且乘积等于vec_len再调用impl.applies(...)确认该实现可用第一个对所有操作数都通过三关检查的候选胜出。若没有任何打包实现可用则回退到标量vec_len1、vec_implNone、使用未切分的原始布局操作退化为每线程循环内的标量计算。从实现细节看reg.py 的_emit_induced发射前还会做 region 切片与置换对齐align_operands_to_anchor随后剥离线程轴把每个操作数视为长度为per_thread_total的扁平 1D 视图通过_make_views_meta用T.decl_buffer建立别名视图以标量 offset 索引——这样 codegen 永远不会在T.vectorized内部看到多维get_indices。外层用T.serial循环而非T.unroll以避免为每个迭代物化临时的int[1]缓冲区导致代码膨胀ptxas 会对静态界循环自动展开。3.4 生成的 TIRx IR 与 CUDA标量回退路径生成的 TIRx IR 形态buffer[f] Tx.sqrt(buffer_1[f]) # 对 lane bundle 中每个元素 f生成的 CUDA每线程局部元素r_local_ptr[f_2] sqrtf(r_local_ptr[f_2]); // per-thread local element文档注明该结果已在sm_100a上验证输出等于sqrt(A)。3.5 输入如何改变算法inputeffectop一元 → 逐元素sqrtf/expf/ …二元 → 逐元素a bfma→a * b cdtype决定标量运算的选择当注册了打包形式如f32x2、双元素 cast时也会选用打包实现local-buffer layout锚点的线程轴决定分区每 lane 的 bundle 越宽串行循环越长四、smem变体shared 作用域的 synthesized 分区路径当元素级操作的所有缓冲区操作数包括输出都在共享内存中时派发落入smem变体。其实现位于 python/tvm/backend/cuda/tile_primitive/elementwise/smem.py设计上镜像了 copy 的 gmem_smem 路径vec_auto的全局/共享内存路径。4.1 接受条件is_smem_ewise谓词def check(op_call, sctx): if not sctx.is_target(cuda): return False, non-cuda target if sctx.scope_kind not in (thread, warp, warpgroup, cta): ... ok, reason _all_threads_active(sctx) # 完整作用域全部线程活跃 plan, msg spec.parse(op_call) # 解析操作的 operands for br in buffer_regions(plan): if not br.source.scope().startswith(shared): # 每个缓冲区操作数都必须是 shared* return False, foperand scope {br.source.scope()} ! shared* if br.source.layout is None: ... # spec.check_extrasdtype 规则与锚点布局校验注意这里的作用域匹配是startswith(shared)即shared以及可能的shared.dyn等变体都满足。PropertyRequirementtarget / scope / prioritycudathread/warp/warpgroup/cta全部线程活跃优先级10operands每个缓冲区操作数含输出都在shared*fill、二元操作与fma允许标量源op父页面列出的任意 CUDA 元素级OpSpec一元sqrt/exp/zero…、二元add/mul…、fmaspec.check_extras校验 dtype 组合layout操作数都带布局。分区由作用域的线程数合成dtype、逻辑最内层区域 extent、每线程元素数共同约束调度块宽度4.2 演示程序文档的演示改编自test_unary.py一个 CTA 对32×32的float32共享 tile 求sqrtCTA 共 256 线程因此分区只需一轮s_layout TileLayout(S[(32, 32)]); full (slice(0, 32), slice(0, 32)) Tx.prim_func def unary_op(A_ptr: Tx.handle): A Tx.match_buffer(A_ptr, (32, 32), float32, layouts_layout) Tx.device_entry(); Tx.cta_id([1]); Tx.warp_id([8]); Tx.lane_id([32]); Tx.thread_id([256]) A_smem Tx.alloc_buffer((32, 32), float32, scopeshared, layouts_layout) Tx.tile.cta.copy(A_smem[full], A[full]) Tx.tile.cta.sqrt(A_smem[full], A_smem[full]) # 元素级 smem 派发 Tx.tile.cta.copy(A[full], A_smem[full])4.3 算法步骤解析操作并检查操作数spec.parse把调用变成 plan输入、输出、操作谓词确认每个缓冲区操作数都在 shared 作用域并做右对齐广播兼容检查anchor 为plan.dst每个源须与目的区形状兼容extent 相等或为 1。合成分区与 copy 的 gmem_smem 路径一致根据作用域的线程数把区域拆成[outer, threads, vec]。候选宽度必须同时整除每线程元素数和每个操作数的逻辑最内层区域 extent。_max_layout_vec的实现要点是它只从候选集合{128, 64, 32, 16, 8}位宽按 dtype 位数换算成元素个数n中挑选且不检查物理布局的 stride/连续性。本例中 dense 单位布局下32×32 1024个float32由 256 线程处理得到vec 4、outer 1。逐元素施加操作与 copy 不同这里不是拷贝而是每个 (thread, round) 读取自己的vec个元素、施加操作、再写回——向量化执行。smen.py的emit_smem还有两个值得注意的工程细节一是线程数取自cuda/common.py的get_thread_cnt读取launch_params[threadIdx.x].dom.extent而非∏ sctx.intra——后者在 cta 作用域遇到 sub-warp 线程数时会静默算出 0二是当前实现假定 1DthreadIdx有断言packed-vec 发射还要求最内层维度 stride 为 1非 swizzle 切片以保证 lane 物理连续。4.4 生成的 TIRx IR 与 CUDAvec 4、outer 1时生成的 TIRx IRfor f in Tx.serial(1): # outer 1 for vec in Tx.vectorized(4): A_smem[tid * 4 vec] Tx.sqrt(A_smem[tid * 4 vec])vec 4的元素束在 CUDA 中成为一次float4访问操作逐分量施加float4 v_ *(float4*)(A_smem_ptr[tid * 4]); __1.x sqrtf(v_.x); __1.y sqrtf(v_.y); __1.z sqrtf(v_.z); __1.w sqrtf(v_.w);文档同样注明已在sm_100a上验证tile 结果等于sqrt(A)。4.5 输入如何改变算法inputeffectop一元 → 逐分量sqrtf/expf/ …二元 → 两个输入逐分量组合a bfma→a * b cdtype约束候选宽度上限每线程元素数与逻辑最内层区域 extent 可能进一步收窄它从而改变轮数。当前宽度选择不检查物理布局连续性scope决定线程轴与线程数从而决定合成的分区形态五、reg 与 smem两种路径的对比与选型要点维度reglocal 路径smemshared 路径操作数作用域全部local全部shared*分区来源诱导induced由锚点布局的线程轴决定线程轴被剥离后每线程得到私有 bundle合成synthesized由作用域线程数生成[outer, threads, vec]划分对齐方式align_operands_to_anchor切片 置换 线程/局部部分规范化replica 须预先一致无需跨操作数对齐共享缓冲按多维权下标 自身布局在 codegen 时解析物理地址向量化机制最宽优先尝试打包实现spec.vec_impls要求切尾物理连续否则标量循环_max_layout_vec从位宽候选集中取最宽整除值发射T.vectorized循环同步无额外同步需求各线程私有发射前/后需要作用域同步emit_scope_sync典型场景数据先落入寄存器如warp.copy到 local tile后的就地/异地图运算直接在共享内存 tile 上就地变换如归一化、激活一句话选型数据在寄存器里就地算走reg数据在共享内存里批量向量化算走smem。两者的判别是互斥的——只要所有缓冲区操作数都属于同一种作用域派发就是确定的。六、源码与测试佐证从注册到验证的完整证据链元素级原语的实现与验证可以归纳为三层证据链注册层register.py 为ALL_OPS中每个操作生成reg、smem两个register_dispatch条目variantpriority10when谓词reg/smem两个模块各自导出emit_*与is_*_ewise。操作层操作语义集中在 ops/unary.py8 个一元操作、ops/binary.py5 个二元操作含交换律与f32x2VecImpl 注册以及 ops/fma.py 与 ops/cast.py。打包向量化实现集中在 vec_emit/ 子目录如binary_f32x2.py、cast_vec2.py、fma_f32x2.py。测试层测试位于 tests/python/tirx/operator/tile_primitive/cuda/elementwise/按操作族分为test_unary.py、test_binary.py、test_fma.py。以 test_unary.py 为例它用pytest.mark.parametrize覆盖(32,32)与带偏移的(32,8,12)两种张量形状、zero/sqrt两种操作、float16→float16/float32→float16/float32→bfloat16三种 dtype 组合并分别构造 in-place同一 buffer与异地址A_smem与B_smem双缓冲两种 kernel通过Tx.cta.sqrt(...)等调用走 smem 变体最后与 NumPy 参考实现比对。测试均带pytest.mark.gpu与skipif(not env.has_cuda())条件需要真实 CUDA 环境执行。这套结构印证了文档中的描述变体判别只关心作用域操作语义完全由 OpSpec 描述发射算法induced / synthesized随后者决定。七、限制与边界最后需要明确当前实现的边界避免误用操作覆盖是当前集合全局Tx.tile目录中的minimum、memset、select等构造器目前没有上述 CUDAreg/smem变体。一次调用必须命中目标后端已注册的变体才能派发成功否则需要在调用侧改用其他可用原语。reg路径的布局要求锚点布局不能是 swizzlewarp/warpgroup/cta作用域下锚点的线程轴必须铺满完整作用域线程链各操作数 replica 签名必须一致。smem路径的向量化前提_max_layout_vec不检查物理布局步长而 packed 发射要求最内层维度 stride 为 1非 swizzle 切片以保证连续性当前实现还假定 1DthreadIdx。打包实现的可用性取决于平台如 ops/binary.py 所示add/sub/mul为sm_100注册了f32x2打包实现而fdiv/maximum没有打包 PTXmax本身降到单条FMNMX/max.f32精确无舍入无需打包变体——向量化能力随架构与 dtype 变化回退路径标量循环始终存在。赞分享模型编译深度学习推理引擎【免费下载链接】tvmOpen Machine Learning Compiler Framework项目地址https://gitcode.com/gh_mirrors/tv/tvm点击查看免费下载相关推荐TVM TIRx Tile Primitives 完全指南可调度的 CUDA 硬件原语与分发机制TVM TIRx Tile Primitives 完全指南可调度的 CUDA 硬件原语与分发机制 本指南系统讲解 TVM 中 TIRx 方言的 Tile Pr模型编译深度学习推理引擎TVM TIRx 降级流水线深度解析从 Tile 原语到 CUDA Kernel 的完整编译路径TVM TIRx 降级流水线深度解析从 Tile 原语到 CUDA Kernel 的完整编译路径 导读 本文围绕 TIRx 降级流水线文档 https://模型编译深度学习推理引擎TVM TIRx 异步拷贝copy_asyncTile 原语深度指南从 cp.async 到 TMA 与 tcgen05 的完整变体谱系TVM TIRx 异步拷贝copy_asyncTile 原语深度指南从 cp.async 到 TMA 与 tcgen05 的完整变体谱系 本篇指南系统讲解模型编译深度学习推理引擎创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 TVM Tirx CUDA 元素级 Tile 原语（elementwise）深度解析：reg 与 smem 双变体派发机制
模型编译深度学习推理引擎【免费下载链接】tvmOpen Machine Learning Compiler Framework项目地址https://gitcode.com/gh_mirrors/tv/tvm点击查看免费下载本文以 docs/tirx/tile_primitives/elementwise.rst 为主线系统讲解 TVM Tirxtile IR extension中 CUDA 后端元素级 tile 原语的完整设计从cast、fill、一元zero、reciprocal、sqrt、exp、exp2、log2、silu、二元add、sub、mul、fdiv、maximum到fma的操作集合到每个操作注册reg/smem两个变体、以操作数存储作用域作为互斥判别器的派发模型。读完本文你将掌握何时走寄存器路径、何时走共享内存路径、两种路径各自的分区生成算法induced vs synthesized、生成的 TIRx IR 与 CUDA 代码形态以及如何通过源码和测试用例验证这些行为。一、总览元素级 tile 原语的操作集与双变体注册模型在 Tirx 的Tx.tile目录中元素级elementwise原语覆盖了深度学习算子中最基础的一类对 tile 中每个元素独立施加同样的计算。当前 CUDA 元素级集合包含类型转换cast填充fill一元操作zero、reciprocal、sqrt、exp、exp2、log2、silu二元操作add、sub、mul、fdiv、maximum融合乘加fma。每个操作在 CUDA 后端注册两个变体——reg与smem优先级均为10。二者的判别完全由**缓冲区操作数的存储作用域storage scope**决定且是互斥的reg变体所有缓冲区操作数都在local作用域寄存器smem变体所有缓冲区操作数都在shared作用域共享内存。标量输入scalar inputs不拥有存储作用域因此在这两个变体中都被允许出现具体哪些操作允许标量取决于该操作的 authoring API见下文。每个操作都由一个OpSpec描述。OpSpec的核心是一套parse逻辑它负责构建目的区destination 源列表source list并携带可选的 dtype 检查check_extras以及逐元素应用的标量表达式compute_scalar。从源码看这一模型在 python/tvm/backend/cuda/tile_primitive/elementwise/register.py 中实现得十分简洁——遍历ALL_OPS中的每个OpSpec分别调用_register_reg与_register_smem注册两个派发变体例如register_dispatch( spec.name, cuda, variantreg, priority10, when[predicate(f{spec.name}_reg, is_reg_ewise(spec))], ) def _dispatch(op: TilePrimitiveCall, sctx: DispatchContext, _specspec) - PrimFunc: return emit_reg(op, _spec, sctx)smem变体与之对称只是variantsmem、谓词换成is_smem_ewise(spec)、发射函数换成emit_smem。该文件注释明确指出这种双变体模型模仿了 copy 原语的 PR-640 设计——scope-pair作用域组合决定派发选择底层算法induced 诱导 vs synthesized 合成随之而来。两个变体的差异可先由下表概览VariantOperandsLoweringreg所有缓冲区操作数均为local分区由本地缓冲区的布局诱导产生操作逐线程应用smem所有缓冲区操作数均为shared合成[outer, threads, vec]分区操作逐向量化元素应用值得注意的限制是全局Tx.tile目录中的其他构造器包括minimum、memset和select目前没有这些 CUDAreg/smem变体。也就是说一次Tx.tile.*调用必须由其所选目标后端注册的变体来承接否则无法派发。二、OpSpec解析、dtype 检查与逐元素标量表达式OpSpec是理解整个元素级派发机制的关键抽象。它把一个 tile 原语调用拆成三个层次parse(op_call) - (Plan, msg)把T.op(dst, src, ...)形式的调用解析为一个PlanPlan记录dst目的 TensorRegion与srcs源列表每个源要么是一个 buffer region要么是一个标量。check_extras(extras, compute_dtype)可选的 dtype 规则校验。例如一元操作的scale/bias标量必须与 compute dtype 一致。compute_scalar(src_vals, extras, dst_dtype)逐元素应用的标量表达式即真正施加在每个元素上的数学运算。以 python/tvm/backend/cuda/tile_primitive/elementwise/ops/unary.py 为例一元操作统一采用T.unary(dst, src[, bias, scale])签名bias/scale可选silu为保持历史行为忽略bias/scale。其UNARY_OPS字典清晰展示了 OpSpec 的四元组结构UNARY_OPS: dict[str, OpSpec] { zero: OpSpec(zero, _parse_unary, _compute_zero, _check_unary_extras), fill: OpSpec(fill, _parse_unary, _compute_fill, _check_unary_extras), reciprocal: OpSpec(reciprocal, _parse_unary, _compute_reciprocal, _check_unary_extras), sqrt: OpSpec(sqrt, _parse_unary, _with_bias_scale(T.sqrt), _check_unary_extras), exp: OpSpec(exp, _parse_unary, _with_bias_scale(T.exp), _check_unary_extras), exp2: OpSpec(exp2, _parse_unary, _with_bias_scale(T.exp2), _check_unary_extras), log2: OpSpec(log2, _parse_unary, _with_bias_scale(T.log2), _check_unary_extras), silu: OpSpec(silu, _parse_unary, _compute_silu, _check_unary_extras), }其中_with_bias_scale是一个高阶包装先施加scale乘法、再施加biasbuffer 或常量最后调用裸运算如T.exp_compute_zero恒返回0.0_compute_reciprocal计算1.0 / x_compute_silu计算x / (1 exp(0 - x))。_parse_unary支持源既可以是TensorRegion也可以是 primitive 表达式标量并处理bias为 buffer 与常量两种形态。二元操作在 python/tvm/backend/cuda/tile_primitive/elementwise/ops/binary.py 中实现还包含常量左操作数的交换律commute逻辑_COMMUTATIVE {add, mul, maximum}——当左源是常量时对可交换操作将其移到右侧不可交换操作则拒绝。两个源都是常量时直接报错both inputs are constants。广播extent1 的维度在派发的_broadcast_lift布局层面处理解析器只是按原样记录每个源。三、reg变体local 作用域的 induced 分区路径当一次元素级操作的所有缓冲区操作数都位于local寄存器作用域时派发落入reg变体。其实现位于 python/tvm/backend/cuda/tile_primitive/elementwise/reg.py设计上镜像了 copy 的 reg 路径vec_auto的寄存器路径。3.1 接受条件is_reg_ewise谓词is_reg_ewise(spec)返回一个检查函数它依次验证def check(op_call, sctx): if not sctx.is_target(cuda): return False, non-cuda target if sctx.scope_kind not in (thread, warp, warpgroup, cta): ... ok, reason _all_threads_active(sctx) plan, msg spec.parse(op_call) for br in buffer_regions(plan): if br.source.scope() ! local: # 每个缓冲区操作数都必须是 local return False, foperand scope {br.source.scope()} ! local if br.source.layout is None: ... # spec.check_extrasdtype 规则 # pick_anchor _validate_anchor_layout锚点必须是 TileLayout、非 swizzle # _validate_scope_level_anchorwarp/warpgroup/cta 作用域下锚点的线程轴 # 必须覆盖完整作用域线程链 # NumPy 风格右对齐形状广播检查 # _check_layout_operands_agree各操作数 replica 签名一致将接受条件整理为属性表PropertyRequirementtarget / scope / prioritycudathread/warp/warpgroup/cta全部线程活跃优先级10operands每个缓冲区操作数都在localfill、二元操作与fma允许标量源op父页面列出的任意 CUDA 元素级OpSpec一元 / 二元 /fmaspec.check_extras校验 dtype 组合local-buffer layout锚点布局必须通过校验且其线程轴必须匹配完整的作用域链它诱导分区所有操作数在 thread、per-thread storage、replica 签名上必须一致shapes每个输入区都与目的区 NumPy 风格右对齐广播兼容几个关键校验点的含义值得展开**锚点anchor**由pick_anchor选出通常是目的操作数。_validate_anchor_layout要求锚点布局不是 swizzle、且必须是TileLayout。**作用域级锚点_validate_scope_level_anchor**针对warp/warpgroup/cta作用域经过目标规范化后锚点切片的线程轴只能是该作用域对应的轴laneid/tid_in_wg/tx且按 stride 排序后必须完整铺满该作用域的所有线程stride 依次为 1、e0、e0×e1…直至线程总数不留空隙、不重叠。thread作用域豁免此检查。布局签名一致_check_layout_operands_agree对齐align_operands_to_anchor只能规范 thread 与 local 部分replica 部分不被对齐触碰因此各操作数的 replica 签名轴 key、extent、stride必须逐项一致否则发射结果会是错的。3.2 演示程序文档给出的演示是一个 warp 对32×8的float32局部 tile 求sqrt局部布局为S[(32,8):(1laneid,1)]——即 lanei拥有第i行from tvm.tirx.layout import S, TileLayout, laneid r_layout TileLayout(S[(32, 8) : (1 laneid, 1)]); fs (slice(0, 32), slice(0, 8)) Tx.prim_func def k(A_ptr: Tx.handle, B_ptr: Tx.handle): A Tx.match_buffer(A_ptr, (32, 8), float32); B Tx.match_buffer(B_ptr, (32, 8), float32) Tx.device_entry(); Tx.cta_id([1]); Tx.lane_id([32]); tid Tx.thread_id([32]) A_smem Tx.alloc_buffer((32, 8), float32, scopeshared, layoutTileLayout(S[(32, 8)])) Tx.tile.warp.copy(A_smem[fs], A[fs]); Tx.cuda.cta_sync() R Tx.alloc_buffer((32, 8), float32, scopelocal, layoutr_layout) Tx.tile.warp.copy(R[fs], A_smem[fs]) Tx.tile.warp.sqrt(R[fs], R[fs]) # 元素级 reg 派发 Tx.tile.warp.copy(A_smem[fs], R[fs]); Tx.cuda.cta_sync() Tx.tile.warp.copy(B[fs], A_smem[fs])程序先通过Tx.tile.warp.copy把全局数据搬进共享内存、再搬进按 lane 分行的局部 tileR随后Tx.tile.warp.sqrt(R[fs], R[fs])原地求平方根此处即触发reg变体的元素级派发最后写回全局。3.3 算法步骤emit_reg进入_emit_inducedinduced 路径锚点存在时无锚点的 trivial 情况则退化为每个线程在自己的私有存储上跑完整循环的扁平路径核心流程解析与检查spec.parse构建操作 plan谓词确认每个缓冲区操作数都是 local、校验 NumPy 风格广播、把锚点与完整作用域链对照检查并要求各操作数 thread/local/replica 布局签名兼容。诱导分区从锚点的线程轴此处即laneid出发丢弃线程迭代器给每个线程留下私有 bundle——本例中每个 lane 分到 8 个元素。对每线程 bundle 施加操作_pick_vec_and_carve按最宽优先widest-first枚举spec.vec_impls中注册的打包实现packed implementation如f32x2。对每个候选vec_len先尝试_carve_tail切出每个操作数仅含内存迭代器的布局尾部验证切出的尾部物理连续stride-1 链且乘积等于vec_len再调用impl.applies(...)确认该实现可用第一个对所有操作数都通过三关检查的候选胜出。若没有任何打包实现可用则回退到标量vec_len1、vec_implNone、使用未切分的原始布局操作退化为每线程循环内的标量计算。从实现细节看reg.py 的_emit_induced发射前还会做 region 切片与置换对齐align_operands_to_anchor随后剥离线程轴把每个操作数视为长度为per_thread_total的扁平 1D 视图通过_make_views_meta用T.decl_buffer建立别名视图以标量 offset 索引——这样 codegen 永远不会在T.vectorized内部看到多维get_indices。外层用T.serial循环而非T.unroll以避免为每个迭代物化临时的int[1]缓冲区导致代码膨胀ptxas 会对静态界循环自动展开。3.4 生成的 TIRx IR 与 CUDA标量回退路径生成的 TIRx IR 形态buffer[f] Tx.sqrt(buffer_1[f]) # 对 lane bundle 中每个元素 f生成的 CUDA每线程局部元素r_local_ptr[f_2] sqrtf(r_local_ptr[f_2]); // per-thread local element文档注明该结果已在sm_100a上验证输出等于sqrt(A)。3.5 输入如何改变算法inputeffectop一元 → 逐元素sqrtf/expf/ …二元 → 逐元素a bfma→a * b cdtype决定标量运算的选择当注册了打包形式如f32x2、双元素 cast时也会选用打包实现local-buffer layout锚点的线程轴决定分区每 lane 的 bundle 越宽串行循环越长四、smem变体shared 作用域的 synthesized 分区路径当元素级操作的所有缓冲区操作数包括输出都在共享内存中时派发落入smem变体。其实现位于 python/tvm/backend/cuda/tile_primitive/elementwise/smem.py设计上镜像了 copy 的 gmem_smem 路径vec_auto的全局/共享内存路径。4.1 接受条件is_smem_ewise谓词def check(op_call, sctx): if not sctx.is_target(cuda): return False, non-cuda target if sctx.scope_kind not in (thread, warp, warpgroup, cta): ... ok, reason _all_threads_active(sctx) # 完整作用域全部线程活跃 plan, msg spec.parse(op_call) # 解析操作的 operands for br in buffer_regions(plan): if not br.source.scope().startswith(shared): # 每个缓冲区操作数都必须是 shared* return False, foperand scope {br.source.scope()} ! shared* if br.source.layout is None: ... # spec.check_extrasdtype 规则与锚点布局校验注意这里的作用域匹配是startswith(shared)即shared以及可能的shared.dyn等变体都满足。PropertyRequirementtarget / scope / prioritycudathread/warp/warpgroup/cta全部线程活跃优先级10operands每个缓冲区操作数含输出都在shared*fill、二元操作与fma允许标量源op父页面列出的任意 CUDA 元素级OpSpec一元sqrt/exp/zero…、二元add/mul…、fmaspec.check_extras校验 dtype 组合layout操作数都带布局。分区由作用域的线程数合成dtype、逻辑最内层区域 extent、每线程元素数共同约束调度块宽度4.2 演示程序文档的演示改编自test_unary.py一个 CTA 对32×32的float32共享 tile 求sqrtCTA 共 256 线程因此分区只需一轮s_layout TileLayout(S[(32, 32)]); full (slice(0, 32), slice(0, 32)) Tx.prim_func def unary_op(A_ptr: Tx.handle): A Tx.match_buffer(A_ptr, (32, 32), float32, layouts_layout) Tx.device_entry(); Tx.cta_id([1]); Tx.warp_id([8]); Tx.lane_id([32]); Tx.thread_id([256]) A_smem Tx.alloc_buffer((32, 32), float32, scopeshared, layouts_layout) Tx.tile.cta.copy(A_smem[full], A[full]) Tx.tile.cta.sqrt(A_smem[full], A_smem[full]) # 元素级 smem 派发 Tx.tile.cta.copy(A[full], A_smem[full])4.3 算法步骤解析操作并检查操作数spec.parse把调用变成 plan输入、输出、操作谓词确认每个缓冲区操作数都在 shared 作用域并做右对齐广播兼容检查anchor 为plan.dst每个源须与目的区形状兼容extent 相等或为 1。合成分区与 copy 的 gmem_smem 路径一致根据作用域的线程数把区域拆成[outer, threads, vec]。候选宽度必须同时整除每线程元素数和每个操作数的逻辑最内层区域 extent。_max_layout_vec的实现要点是它只从候选集合{128, 64, 32, 16, 8}位宽按 dtype 位数换算成元素个数n中挑选且不检查物理布局的 stride/连续性。本例中 dense 单位布局下32×32 1024个float32由 256 线程处理得到vec 4、outer 1。逐元素施加操作与 copy 不同这里不是拷贝而是每个 (thread, round) 读取自己的vec个元素、施加操作、再写回——向量化执行。smen.py的emit_smem还有两个值得注意的工程细节一是线程数取自cuda/common.py的get_thread_cnt读取launch_params[threadIdx.x].dom.extent而非∏ sctx.intra——后者在 cta 作用域遇到 sub-warp 线程数时会静默算出 0二是当前实现假定 1DthreadIdx有断言packed-vec 发射还要求最内层维度 stride 为 1非 swizzle 切片以保证 lane 物理连续。4.4 生成的 TIRx IR 与 CUDAvec 4、outer 1时生成的 TIRx IRfor f in Tx.serial(1): # outer 1 for vec in Tx.vectorized(4): A_smem[tid * 4 vec] Tx.sqrt(A_smem[tid * 4 vec])vec 4的元素束在 CUDA 中成为一次float4访问操作逐分量施加float4 v_ *(float4*)(A_smem_ptr[tid * 4]); __1.x sqrtf(v_.x); __1.y sqrtf(v_.y); __1.z sqrtf(v_.z); __1.w sqrtf(v_.w);文档同样注明已在sm_100a上验证tile 结果等于sqrt(A)。4.5 输入如何改变算法inputeffectop一元 → 逐分量sqrtf/expf/ …二元 → 两个输入逐分量组合a bfma→a * b cdtype约束候选宽度上限每线程元素数与逻辑最内层区域 extent 可能进一步收窄它从而改变轮数。当前宽度选择不检查物理布局连续性scope决定线程轴与线程数从而决定合成的分区形态五、reg 与 smem两种路径的对比与选型要点维度reglocal 路径smemshared 路径操作数作用域全部local全部shared*分区来源诱导induced由锚点布局的线程轴决定线程轴被剥离后每线程得到私有 bundle合成synthesized由作用域线程数生成[outer, threads, vec]划分对齐方式align_operands_to_anchor切片 置换 线程/局部部分规范化replica 须预先一致无需跨操作数对齐共享缓冲按多维权下标 自身布局在 codegen 时解析物理地址向量化机制最宽优先尝试打包实现spec.vec_impls要求切尾物理连续否则标量循环_max_layout_vec从位宽候选集中取最宽整除值发射T.vectorized循环同步无额外同步需求各线程私有发射前/后需要作用域同步emit_scope_sync典型场景数据先落入寄存器如warp.copy到 local tile后的就地/异地图运算直接在共享内存 tile 上就地变换如归一化、激活一句话选型数据在寄存器里就地算走reg数据在共享内存里批量向量化算走smem。两者的判别是互斥的——只要所有缓冲区操作数都属于同一种作用域派发就是确定的。六、源码与测试佐证从注册到验证的完整证据链元素级原语的实现与验证可以归纳为三层证据链注册层register.py 为ALL_OPS中每个操作生成reg、smem两个register_dispatch条目variantpriority10when谓词reg/smem两个模块各自导出emit_*与is_*_ewise。操作层操作语义集中在 ops/unary.py8 个一元操作、ops/binary.py5 个二元操作含交换律与f32x2VecImpl 注册以及 ops/fma.py 与 ops/cast.py。打包向量化实现集中在 vec_emit/ 子目录如binary_f32x2.py、cast_vec2.py、fma_f32x2.py。测试层测试位于 tests/python/tirx/operator/tile_primitive/cuda/elementwise/按操作族分为test_unary.py、test_binary.py、test_fma.py。以 test_unary.py 为例它用pytest.mark.parametrize覆盖(32,32)与带偏移的(32,8,12)两种张量形状、zero/sqrt两种操作、float16→float16/float32→float16/float32→bfloat16三种 dtype 组合并分别构造 in-place同一 buffer与异地址A_smem与B_smem双缓冲两种 kernel通过Tx.cta.sqrt(...)等调用走 smem 变体最后与 NumPy 参考实现比对。测试均带pytest.mark.gpu与skipif(not env.has_cuda())条件需要真实 CUDA 环境执行。这套结构印证了文档中的描述变体判别只关心作用域操作语义完全由 OpSpec 描述发射算法induced / synthesized随后者决定。七、限制与边界最后需要明确当前实现的边界避免误用操作覆盖是当前集合全局Tx.tile目录中的minimum、memset、select等构造器目前没有上述 CUDAreg/smem变体。一次调用必须命中目标后端已注册的变体才能派发成功否则需要在调用侧改用其他可用原语。reg路径的布局要求锚点布局不能是 swizzlewarp/warpgroup/cta作用域下锚点的线程轴必须铺满完整作用域线程链各操作数 replica 签名必须一致。smem路径的向量化前提_max_layout_vec不检查物理布局步长而 packed 发射要求最内层维度 stride 为 1非 swizzle 切片以保证连续性当前实现还假定 1DthreadIdx。打包实现的可用性取决于平台如 ops/binary.py 所示add/sub/mul为sm_100注册了f32x2打包实现而fdiv/maximum没有打包 PTXmax本身降到单条FMNMX/max.f32精确无舍入无需打包变体——向量化能力随架构与 dtype 变化回退路径标量循环始终存在。赞分享模型编译深度学习推理引擎【免费下载链接】tvmOpen Machine Learning Compiler Framework项目地址https://gitcode.com/gh_mirrors/tv/tvm点击查看免费下载相关推荐TVM TIRx Tile Primitives 完全指南可调度的 CUDA 硬件原语与分发机制TVM TIRx Tile Primitives 完全指南可调度的 CUDA 硬件原语与分发机制 本指南系统讲解 TVM 中 TIRx 方言的 Tile Pr模型编译深度学习推理引擎TVM TIRx 降级流水线深度解析从 Tile 原语到 CUDA Kernel 的完整编译路径TVM TIRx 降级流水线深度解析从 Tile 原语到 CUDA Kernel 的完整编译路径 导读 本文围绕 TIRx 降级流水线文档 https://模型编译深度学习推理引擎TVM TIRx 异步拷贝copy_asyncTile 原语深度指南从 cp.async 到 TMA 与 tcgen05 的完整变体谱系TVM TIRx 异步拷贝copy_asyncTile 原语深度指南从 cp.async 到 TMA 与 tcgen05 的完整变体谱系 本篇指南系统讲解模型编译深度学习推理引擎创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考