torch-catlass 测试框架设计文档1. 总览tests/optest是 CATLASS 示例算子接入 PyTorch 的端到端测试框架。框架将 CATLASS AscendC kernel 封装为torch.ops.catlass.*算子并通过 Python 包torch_catlass提供测试入口。框架按职责分为五层Python API torch_catlass.ops.* | v Python package loader load kernel libs and libcatlass_torch.so | v PyTorch C extension register torch.ops.catlass.* | v Kernel adapter convert Tensor arguments to CATLASS kernel params | v Kernel implementation prebuilt kernel or JIT compiled template核心原则Python 层只负责用户接口、动态库加载和轻量参数转换。C extension 层负责 PyTorch op 注册和 NPU dispatch。adapter 层负责 Tensor 到 kernel ABI 的转换。kernel 层负责 AscendC/CATLASS 代码执行。JIT 子系统负责模板参数宏生成、编译、缓存和动态加载。2. 目录模块tests/optest/ ├── pyproject.toml ├── CMakeLists.txt ├── build.sh ├── docs/ │ └── design.md ├── include/ │ ├── catlass_kernel.h │ └── catlass_torch.h ├── torch_catlass/ │ ├── __init__.py │ ├── _version.py │ └── ops/ ├── src/ │ ├── catlass_torch.cpp │ ├── common/ │ └── include/ ├── utils/ │ ├── CMakeLists.txt │ ├── include/ │ ├── kernel_utils.cpp │ ├── torch_utils.cpp │ └── type_utils.hpp ├── kernels/ │ ├── CMakeLists.txt │ ├── common/ │ ├── include/ │ ├── jit/ │ └── 00_basic_matmul/ └── tests/ └── test_00_basic_matmul.py模块责任torch_catlass/Python 包入口、动态库加载、用户侧 op wrapperinclude/框架和 kernel 共享的公共 ABI 声明src/PyTorch C extension 和 torch op 注册utils/dtype/layout/Tensor 工具函数拆分 torch 依赖和纯 ACL 依赖kernels/kernel 构建、JIT compiler、JIT template、kernel entrytests/pytest 集成测试3. Python 包模块3.1 包初始化torch_catlass/__init__.py在 import 时完成运行时初始化从_version.py读取构建期版本信息。设置TORCH_CATLASS_VERSION供 JIT 编译阶段注入版本宏。设置TORCH_CATLASS_PKG_DIR供 JIT compiler 定位安装后的 headers 和 templates。加载 JIT compiler 和 JIT kernel entry 库。根据当前 NPU 架构加载 arch-specific kernel 库。调用torch.ops.load_library()加载libcatlass_torch.so注册torch.ops.catlass.*。动态库加载顺序lib/jit/libcatlass_kernel_jit_compiler.so lib/jit/libcatlass_kernel_jit.so lib/arch/*.so lib/libcatlass_torch.soJIT compiler 和 kernel entry 必须先于 PyTorch extension 加载保证 extension 中引用的 kernel 符号可解析。3.2 架构识别Python loader 通过torch_npu.npu.get_device_name()识别设备并映射为 CATLASS arch id设备名archAscend910B.*2201Ascend910_932201Ascend950PR/Ascend950DT3510当torch_npu.npu.device_count()为 0 时loader 抛出明确错误。测试代码在无 NPU 环境下通过 pytest skip 处理避免在 collection 阶段触发 torch-npu 内部错误。3.3 Python op wrappertorch_catlass/ops/保存 Python 用户接口。以basic_matmul为例torch_catlass.basic_matmul( mat1, mat2, outDTypefloat16, transAFalse, transBFalse, formatAFalse, formatBFalse, )Python wrapper 只做用户友好的轻量转换例如 dtype 字符串白名单解析。shape 推导、输出分配、stream 获取和 kernel launch 都在 C 层完成避免 Python 和 C 维护两套语义。4. PyTorch C Extension 模块4.1 注册入口src/catlass_torch.cpp是 PyTorch extension 的注册入口using BasicMatmulOp MatmulLikeCatlassKernel::BasicMatmul; static auto basic_matmul BasicMatmulOp::Run; REGISTER_TORCH_FUNC(basic_matmul);REGISTER_TORCH_FUNC位于src/include/common/register.h注册流程为创建或复用torch::Librarynamespace 固定为catlass。通过 PyTorch schema inference 从 C 函数签名生成 schema。将实现注册到c10::DispatchKey::PrivateUse1。PrivateUse1是 torch-npu 使用的 NPU dispatch key。Python 调用torch.ops.catlass.basic_matmul时PyTorch 根据输入 Tensor device 走到该 backend 实现。4.2 kernel launch 包装RUN_NPU_FUNC位于src/include/common/run_npu_func.h。它通过 torch-npu 的OpCommand::RunOpApiV2执行 kernel launchlaunch 前检查函数指针是否为空。将 C 异常转为 ACL error code。将 kernel 调用交给 torch-npu runtime 管理。5. Matmul Adapter 模块src/include/template/matmul.h提供MatmulLikeKernelFunc。该模板封装 matmul 类算子的通用流程Run() ├─ GetKernelInfo() ├─ AllocOutput() ├─ get current NPU stream ├─ get AIC core count └─ RUN_NPU_FUNC(KernelFunc, ...)5.1 参数拆分Matmul 参数拆分为两类参数结构用途是否参与 JIT 编译MatmulTParamsdtype、layout、transpose 等模板参数是MatmulParamsM/N/K、输入输出地址等运行时参数否这种拆分保证 dtype/layout 变化会生成新的模板实例而 shape 和 Tensor 地址变化不会导致重复编译。5.2GetKernelInfoGetKernelInfo()负责将 PyTorch Tensor 转为 kernel 参数将torch::Dtype转为aclDataType。根据transA和transB推导m/n/k。检查两个输入矩阵的 K 维一致。填充MatmulTParams。填充MatmulParams。将输入 Tensor storage 地址写入params.inputAddr。5.3AllocOutputAllocOutput()根据params.m、params.n和tParams.elementC创建输出 Tensor并将其 storage 地址写入params.outputAddr[0]。输出 Tensor 生命周期由 PyTorch 管理kernel ABI 只接收裸地址。6. 公共 ABI 模块include/catlass_kernel.h定义 C wrapper 和 kernel 实现之间共享的数据结构和函数声明。6.1 matmul ABIMatmulTParams表示编译期参数elementAelementBelementCtransAtransBtransCuseNzAuseNzBuseNzCMatmulParams表示运行期参数mnkinputAddroutputAddrkernel entry 签名保持固定void BasicMatmul( uint32_t blockNum, aclrtStream stream, const MatmulTParams tParams, const MatmulParams params);6.2 扩展 ABIcatlass_kernel.h中还保留了 grouped matmul、quant matmul、conv、flash attention 等参数结构和函数声明。新增算子时优先复用已有 ABI 结构当参数语义明显不同再新增独立结构。7. Utils 模块utils/将工具函数拆成两个 targettarget文件依赖用途catlass_kernel_utilskernel_utils.cppACLJIT compiler 使用的 dtype 到 bisheng type 转换catlass_torch_utilstorch_utils.cppACL、torch、torch-npuPyTorch wrapper 使用的 Tensor/dtype/layout 工具7.1 dtype 映射utils/type_utils.hpp维护 dtype 映射表canonical string nametorch::DtypeaclDataTypebisheng C type token映射表按 index 对齐TypeCastS, T()通过查表完成不同表示之间的转换。部分 dtype 在 torch-npu 随包 ACL 头中没有枚举名但 ABI 数值与 CANN 定义一致。此类 dtype 使用static_castaclDataType(value)表达避免混用系统 CANN ACL 头和 torch-npu 随包 ACL 头导致重复定义。7.2 Tensor 工具torch_utils.cpp提供GetOutputTensor()在当前 NPU device 上创建 ND 格式输出 Tensor。TypeStrToTorchDtype()字符串到 torch dtype。TorchDtypeToAclDtype()torch dtype 到 ACL dtype。AclDtypeToTorchDtype()ACL dtype 到 torch dtype。GetTransposeStatus()根据 tensor stride 和 NPU format 判断矩阵布局。8. JIT 子系统JIT 子系统由四部分组成组件文件责任JIT entrykernels/00_basic_matmul/basic_matmul.cpp稳定 kernel 入口负责获取 JIT 函数并调用JIT templatekernels/00_basic_matmul/basic_matmul_impl.cpp被运行时编译的 CATLASS kernel 模板JIT compilerkernels/jit/jit_compiler.cpp编译、缓存、加载.somacro generatorkernels/include/jit_macro_generator.h将模板参数转为-D宏8.1 JIT entryJIT entry 固定编译进libcatlass_kernel_jit.so。以BasicMatmul为例auto* entry JitCompiler::instance().getKernel( basic_matmul_impl.cpp, JitMacroGeneratorMatmulTParams::generate(basic_matmul, tParams)); entry(blockNum, stream, tParams, params);entry 的职责是连接 stable ABI 和 runtime-compiled template不承载具体 GEMM 模板逻辑。8.2 JIT templateJIT template 使用宏注入类型和布局宏语义CATLASS_JIT_ELEMENT_AA 元素类型CATLASS_JIT_ELEMENT_BB 元素类型CATLASS_JIT_ELEMENT_CC 元素类型CATLASS_JIT_LAYOUT_AA layoutCATLASS_JIT_LAYOUT_BB layoutCATLASS_JIT_LAYOUT_CC layoutCATLASS_JIT_KERNEL_NAMEdevice kernel 符号名模板导出稳定 C ABIextern C void run( uint32_t blockNum, aclrtStream stream, const MatmulTParams* tParams, const MatmulParams* params);JIT loader 固定解析run符号。device kernel 名只用于编译产物可读性和 profiling。8.3 宏生成JitMacroGeneratorTParams是模板策略类。默认模板不生成任何宏具体参数类型通过特化实现。JitMacroGeneratorMatmulTParams生成CATLASS_KERNEL_NAMECATLASS_JIT_ELEMENT_ACATLASS_JIT_ELEMENT_BCATLASS_JIT_ELEMENT_CCATLASS_JIT_LAYOUT_ACATLASS_JIT_LAYOUT_BCATLASS_JIT_LAYOUT_CCATLASS_JIT_KERNEL_NAME新增非 matmul JIT kernel 时应新增对应参数结构的JitMacroGenerator特化。8.4 编译和缓存JitCompiler是进程级单例。初始化内容包括JIT cache 目录。bisheng/ccec 路径。当前 NPU arch。JIT template 根目录。缓存分两层内存缓存loaded_保存SharedLib和run指针。磁盘缓存保存编译后的.so。cache key 由 kernel 名、arch 和完整宏集合构成CATLASS_KERNEL_NAME_archarch_macro-name_macro-value...宏按 key 排序后拼接保证unordered_map遍历顺序不影响缓存路径。8.5 环境变量环境变量作用CATLASS_JIT_LOG_LEVELJIT 日志等级0/1/2TORCH_CATLASS_CACHE_DIRJIT 磁盘缓存目录MS_SANITIZE_MEMORY启用 msSanitizer 编译选项TORCH_CATLASS_VERSION注入 package/CATLASS 版本ASCEND_HOME_PATH查找 Ascend compiler 和 runtime 库TORCH_CATLASS_PKG_DIR安装后的 Python 包根目录JIT 编译使用的 NPU arch 只通过 AscendC platform API 获取即GetCurrentNPUArch()。运行时不支持通过环境变量覆盖 arch避免编译参数和实际设备不一致。环境变量分为外部配置和包内注入两类。外部配置只保留ASCEND_HOME_PATH、TORCH_CATLASS_CACHE_DIR、CATLASS_JIT_LOG_LEVEL和MS_SANITIZE_MEMORYTORCH_CATLASS_VERSION和TORCH_CATLASS_PKG_DIR由 Python loader 在 import 时设置。JIT template 路径固定由TORCH_CATLASS_PKG_DIR/jit/templates/推导compiler 优先从ASCEND_HOME_PATH的标准目录查找未找到时回退到PATH中的ccec。9. Kernel 构建模块kernels/CMakeLists.txt提供add_kernel()统一 JIT 和 prebuilt kernel 的构建入口。9.1 JIT kerneladd_kernel( NAME basic_matmul NPU_ARCH_LIST 2201 KERNEL_TYPE jit ${CMAKE_CURRENT_SOURCE_DIR}/basic_matmul.cpp TEMPLATE ${CMAKE_CURRENT_SOURCE_DIR}/basic_matmul_impl.cpp)JIT kernel 构建流程entry 源文件加入统一libcatlass_kernel_jit.so。template 文件安装到jit/templates/。jit_try_compile_check()在构建期检查 template 可被 bisheng 编译。运行时由JitCompiler根据模板参数编译具体.so。9.2 prebuilt kernelprebuilt kernel 按 arch 构建独立动态库lib/arch/libcatlass_kernel_arch_name.soprebuilt 模式用于固定参数组合或无需运行时模板编译的 kernel。启用 msSanitizer 时可额外构建_ms变体。10. 顶层构建模块10.1 Python 构建入口build.sh是主要构建入口推导 package 版本。写入torch_catlass/_version.py。自动探测当前 Python 环境中的 Torch CMake 目录。调用 pip/scikit-build-core 驱动 CMake 构建。常用命令bash build.sh --skip-wheel bash build.sh --build-type Debug --skip-wheel bash build.sh --clean10.2 CMake target顶层CMakeLists.txt负责查找 ASC、Python、Torch。设置 C17、PIC、compile commands。安装 public headers。安装 CATLASS headers 到 Python 包内的 JIT include tree。添加kernels、utils、src子目录。关键 targettarget输出说明catlass_kernel_utilsstatic libJIT compiler 依赖的纯 ACL 工具catlass_torch_utilsstatic libtorch wrapper 依赖的 Tensor 工具catlass_kernel_jit_compilershared libJIT 编译器catlass_kernel_jitshared libJIT entry 集合catlass_torchshared libPyTorch extension11. 测试模块pytest 集成测试验证 Python API 到 kernel 执行的完整链路。tests/test_00_basic_matmul.py测试流程检查是否存在可用 Ascend NPU无设备时跳过。构造 NPU fp16 输入。调用torch_catlass.basic_matmul()。用torch.matmul()生成参考结果。校验 shape、dtype、device。用torch.allclose()校验数值。本地静态检查python3 -m py_compile torch_catlass/__init__.py torch_catlass/ops/basic_matmul.py tests/test_00_basic_matmul.py python3 -m ruff check torch_catlass tests完整集成测试bash build.sh --skip-wheel pytest tests/test_00_basic_matmul.py -v -s12. 扩展流程12.1 新增 matmul 类 JIT 算子在kernels/nn_name/下添加 entry.cpp和 template.cpp。在 entry 中调用JitCompiler::instance().getKernel()。复用或扩展JitMacroGeneratorMatmulTParams。在kernels/nn_name/CMakeLists.txt中调用add_kernel()。在include/catlass_kernel.h声明 kernel entry。在src/catlass_torch.cpp使用MatmulLikeKernel注册 torch op。在torch_catlass/ops/添加 Python wrapper。在tests/添加 pytest与 PyTorch 参考实现比对。12.2 新增非 matmul 算子非 matmul 算子应新增独立 adapter而不是扩展MatmulLikesrc/include/template/op_family.h ├─ GetKernelInfo() ├─ AllocOutput() └─ Run()同时新增对应的参数结构和JitMacroGenerator特化保持参数解析、宏生成和 kernel ABI 各自独立。12.3 新增 dtype 或 layout新增 dtype/layout 时需要同步更新utils/type_utils.hppJitMacroGenerator对应特化JIT template 中的默认宏和类型别名pytest 参数覆盖dtype 映射应优先使用当前编译环境可见的 enum当 torch-npu 随包 ACL 头未暴露 enum 名但 ABI 数值稳定时可使用static_castaclDataType(value)保持兼容。【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
torch-catlass 测试框架设计文档
torch-catlass 测试框架设计文档1. 总览tests/optest是 CATLASS 示例算子接入 PyTorch 的端到端测试框架。框架将 CATLASS AscendC kernel 封装为torch.ops.catlass.*算子并通过 Python 包torch_catlass提供测试入口。框架按职责分为五层Python API torch_catlass.ops.* | v Python package loader load kernel libs and libcatlass_torch.so | v PyTorch C extension register torch.ops.catlass.* | v Kernel adapter convert Tensor arguments to CATLASS kernel params | v Kernel implementation prebuilt kernel or JIT compiled template核心原则Python 层只负责用户接口、动态库加载和轻量参数转换。C extension 层负责 PyTorch op 注册和 NPU dispatch。adapter 层负责 Tensor 到 kernel ABI 的转换。kernel 层负责 AscendC/CATLASS 代码执行。JIT 子系统负责模板参数宏生成、编译、缓存和动态加载。2. 目录模块tests/optest/ ├── pyproject.toml ├── CMakeLists.txt ├── build.sh ├── docs/ │ └── design.md ├── include/ │ ├── catlass_kernel.h │ └── catlass_torch.h ├── torch_catlass/ │ ├── __init__.py │ ├── _version.py │ └── ops/ ├── src/ │ ├── catlass_torch.cpp │ ├── common/ │ └── include/ ├── utils/ │ ├── CMakeLists.txt │ ├── include/ │ ├── kernel_utils.cpp │ ├── torch_utils.cpp │ └── type_utils.hpp ├── kernels/ │ ├── CMakeLists.txt │ ├── common/ │ ├── include/ │ ├── jit/ │ └── 00_basic_matmul/ └── tests/ └── test_00_basic_matmul.py模块责任torch_catlass/Python 包入口、动态库加载、用户侧 op wrapperinclude/框架和 kernel 共享的公共 ABI 声明src/PyTorch C extension 和 torch op 注册utils/dtype/layout/Tensor 工具函数拆分 torch 依赖和纯 ACL 依赖kernels/kernel 构建、JIT compiler、JIT template、kernel entrytests/pytest 集成测试3. Python 包模块3.1 包初始化torch_catlass/__init__.py在 import 时完成运行时初始化从_version.py读取构建期版本信息。设置TORCH_CATLASS_VERSION供 JIT 编译阶段注入版本宏。设置TORCH_CATLASS_PKG_DIR供 JIT compiler 定位安装后的 headers 和 templates。加载 JIT compiler 和 JIT kernel entry 库。根据当前 NPU 架构加载 arch-specific kernel 库。调用torch.ops.load_library()加载libcatlass_torch.so注册torch.ops.catlass.*。动态库加载顺序lib/jit/libcatlass_kernel_jit_compiler.so lib/jit/libcatlass_kernel_jit.so lib/arch/*.so lib/libcatlass_torch.soJIT compiler 和 kernel entry 必须先于 PyTorch extension 加载保证 extension 中引用的 kernel 符号可解析。3.2 架构识别Python loader 通过torch_npu.npu.get_device_name()识别设备并映射为 CATLASS arch id设备名archAscend910B.*2201Ascend910_932201Ascend950PR/Ascend950DT3510当torch_npu.npu.device_count()为 0 时loader 抛出明确错误。测试代码在无 NPU 环境下通过 pytest skip 处理避免在 collection 阶段触发 torch-npu 内部错误。3.3 Python op wrappertorch_catlass/ops/保存 Python 用户接口。以basic_matmul为例torch_catlass.basic_matmul( mat1, mat2, outDTypefloat16, transAFalse, transBFalse, formatAFalse, formatBFalse, )Python wrapper 只做用户友好的轻量转换例如 dtype 字符串白名单解析。shape 推导、输出分配、stream 获取和 kernel launch 都在 C 层完成避免 Python 和 C 维护两套语义。4. PyTorch C Extension 模块4.1 注册入口src/catlass_torch.cpp是 PyTorch extension 的注册入口using BasicMatmulOp MatmulLikeCatlassKernel::BasicMatmul; static auto basic_matmul BasicMatmulOp::Run; REGISTER_TORCH_FUNC(basic_matmul);REGISTER_TORCH_FUNC位于src/include/common/register.h注册流程为创建或复用torch::Librarynamespace 固定为catlass。通过 PyTorch schema inference 从 C 函数签名生成 schema。将实现注册到c10::DispatchKey::PrivateUse1。PrivateUse1是 torch-npu 使用的 NPU dispatch key。Python 调用torch.ops.catlass.basic_matmul时PyTorch 根据输入 Tensor device 走到该 backend 实现。4.2 kernel launch 包装RUN_NPU_FUNC位于src/include/common/run_npu_func.h。它通过 torch-npu 的OpCommand::RunOpApiV2执行 kernel launchlaunch 前检查函数指针是否为空。将 C 异常转为 ACL error code。将 kernel 调用交给 torch-npu runtime 管理。5. Matmul Adapter 模块src/include/template/matmul.h提供MatmulLikeKernelFunc。该模板封装 matmul 类算子的通用流程Run() ├─ GetKernelInfo() ├─ AllocOutput() ├─ get current NPU stream ├─ get AIC core count └─ RUN_NPU_FUNC(KernelFunc, ...)5.1 参数拆分Matmul 参数拆分为两类参数结构用途是否参与 JIT 编译MatmulTParamsdtype、layout、transpose 等模板参数是MatmulParamsM/N/K、输入输出地址等运行时参数否这种拆分保证 dtype/layout 变化会生成新的模板实例而 shape 和 Tensor 地址变化不会导致重复编译。5.2GetKernelInfoGetKernelInfo()负责将 PyTorch Tensor 转为 kernel 参数将torch::Dtype转为aclDataType。根据transA和transB推导m/n/k。检查两个输入矩阵的 K 维一致。填充MatmulTParams。填充MatmulParams。将输入 Tensor storage 地址写入params.inputAddr。5.3AllocOutputAllocOutput()根据params.m、params.n和tParams.elementC创建输出 Tensor并将其 storage 地址写入params.outputAddr[0]。输出 Tensor 生命周期由 PyTorch 管理kernel ABI 只接收裸地址。6. 公共 ABI 模块include/catlass_kernel.h定义 C wrapper 和 kernel 实现之间共享的数据结构和函数声明。6.1 matmul ABIMatmulTParams表示编译期参数elementAelementBelementCtransAtransBtransCuseNzAuseNzBuseNzCMatmulParams表示运行期参数mnkinputAddroutputAddrkernel entry 签名保持固定void BasicMatmul( uint32_t blockNum, aclrtStream stream, const MatmulTParams tParams, const MatmulParams params);6.2 扩展 ABIcatlass_kernel.h中还保留了 grouped matmul、quant matmul、conv、flash attention 等参数结构和函数声明。新增算子时优先复用已有 ABI 结构当参数语义明显不同再新增独立结构。7. Utils 模块utils/将工具函数拆成两个 targettarget文件依赖用途catlass_kernel_utilskernel_utils.cppACLJIT compiler 使用的 dtype 到 bisheng type 转换catlass_torch_utilstorch_utils.cppACL、torch、torch-npuPyTorch wrapper 使用的 Tensor/dtype/layout 工具7.1 dtype 映射utils/type_utils.hpp维护 dtype 映射表canonical string nametorch::DtypeaclDataTypebisheng C type token映射表按 index 对齐TypeCastS, T()通过查表完成不同表示之间的转换。部分 dtype 在 torch-npu 随包 ACL 头中没有枚举名但 ABI 数值与 CANN 定义一致。此类 dtype 使用static_castaclDataType(value)表达避免混用系统 CANN ACL 头和 torch-npu 随包 ACL 头导致重复定义。7.2 Tensor 工具torch_utils.cpp提供GetOutputTensor()在当前 NPU device 上创建 ND 格式输出 Tensor。TypeStrToTorchDtype()字符串到 torch dtype。TorchDtypeToAclDtype()torch dtype 到 ACL dtype。AclDtypeToTorchDtype()ACL dtype 到 torch dtype。GetTransposeStatus()根据 tensor stride 和 NPU format 判断矩阵布局。8. JIT 子系统JIT 子系统由四部分组成组件文件责任JIT entrykernels/00_basic_matmul/basic_matmul.cpp稳定 kernel 入口负责获取 JIT 函数并调用JIT templatekernels/00_basic_matmul/basic_matmul_impl.cpp被运行时编译的 CATLASS kernel 模板JIT compilerkernels/jit/jit_compiler.cpp编译、缓存、加载.somacro generatorkernels/include/jit_macro_generator.h将模板参数转为-D宏8.1 JIT entryJIT entry 固定编译进libcatlass_kernel_jit.so。以BasicMatmul为例auto* entry JitCompiler::instance().getKernel( basic_matmul_impl.cpp, JitMacroGeneratorMatmulTParams::generate(basic_matmul, tParams)); entry(blockNum, stream, tParams, params);entry 的职责是连接 stable ABI 和 runtime-compiled template不承载具体 GEMM 模板逻辑。8.2 JIT templateJIT template 使用宏注入类型和布局宏语义CATLASS_JIT_ELEMENT_AA 元素类型CATLASS_JIT_ELEMENT_BB 元素类型CATLASS_JIT_ELEMENT_CC 元素类型CATLASS_JIT_LAYOUT_AA layoutCATLASS_JIT_LAYOUT_BB layoutCATLASS_JIT_LAYOUT_CC layoutCATLASS_JIT_KERNEL_NAMEdevice kernel 符号名模板导出稳定 C ABIextern C void run( uint32_t blockNum, aclrtStream stream, const MatmulTParams* tParams, const MatmulParams* params);JIT loader 固定解析run符号。device kernel 名只用于编译产物可读性和 profiling。8.3 宏生成JitMacroGeneratorTParams是模板策略类。默认模板不生成任何宏具体参数类型通过特化实现。JitMacroGeneratorMatmulTParams生成CATLASS_KERNEL_NAMECATLASS_JIT_ELEMENT_ACATLASS_JIT_ELEMENT_BCATLASS_JIT_ELEMENT_CCATLASS_JIT_LAYOUT_ACATLASS_JIT_LAYOUT_BCATLASS_JIT_LAYOUT_CCATLASS_JIT_KERNEL_NAME新增非 matmul JIT kernel 时应新增对应参数结构的JitMacroGenerator特化。8.4 编译和缓存JitCompiler是进程级单例。初始化内容包括JIT cache 目录。bisheng/ccec 路径。当前 NPU arch。JIT template 根目录。缓存分两层内存缓存loaded_保存SharedLib和run指针。磁盘缓存保存编译后的.so。cache key 由 kernel 名、arch 和完整宏集合构成CATLASS_KERNEL_NAME_archarch_macro-name_macro-value...宏按 key 排序后拼接保证unordered_map遍历顺序不影响缓存路径。8.5 环境变量环境变量作用CATLASS_JIT_LOG_LEVELJIT 日志等级0/1/2TORCH_CATLASS_CACHE_DIRJIT 磁盘缓存目录MS_SANITIZE_MEMORY启用 msSanitizer 编译选项TORCH_CATLASS_VERSION注入 package/CATLASS 版本ASCEND_HOME_PATH查找 Ascend compiler 和 runtime 库TORCH_CATLASS_PKG_DIR安装后的 Python 包根目录JIT 编译使用的 NPU arch 只通过 AscendC platform API 获取即GetCurrentNPUArch()。运行时不支持通过环境变量覆盖 arch避免编译参数和实际设备不一致。环境变量分为外部配置和包内注入两类。外部配置只保留ASCEND_HOME_PATH、TORCH_CATLASS_CACHE_DIR、CATLASS_JIT_LOG_LEVEL和MS_SANITIZE_MEMORYTORCH_CATLASS_VERSION和TORCH_CATLASS_PKG_DIR由 Python loader 在 import 时设置。JIT template 路径固定由TORCH_CATLASS_PKG_DIR/jit/templates/推导compiler 优先从ASCEND_HOME_PATH的标准目录查找未找到时回退到PATH中的ccec。9. Kernel 构建模块kernels/CMakeLists.txt提供add_kernel()统一 JIT 和 prebuilt kernel 的构建入口。9.1 JIT kerneladd_kernel( NAME basic_matmul NPU_ARCH_LIST 2201 KERNEL_TYPE jit ${CMAKE_CURRENT_SOURCE_DIR}/basic_matmul.cpp TEMPLATE ${CMAKE_CURRENT_SOURCE_DIR}/basic_matmul_impl.cpp)JIT kernel 构建流程entry 源文件加入统一libcatlass_kernel_jit.so。template 文件安装到jit/templates/。jit_try_compile_check()在构建期检查 template 可被 bisheng 编译。运行时由JitCompiler根据模板参数编译具体.so。9.2 prebuilt kernelprebuilt kernel 按 arch 构建独立动态库lib/arch/libcatlass_kernel_arch_name.soprebuilt 模式用于固定参数组合或无需运行时模板编译的 kernel。启用 msSanitizer 时可额外构建_ms变体。10. 顶层构建模块10.1 Python 构建入口build.sh是主要构建入口推导 package 版本。写入torch_catlass/_version.py。自动探测当前 Python 环境中的 Torch CMake 目录。调用 pip/scikit-build-core 驱动 CMake 构建。常用命令bash build.sh --skip-wheel bash build.sh --build-type Debug --skip-wheel bash build.sh --clean10.2 CMake target顶层CMakeLists.txt负责查找 ASC、Python、Torch。设置 C17、PIC、compile commands。安装 public headers。安装 CATLASS headers 到 Python 包内的 JIT include tree。添加kernels、utils、src子目录。关键 targettarget输出说明catlass_kernel_utilsstatic libJIT compiler 依赖的纯 ACL 工具catlass_torch_utilsstatic libtorch wrapper 依赖的 Tensor 工具catlass_kernel_jit_compilershared libJIT 编译器catlass_kernel_jitshared libJIT entry 集合catlass_torchshared libPyTorch extension11. 测试模块pytest 集成测试验证 Python API 到 kernel 执行的完整链路。tests/test_00_basic_matmul.py测试流程检查是否存在可用 Ascend NPU无设备时跳过。构造 NPU fp16 输入。调用torch_catlass.basic_matmul()。用torch.matmul()生成参考结果。校验 shape、dtype、device。用torch.allclose()校验数值。本地静态检查python3 -m py_compile torch_catlass/__init__.py torch_catlass/ops/basic_matmul.py tests/test_00_basic_matmul.py python3 -m ruff check torch_catlass tests完整集成测试bash build.sh --skip-wheel pytest tests/test_00_basic_matmul.py -v -s12. 扩展流程12.1 新增 matmul 类 JIT 算子在kernels/nn_name/下添加 entry.cpp和 template.cpp。在 entry 中调用JitCompiler::instance().getKernel()。复用或扩展JitMacroGeneratorMatmulTParams。在kernels/nn_name/CMakeLists.txt中调用add_kernel()。在include/catlass_kernel.h声明 kernel entry。在src/catlass_torch.cpp使用MatmulLikeKernel注册 torch op。在torch_catlass/ops/添加 Python wrapper。在tests/添加 pytest与 PyTorch 参考实现比对。12.2 新增非 matmul 算子非 matmul 算子应新增独立 adapter而不是扩展MatmulLikesrc/include/template/op_family.h ├─ GetKernelInfo() ├─ AllocOutput() └─ Run()同时新增对应的参数结构和JitMacroGenerator特化保持参数解析、宏生成和 kernel ABI 各自独立。12.3 新增 dtype 或 layout新增 dtype/layout 时需要同步更新utils/type_utils.hppJitMacroGenerator对应特化JIT template 中的默认宏和类型别名pytest 参数覆盖dtype 映射应优先使用当前编译环境可见的 enum当 torch-npu 随包 ACL 头未暴露 enum 名但 ABI 数值稳定时可使用static_castaclDataType(value)保持兼容。【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
