catlass_cppgen【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass1. 项目介绍catlass_cppgen是一个基于 Python 的代码生成框架用于构建和生成 CATLASS 高性能算子。该框架提供了灵活的接口来定义算子参数、选择优化策略并自动生成对应的 C 核函数代码。1.1 主要特性算子代码生成通过 Python API 定义算子参数自动生成优化的 C 核函数代码灵活的调优接口支持自定义TileShape、DispatchPolicy等参数多架构支持支持多种硬件架构包括AtlasA2/A3, Ascend950类型安全提供完整的数据类型和布局抽象1.2 工程结构以下是本项目目录结构说明./catlass_cppgen ├── catlass_cppgen │ ├── __init__.py │ ├── catlass # CATLASS 相关特性组件 │ │ ├── __init__.py │ │ ├── arch # 代际声明 │ │ ├── evg # EVG 特性承载 │ │ ├── evg_extension.py │ │ ├── gemm │ │ ├── gemm_coord.py │ │ ├── layout │ │ └── library.py │ ├── common # 通用组件 │ │ ├── __init__.py │ │ ├── data_type.py │ │ ├── op_tensor.py │ │ ├── typing.py │ │ └── utils.py │ ├── kernel │ │ ├── __init__.py │ │ ├── gemm # GEMM 类算子特化类 │ │ ├── group_gemm # Group GEMM 类算子特化类 │ │ ├── kernel_base.py │ │ └── visitor_kernel_base.py │ ├── op # 算子Kernel基类 │ │ ├── __init__.py │ │ ├── gemm.py │ │ ├── group_gemm.py │ │ └── op.py │ └── _version.py ├── docs # API 文档 │ ├── evg_api.md │ ├── kernel_api.md │ └── optensor_api.md ├── tests # 单元测试组件 │ ├── catlass # 面向 CATLASS 相关特性的测试件 │ ├── common # 面向通用组件的测试件(类型、排布) │ └── op # 面向算子kernel生成的测试件 ├── pyproject.toml ├── README.md # 主README文档 └── uv.lock2. 支持的算子2.1 GEMM 类矩阵乘法算子类型Kernel 类主要特性切分轴基础矩阵乘法BasicMatmulKernel• 输入张量 A 和 B 为 2 维•alpha 1.0且beta 0.0• 支持可选的 Bias 参数无批处理矩阵乘法BatchedMatmulKernel• 输入张量 A 和 B 为 3 维batchCount, M, K和batchCount, K, N• 所有批次共享相同的矩阵维度•alpha 1.0且beta 0.0无EVG Visitor 矩阵乘法BasicMatmulTlaVisitorKernel• 支持CATLASS模板库后处理框架 EVGEpilogue Visitor Graph无多核 Split-KMultiCoreSplitkMatmulKernel• 输入张量 A 和 B 为 2 维• 优化动作沿 K 方向多核切分• 支持可选的 Bias 参数K尾块多核 Split-KTailMultiCoreSplitkMatmulKernel• 输入张量 A 和 B 为 2 维• 多核切K的尾块优化变体• 支持可选的 Bias 参数KStream-KStreamkMatmulKernel• 输入张量 A 和 B 为 2 维• 优化动作Stream-K 调度策略• 支持可选的 Bias 参数K2.2 Group GEMM 类分组矩阵乘法算子类型Kernel 类主要特性切分轴分组矩阵乘M 轴切分GroupedMatmulSliceMKernel多组不同 M 维度的矩阵乘法M2.3 EVG 后处理 (Epilogue Visitor Graph)支持通过EVG(Epilogue Visitor Graph)框架实现后处理功能支持的后处理类别包括单一计算环节可通过运算符或函数调用表达以下算子 | 类别 | 算子 | 写法示例 | |------|------|----------| | 二元运算 | add |accum bias| | 二元运算 | sub |accum - bias| | 二元运算 | mul |accum * scale| | 二元运算 | div |accum / scale| | 激活函数 | relu |relu(accum)| | 激活函数 | leakyRelu |leakyRelu(accum, alpha)| | 激活函数 | Prelu |Prelu(accum, weight)| | 激活函数 | sigmoid |sigmoid(accum)| | 激活函数 | silu |silu(accum)| | 比较/选择 | max / min |max(a, b)/min(a, b)| | 类型转换 | cast |cast(accum, float16, float)| | 常量 | constant |constant(1.0, float)|组合计算支持多个计算节点拼接广播计算支持行广播计算。3. 安装3.1 从源码安装构建分发包pip install build python -m build这会在dist/目录下生成.whl和.tar.gz文件。安装分发包pip install dist/catlass_cppgen-*.whl或者pip install dist/catlass_cppgen-*.tar.gz直接安装开发模式pip install -e .3.2 从本地目录安装如果您想直接从项目目录安装pip install .4. 使用示例当前catlass_cppgen支持 matmul、grouped_matmul 以及 EVG 后处理特性的代码生成。以下是应用cppgen的环节示意Gemm / GroupGemm算子规划 ↓ get_kernels() Kernel 对象调优与特性查询 ↓ tune() / to_evg() 配置完成的 Kernel详细参考使用示例和 API 文档请参考下述文档Kernel API 基础文档OpTensor API 基础文档EVG API 基础文档4.1 基础 GEMM以下是一个基础的创建 matmul 算子cppgen对象的示例from catlass_cppgen.op.gemm import Gemm from catlass_cppgen.common.op_tensor import OpTensor from catlass_cppgen.common.data_type import DataType from catlass_cppgen.catlass.layout.layout import RowMajor from catlass_cppgen.catlass.gemm_coord import GemmShape from catlass_cppgen.catlass.arch.arch import Arch from catlass_cppgen.catlass.gemm.dispatch_policy import MmadPingpong # 1. 描述输入张量无需绑定底层数据 a OpTensor.from_shape_stride((128, 256), (256, 1), DataType.FLOAT) b OpTensor.from_shape_stride((256, 384), (384, 1), DataType.FLOAT) # 2. 创建算子并获取 Kernel gemm Gemm(atlas_archArch.Ascend950, elementDataType.FLOAT, layoutRowMajor, Aa, Bb) kernels gemm.get_kernels() # 3. [可选] 可以显式指定使用的Kernel组件以BasicMatmulKernel为例 from catlass_cppgen.kernel.gemm import BasicMatmulKernel kernel find_kernel_by_type(kernels, BasicMatmulKernel) # 非定向指定: kernel kernels[0] # 4. [可选] 调优 Tile 形状与调度策略 kernel.tune( GemmShape(128, 256, 64), GemmShape(128, 256, 64), dispatch_policyMmadPingpong(arch_tagArch.Ascend950), ) # 5. [可选] matmul hello world: 核函数的代码生成 print(f[Kernel] \n{kernel.gen_kernel_template()})得到kernel对象后可以调用gen_kernel_template()gen_params_device()等方法针对核函数和参数绑定做代码生成。4.2 Group GEMM以下是建立 matmul 算子cppgen对象的示例from catlass_cppgen.op.group_gemm import GroupGemm from catlass_cppgen.catlass.layout.layout import VectorLayout # ... # 1. 创建groupList 张量 groupList OpTensor(dtypeDataType.INT64, layoutVectorLayout(4), shape(4,)) # 2. 建立 group matmul 对象 group_gemm GroupGemm(atlas_archArch.Ascend950, Aa, Bb_3d, groupListgroupList) # 3. [可选] 取得kernel并做 Tiling 调优 kernels group_gemm.get_kernels() kernels[0].tune(GemmShape(256, 256, 256), GemmShape(256, 256, 64))4.3 EVG 后处理# ... # 1. 进行evg对象声明 # - fn_src: EVG对象函数头 # - example_inputs: 后处理过程中涉及的name:tensor键值对 evg_config { fn_src: def epilogue(accum, bias):\n return relu(accum bias), example_inputs: { accum: OpTensor.from_shape_stride((128, 256), (256, 1), DataType.FLOAT), bias: OpTensor.from_shape_stride((1, 256), (256, 1), DataType.FLOAT), result: OpTensor.from_shape_stride((128, 256), (256, 1), DataType.FLOAT), }, } # 2. 创建 matmul 算子对象并获取 Kernel gemm Gemm(atlas_archArch.Ascend950, evg_configevg_config, Aa, Bb) kernel gemm.get_kernels()[0] # is_support_evg 特性为 True assert kernel.is_support_evg # [可选] 3. 进行 Tile 形状调优 kernel.tune(GemmShape(128, 256, 64), GemmShape(128, 256, 64))支持二元运算add/sub/mul/div、激活函数relu/silu/sigmoid/leakyRelu/prelu、类型转换cast、常量constant等可串联组合使用。【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考