模型编译深度学习推理引擎【免费下载链接】tvmOpen Machine Learning Compiler Framework项目地址https://gitcode.com/gh_mirrors/tv/tvm点击查看免费下载导读本文聚焦 TVM TIRx CUDA 后端归约算子sum/max/min的local变体——即当源src与目标dst缓冲区都位于local寄存器/线程私有存储作用域时编译器如何完成算子分发、布局校验与代码生成。你将掌握该变体的三层递进算法路径线程级顺序归约、warp 级 laneid shard→replica 专用 shuffle 归约、通用 warp/warpgroup 视图归约、每个输入的语义影响以及对应源码位置与测试用例的验证方式可直接用于阅读或扩展 TIRx 归约管线。变体定位归约算子家族中的 local 路径在 TIRx 的归约算子体系中sum、max、min三种操作都通过register_dispatch在 CUDA 目标上注册了多个变体按操作数存储作用域与硬件能力区分见 reduction 索引页变体优先级降级策略reduction/local10local 缓冲区 src/dst线程级顺序归约可选 warp shufflereduction/shared10shared 缓冲区 src/dst自适应分组__shfl_xor树reduction/sm100_packedpacked_add_sum/3input_maxmin20CUDA SM100 线程级 fp32 ≥8 元素打包add.f32x2/max3/min3本文讨论的local变体位于 python/tvm/backend/cuda/tile_primitive/reduction/local.py是三个变体中唯一要求 src 与 dst 均为local作用域的实现。接受条件predicate 双层校验注册声明local变体在local.py末尾通过循环注册三个算子共用同一套校验逻辑见 local.py#L471-L489for op_name, op_type in [ (sum, ReduceOpType.SUM), (max, ReduceOpType.MAX), (min, ReduceOpType.MIN), ]: register_dispatch( op_name, cuda, variantlocal, priority10, when[ predicate(storage_scope, _match_reduction_storage_scope, expected_scope[local]), predicate(local_valid, validate_reduction_local), ], ) def _local_dispatch(op: TilePrimitiveCall, sctx: DispatchContext, _op_typeop_type) - PrimFunc: op TilePrimitiveCall.downcast(op) return reduction_local_impl(op, _op_type, sctx)第一层_match_reduction_storage_scope定义于 reduction/utils.py#L61-L71要求src 与 dst 的作用域同时匹配local第二层validate_reduction_local再深入检查执行作用域、布局与轴信息。详细要求validate_reduction_locallocal.py#L152-L224逐条落实以下约束属性要求target / prioritycuda优先级10操作数作用域src 与 dst 均为local且 dtype 相等src.dtype ! dst.dtype直接拒绝执行作用域thread顺序、线程本地warp/warpgroup需要合法且非 swizzle 的TileLayout。warp 作用域下若 src/dst 呈 laneid shard→replica 模式则自动选择专用 shuffle 路径否则thread_reduceTrue可为通用路径附加 shuffle 步骤。warpgroup 作用域拒绝thread_reduceTrue报错thread_reduceTrue is only supported in warp scope; warpgroup local reduction is thread-local only形状thread 作用域下校验器不检查轴、不比较 src/dst 尺寸轴分析与循环构造推迟到 emit 阶段宽作用域warp/warpgroup视图归约要求空间维的布局结构匹配线程/局部 extent 一致被归约维在 dst 中的局部 extent 为 1从源码结构看warp 作用域的校验顺序有一个值得注意的细节先尝试_analyze_shuffle_reduce识别 laneid shard→replica 模式并提前放行因为该模式中 laneid 出现在 dst 的 replica广播里会被后续通用布局校验拒绝见 local.py#L188-L194 的注释。演示程序线程级 4 元素向量归约原文档给出的演示程序出自 test_reduction.py在单个线程内把一个 4 元素float32local 向量归约为标量Tx.prim_func def test_func(A_ptr: Tx.handle, B_ptr: Tx.handle): A Tx.match_buffer(A_ptr, [4], float32, layoutTileLayout(S[(4,)])) B Tx.match_buffer(B_ptr, [1], float32, layoutTileLayout(S[(1,)])) Tx.device_entry(); Tx.cta_id([1]); Tx.thread_id([1]) A_local Tx.alloc_buffer([4], float32, scopelocal) B_local Tx.alloc_buffer([1], float32, scopelocal) for i in Tx.serial(4): A_local[i] A[i] Tx.tile.sum(B_local, A_local, accumFalse) # reduction local dispatch B[0] B_local[0]由于归约长度只有 4 8此用例停留在local路径而非跳到 SM100 的packed_add_sum/3input_maxmin快速路径参见 reduction/sm100_packed 文档与 reduction/sm100_packed.py。归约算子其余参数reduce_axes、accum、config由_reduction_args统一解析见 reduction/utils.py#L48-L58。算法三条路径的选择逻辑reduction_local_impllocal.py#L403-L464根据执行作用域与布局形态选择路径路径一线程级顺序归约_emit_reduction_local_thread_wise当执行作用域为thread时走此路径local.py#L227-L283对输出位置做空间循环每个位置先初始化为算子单位元sum→0.0、max→T.min_value(dtype)、min→T.max_value(dtype)见 reduction/utils.py#L40-L45除非accumTrue随后内层归约循环逐元素累积源数据。该路径完全没有跨线程通信for spa in Tx.serial(spatial_len): if not accum: dst[spa] identity for red in Tx.serial(reduction_len): dst[spa] op(dst[spa], src[spa, red])其中spatial_len与reduction_len分别是空间维与归约维 extent 的乘积负轴先经_analyze_axes归一化为非负下标reduction/utils.py#L74-L82。_emit_reduction_local_thread_wise的 docstring 给出了一个 2D 示例Tx.sum(B_local[0:2, 0:3], A_local[0:2, 0:3, 0:4], [-1], False)展开为spatial_len6、reduction_len4的双层循环local.py#L24-L34。路径二专用 laneid shard→replica shuffle 归约_gen_warp_shuffle_reduce当 src 具有全跨度32 lanes的 laneid shard、dst 具有 2 的幂次 laneid replica 时_analyze_shuffle_reducelocal.py#L90-L127返回(reduce_width, local_elems)其中reduce_width参与每组归约的 lane 数dst replica 中 laneid 迭代子 extent 的乘积必须为正、≤32 且为 2 的幂local_elemssrc 中非 laneid shard extent 的乘积即每个线程持有的元素个数swizzle 布局直接返回None不支持。命中后由_gen_warp_shuffle_reducelocal.py#L130-L149生成实现每个 lane 先把本线程对应的局部值复制到 dst 视图src/dst为同一缓冲区时跳过复制再对每个元素执行T.cuda.warp_reduce(dst_local[k], op_str, reduce_width)。该路径由布局自动触发与thread_reduce无关也不先运行通用局部轴循环同时它不分支处理accum参数因此总是用 shuffle 结果覆盖 dst。路径三通用 warp/warpgroup 视图归约_emit_reduction_local_view未命中专用模式时warp/warpgroup 走此路径local.py#L286-L400通过get_local_region从TileLayout中分解出每个线程的局部区域把 src 的局部归约维逐一归约进 dst 的每个位置。在 warp 作用域下thread_reduceTrue会额外发出显式的tvm_warp_shuffle_xor步骤掩码来自_compute_shuffle_masksreduction/utils.py#L111-L128——对归约维中每个线程迭代子按其 stride 生成stride * 2^ii 从 0 到 log2(extent)的 XOR 掩码并升序排列同时用T.tvm_warp_activemask()取活跃掩码。warpgroup 作用域只支持纯 local 部分不生成 shuffle。该路径还有一个关键组合accumTrue且需要 shuffle 时实现会先在归约前把旧 dst 值保存到临时old_val局部缓冲区完成局部归约与 shuffle 后再与新结果结合保证累积语义正确local.py#L355-L378。生成的 IR 与 CUDA 代码对演示中的 4 元素线程级归约生成的 TIRx IR 为for spa in Tx.serial(1): dst[...] Tx.float32(0) for red in Tx.serial(4): dst[...] dst[...] src[...] # op sum对应的 CUDA C 为for (int red 0; red 4; red) B_local_ptr[0] B_local_ptr[0] A_local_ptr[red];该用例已在sm_100a上验证B sum(A)成立测试由 test_reduction.py 中的 GPU 测试覆盖。注意 IR 与 CUDA 均不含跨线程通信完全符合路径一的顺序语义。输入如何改变算法行为输入影响opsum→max→maxmin→min以及对应单位元算子到字符串的映射表_REDUCE_OP_TO_STR见 reduction/utils.py#L216exec scopethread→顺序路径warp 下匹配 shard→replica 布局→专用warp_reduce否则 warp/warpgroup 走通用局部视图路径仅thread_reduceTrue时附加 warp shufflewarpgroup 不允许axes / shape决定空间循环与归约循环的 extent_analyze_axes将负轴归一化accum线程级与通用视图路径按上文方式复用旧 dst 值线程级为跳过初始化视图shuffle 为保存旧值后合并专用 shard→replica 路径忽略该标志总是覆盖 dst测试与验证local变体的行为由 tests/python/tirx/operator/tile_primitive/cuda/reduction/test_reduction.py 系统性验证覆盖了三类场景test_reduction_local_thread_wise对多种 shape/axes 组合1D→1D、4D→2D、3D→2D、非 2 的幂、带 offset 切片参数化测试线程级顺序归约并同时参数化op_typesum/max/min、dtypefloat32/float16与accumtest_reduction_local_view_basic/test_reduction_local_view_complex前者验证纯 local 布局的视图归约含切片输入后者使用 WGMMA 风格的多层 tile 布局并参数化thread_reduceshuffle 开关与accum还演示了“先thread_reduceFalse归约、再用Tx.warp.sum(red_view, red_view, thread_reduceTrue)补一步 shuffle”的用法test_reduction_op_warp_shuffle系列覆盖 laneid shard→replica 专用路径——32 lane 全 warp 归约到 1 值并广播、每线程多元素4 元素×32 lane、稀疏存储槽storage span 7 但仅 4 元素等边界情形另有test_reduction_warp_shuffle_multi_warp_loop验证循环内 thread→warp 作用域交替的跨 warp 归约组合。这些测试在tir_pipelinetirx下编译并在真实 CUDA 设备上用tvm.testing.assert_allclose与 NumPy 参考结果比对。此外 test_dispatcher.py 覆盖分发框架本身的注册与 predicate 判定逻辑。小结local变体是 TIRx 归约体系中最贴近硬件线程模型的实现线程级路径完全避免通信专用 shard→replica 路径把归约折叠进单条warp_reduce通用视图路径则以布局分解为代价换取 warp/warpgroup 下对任意 local 视图的归约能力。阅读 local.py 及其配套的 reduction/utils.py 和 test_reduction.py即可完整掌握该变体的分发条件、三条算法路径与验证手段为后续在 TIRx 中编写或扩展归约类 tile primitive 提供直接参考。赞分享模型编译深度学习推理引擎【免费下载链接】tvmOpen Machine Learning Compiler Framework项目地址https://gitcode.com/gh_mirrors/tv/tvm点击查看免费下载相关推荐Apache TVM TIRx CUDA 归约原语Reduction Tile Primitive完全指南local / shared / SM100 packed 三变体的调度、算法与代码生成Apache TVM TIRx CUDA 归约原语Reduction Tile Primitive完全指南local / shared / SM100 p模型编译深度学习推理引擎TVM TIRx tcgen05 张量内存与寄存器异步拷贝copy_async tmem-local 变体tcgen05.ld/st深度解析TVM TIRx tcgen05 张量内存与寄存器异步拷贝copy_async tmem local 变体tcgen05.ld/st深度解析 本篇技术指模型编译深度学习推理引擎TVM TIRx CUDA 分布式共享内存拷贝copy_async 的 dsmem 变体深度解析TVM TIRx CUDA 分布式共享内存拷贝copy_async 的 dsmem 变体深度解析 导读 在 TVM TIRxTile IR eXtensio模型编译深度学习推理引擎创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考