补充cppgen框架-catlass-AtomGit

Apython/catlass_cppgen/README.md+238-0

@@ -0,0 +1,238 @@
	1	+# catlass_cppgen
	CheaterAbec7月6日参考wiki详细说明一下提供的接口和特性 I init__zhb__ 7月8日评论： 7月8日评论：
	2	+
	3	+## 1. 项目介绍
	4	+
	5	+`catlass_cppgen` 是一个基于 Python 的代码生成框架，用于构建和生成 CATLASS 高性能算子。该框架提供了灵活的接口来定义算子参数、选择优化策略，并自动生成对应的 C++ 核函数代码。
	6	+
	7	+### 1.1 主要特性
	8	+
	9	+- 算子代码生成：通过 Python API 定义算子参数，自动生成优化的 C++ 核函数代码
	10	+- 灵活的调优接口：支持自定义TileShape、DispatchPolicy等参数
	11	+- 多架构支持：支持多种硬件架构（包括AtlasA2/A3, Ascend950）
	12	+- 类型安全：提供完整的数据类型和布局抽象
	13	+
	14	+### 1.2 工程结构
	15	+
	16	+以下是本项目目录结构说明：
	17	+
	18	+```plain
	19	+./catlass_cppgen
	20	+├── catlass_cppgen
	21	+│ ├── __init__.py
	22	+│ ├── catlass # CATLASS 相关特性组件
	23	+│ │ ├── __init__.py
	24	+│ │ ├── arch # 代际声明
	25	+│ │ ├── evg # EVG 特性承载
	26	+│ │ ├── evg_extension.py
	27	+│ │ ├── gemm
	28	+│ │ ├── gemm_coord.py
	29	+│ │ ├── layout
	30	+│ │ └── library.py
	31	+│ ├── common # 通用组件
	32	+│ │ ├── __init__.py
	33	+│ │ ├── data_type.py
	34	+│ │ ├── op_tensor.py
	35	+│ │ ├── typing.py
	36	+│ │ └── utils.py
	37	+│ ├── kernel
	38	+│ │ ├── __init__.py
	39	+│ │ ├── gemm # GEMM 类算子特化类
	40	+│ │ ├── group_gemm # Group GEMM 类算子特化类
	41	+│ │ ├── kernel_base.py
	42	+│ │ └── visitor_kernel_base.py
	43	+│ ├── op # 算子Kernel基类
	44	+│ │ ├── __init__.py
	45	+│ │ ├── gemm.py
	46	+│ │ ├── group_gemm.py
	47	+│ │ └── op.py
	48	+│ └── _version.py
	49	+├── docs # API 文档
	50	+│ ├── evg_api.md
	51	+│ ├── kernel_api.md
	52	+│ └── optensor_api.md
	53	+├── tests # 单元测试组件
	54	+│ ├── catlass # 面向 CATLASS 相关特性的测试件
	55	+│ ├── common # 面向通用组件的测试件(类型、排布)
	56	+│ └── op # 面向算子kernel生成的测试件
	57	+├── pyproject.toml
	58	+├── README.md # 主README文档
	59	+└── uv.lock
	60	+```
	61	+
	62	+## 2. 支持的算子
	63	+
	64	+### 2.1 GEMM 类（矩阵乘法）
	65	+
	66	+\| 算子类型 \| Kernel 类 \| 主要特性 \| 切分轴 \|
	67	+\|---------\|----------\|---------\|--------\|
	68	+\| 基础矩阵乘法 \| `BasicMatmulKernel` \| • 输入张量 A 和 B 为 2 维<br>• `alpha = 1.0` 且 `beta = 0.0`<br>• 支持可选的 Bias 参数 \| 无 \|
	69	+\| 批处理矩阵乘法 \| `BatchedMatmulKernel` \| • 输入张量 A 和 B 为 3 维（batchCount, M, K）和（batchCount, K, N）<br>• 所有批次共享相同的矩阵维度<br>• `alpha = 1.0` 且 `beta = 0.0` \| 无 \|
	70	+\| EVG Visitor 矩阵乘法 \| `BasicMatmulTlaVisitorKernel` \| • 支持CATLASS模板库后处理框架 EVG（Epilogue Visitor Graph） \| 无 \|
	71	+\| 多核 Split-K \| `MultiCoreSplitkMatmulKernel` \| • 输入张量 A 和 B 为 2 维<br>• 优化动作：沿 K 方向多核切分<br>• 支持可选的 Bias 参数 \| K \|
	72	+\| 尾块多核 Split-K \| `TailMultiCoreSplitkMatmulKernel` \| • 输入张量 A 和 B 为 2 维<br>• 多核切K的尾块优化变体<br>• 支持可选的 Bias 参数 \| K \|
	73	+\| Stream-K \| `StreamkMatmulKernel` \| • 输入张量 A 和 B 为 2 维<br>• 优化动作：Stream-K 调度策略<br>• 支持可选的 Bias 参数 \| K \|
	74	+
	75	+### 2.2 Group GEMM 类（分组矩阵乘法）
	76	+
	77	+\| 算子类型 \| Kernel 类 \| 主要特性 \| 切分轴 \|
	78	+\|---------\|----------\|---------\|--------\|
	79	+\| 分组矩阵乘（M 轴切分） \| `GroupedMatmulSliceMKernel` \| 多组不同 M 维度的矩阵乘法 \| M \|
	80	+
	81	+### 2.3 EVG 后处理 (Epilogue Visitor Graph)
	82	+
	83	+支持通过EVG(Epilogue Visitor Graph)框架实现后处理功能，支持的后处理类别包括：
	84	+ - 单一计算环节：可通过运算符或函数调用表达以下算子：
	85	+ \| 类别 \| 算子 \| 写法示例 \|
	86	+ \|------\|------\|----------\|
	87	+ \| 二元运算 \| add \| `accum + bias` \|
	88	+ \| 二元运算 \| sub \| `accum - bias` \|
	89	+ \| 二元运算 \| mul \| `accum * scale` \|
	90	+ \| 二元运算 \| div \| `accum / scale` \|
	91	+ \| 激活函数 \| relu \| `relu(accum)` \|
	92	+ \| 激活函数 \| leakyRelu \| `leakyRelu(accum, alpha)` \|
	93	+ \| 激活函数 \| Prelu \| `Prelu(accum, weight)` \|
	94	+ \| 激活函数 \| sigmoid \| `sigmoid(accum)` \|
	95	+ \| 激活函数 \| silu \| `silu(accum)` \|
	96	+ \| 比较/选择 \| max / min \| `max(a, b)` / `min(a, b)` \|
	97	+ \| 类型转换 \| cast \| `cast(accum, "float16", "float")` \|
	98	+ \| 常量 \| constant \| `constant(1.0, "float")` \|
	99	+ - 组合计算：支持多个计算节点拼接；
	100	+ - 广播计算：支持行广播计算。
	101	+
	102	+## 3. 安装
	103	+
	104	+### 3.1 从源码安装
	105	+
	106	+1. 构建分发包：
	107	+ ```bash
	108	+ pip install build
	109	+ python -m build
	110	+ ```
	111	+ 这会在 `dist/` 目录下生成 `.whl` 和 `.tar.gz` 文件。
	112	+
	113	+2. 安装分发包：
	114	+ ```bash
	115	+ pip install dist/catlass_cppgen-*.whl
	116	+ ```
	117	+ 或者：
	118	+ ```bash
	119	+ pip install dist/catlass_cppgen-*.tar.gz
	120	+ ```
	121	+
	122	+3. 直接安装（开发模式）：
	123	+ ```bash
	124	+ pip install -e .
	125	+ ```
	126	+
	127	+### 3.2 从本地目录安装
	128	+
	129	+如果您想直接从项目目录安装：
	130	+```bash
	131	+pip install .
	132	+```
	133	+
	134	+## 4. 使用示例
	135	+
	136	+当前 `catlass_cppgen` 支持 matmul、grouped_matmul 以及 EVG 后处理特性的代码生成。以下是应用`cppgen`的环节示意：
	137	+```plain
	138	+Gemm / GroupGemm（算子规划）
	139	+ ↓ get_kernels()
	140	+Kernel 对象（调优与特性查询）
	141	+ ↓ tune() / to_evg()
	142	+配置完成的 Kernel
	143	+```
	144	+
	145	+详细参考使用示例和 API 文档请参考下述文档：
	146	+ - [Kernel API 基础文档](docs/kernel_api.md)
	147	+ - [OpTensor API 基础文档](docs/input_methods.md)
	148	+ - [EVG API 基础文档](docs/evg_api.md)
	149	+
	150	+### 4.1 基础 GEMM
	151	+
	152	+以下是一个基础的创建 matmul 算子cppgen对象的示例：
	153	+
	154	+```python
	155	+from catlass_cppgen.op.gemm import Gemm
	156	+from catlass_cppgen.common.op_tensor import OpTensor
	157	+from catlass_cppgen.common.data_type import DataType
	158	+from catlass_cppgen.catlass.layout.layout import RowMajor
	159	+from catlass_cppgen.catlass.gemm_coord import GemmShape
	160	+from catlass_cppgen.catlass.arch.arch import Arch
	161	+from catlass_cppgen.catlass.gemm.dispatch_policy import MmadPingpong
	162	+
	163	+# 1. 描述输入张量（无需绑定底层数据）
	164	+a = OpTensor.from_shape_stride((128, 256), (256, 1), DataType.FLOAT)
	165	+b = OpTensor.from_shape_stride((256, 384), (384, 1), DataType.FLOAT)
	166	+
	167	+# 2. 创建算子并获取 Kernel
	168	+gemm = Gemm(atlas_arch=Arch.Ascend950, element=DataType.FLOAT, layout=RowMajor, A=a, B=b)
	169	+kernels = gemm.get_kernels()
	170	+
	171	+# 3. [可选] 可以显式指定使用的Kernel组件（以`BasicMatmulKernel`为例）
	172	+from catlass_cppgen.kernel.gemm import BasicMatmulKernel
	173	+kernel = find_kernel_by_type(kernels, BasicMatmulKernel)
	174	+# 非定向指定: kernel = kernels[0]
	175	+
	176	+# 4. [可选] 调优 Tile 形状与调度策略
	177	+kernel.tune(
	178	+ GemmShape(128, 256, 64),
	179	+ GemmShape(128, 256, 64),
	180	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950),
	181	+)
	182	+
	183	+# 5. [可选] matmul hello world: 核函数的代码生成
	184	+print(f"[Kernel] \n{kernel.gen_kernel_template()}")
	185	+```
	186	+
	187	+得到kernel对象后，可以调用`gen_kernel_template()`，`gen_params_device()`等方法，针对核函数和参数绑定做代码生成。
	188	+
	189	+### 4.2 Group GEMM
	190	+
	191	+以下是建立 matmul 算子cppgen对象的示例：
	192	+
	193	+```python
	194	+from catlass_cppgen.op.group_gemm import GroupGemm
	195	+from catlass_cppgen.catlass.layout.layout import VectorLayout
	196	+
	197	+# ...
	198	+
	199	+# 1. 创建groupList 张量
	200	+groupList = OpTensor(dtype=DataType.INT64, layout=VectorLayout(4), shape=(4,))
	201	+
	202	+# 2. 建立 group matmul 对象
	203	+group_gemm = GroupGemm(atlas_arch=Arch.Ascend950, A=a, B=b_3d, groupList=groupList)
	204	+
	205	+# 3. [可选] 取得kernel并做 Tiling 调优
	206	+kernels = group_gemm.get_kernels()
	207	+kernels[0].tune(GemmShape(256, 256, 256), GemmShape(256, 256, 64))
	208	+```
	209	+
	210	+
	211	+### 4.3 EVG 后处理
	212	+
	213	+```python
	214	+# ...
	215	+
	216	+# 1. 进行evg对象声明
	217	+# - fn_src: EVG对象函数头
	218	+# - example_inputs: 后处理过程中涉及的`name:tensor`键值对
	219	+evg_config = {
	220	+ "fn_src": "def epilogue(accum, bias):\n return relu(accum + bias)",
	221	+ "example_inputs": {
	222	+ "accum": OpTensor.from_shape_stride((128, 256), (256, 1), DataType.FLOAT),
	223	+ "bias": OpTensor.from_shape_stride((1, 256), (256, 1), DataType.FLOAT),
	224	+ "result": OpTensor.from_shape_stride((128, 256), (256, 1), DataType.FLOAT),
	225	+ },
	226	+}
	227	+
	228	+# 2. 创建 matmul 算子对象并获取 Kernel
	229	+gemm = Gemm(atlas_arch=Arch.Ascend950, evg_config=evg_config, A=a, B=b)
	230	+kernel = gemm.get_kernels()[0]
	231	+# `is_support_evg` 特性为 True
	232	+assert kernel.is_support_evg
	233	+
	234	+# [可选] 3. 进行 Tile 形状调优
	235	+kernel.tune(GemmShape(128, 256, 64), GemmShape(128, 256, 64))
	236	+```
	237	+
	238	+支持二元运算（add/sub/mul/div）、激活函数（relu/silu/sigmoid/leakyRelu/prelu）、类型转换（cast）、常量（constant）等，可串联组合使用。

Apython/catlass_cppgen/catlass_cppgen/__init__.py+50-0

@@ -0,0 +1,50 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from catlass_cppgen.op import Gemm, OperationBase
	11	+from catlass_cppgen.common import DataType, get_default_accumulator
	12	+from catlass_cppgen.catlass import (
	13	+ Arch,
	14	+ GemmShape,
	15	+ GemmCoord,
	16	+ Shape,
	17	+ MmadBase,
	18	+ MmadAtlasA2,
	19	+ MmadAtlasA2Async,
	20	+ MmadAtlasA2Pingpong,
	21	+ MmadAtlasA2PingpongSliceKWithPrologue,
	22	+ MmadAtlasA2PingPongWithPrologue,
	23	+ MmadAtlasA2Preload,
	24	+ MmadAtlasA2PreloadAsync,
	25	+ MmadAtlasA2PreloadAsyncWithCallback,
	26	+ GemmAtlasA2,
	27	+ GemvAtlasA2,
	28	+ MmadAtlasA2PingpongBias,
	29	+ MmadAtlasA2FullLoadA,
	30	+ MmadAtlasA2W8A16,
	31	+ MmadAtlasA2DynamicCommon,
	32	+ MmadAtlasA2Small,
	33	+ MmadPingpong,
	34	+ MmadPreloadAsyncWithCallback,
	35	+ MmadMultiBatch,
	36	+)
	37	+from catlass_cppgen.catlass.layout import (
	38	+ Layout,
	39	+ RowMajor,
	40	+ ColumnMajor,
	41	+ PaddingRowMajor,
	42	+ PaddingColumnMajor,
	43	+ VectorLayout,
	44	+ nZ,
	45	+ zN,
	46	+ zZ,
	47	+ nN,
	48	+)
	49	+
	50	+from catlass_cppgen._version import __version__

Apython/catlass_cppgen/catlass_cppgen/_version.py+28-0

@@ -0,0 +1,28 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+"""版本号管理模块，自动生成带时间戳的版本号"""
	11	+from datetime import datetime
	12	+
	13	+# 基础版本号
	14	+BASE_VERSION = "0.1.0"
	15	+
	16	+def get_version():
	17	+ """
	18	+ 生成带时间戳的版本号
	19	+
	20	+ 格式: {BASE_VERSION}+{YYYYMMDDHHMMSS}
	21	+ 例如: 0.1.0+20240101120000
	22	+ """
	23	+ timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
	24	+ return f"{BASE_VERSION}+{timestamp}"
	25	+
	26	+# 在构建时生成版本号
	27	+__version__ = get_version()
	28	+

Apython/catlass_cppgen/catlass_cppgen/catlass/__init__.py+57-0

@@ -0,0 +1,57 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from catlass_cppgen.catlass.arch.arch import Arch
	11	+from catlass_cppgen.catlass.gemm_coord import GemmShape, GemmCoord, Shape
	12	+from catlass_cppgen.catlass.gemm.dispatch_policy import (
	13	+ MmadBase,
	14	+ MmadAtlasA2,
	15	+ MmadAtlasA2Async,
	16	+ MmadAtlasA2Pingpong,
	17	+ MmadAtlasA2PingpongSliceKWithPrologue,
	18	+ MmadAtlasA2PingPongWithPrologue,
	19	+ MmadAtlasA2Preload,
	20	+ MmadAtlasA2PreloadAsync,
	21	+ MmadAtlasA2PreloadAsyncWithCallback,
	22	+ GemmAtlasA2,
	23	+ GemvAtlasA2,
	24	+ MmadAtlasA2PingpongBias,
	25	+ MmadAtlasA2FullLoadA,
	26	+ MmadAtlasA2W8A16,
	27	+ MmadAtlasA2DynamicCommon,
	28	+ MmadAtlasA2Small,
	29	+ MmadPingpong,
	30	+ MmadPreloadAsyncWithCallback,
	31	+ MmadMultiBatch,
	32	+)
	33	+
	34	+__all__ = [
	35	+ "Arch",
	36	+ "GemmShape",
	37	+ "GemmCoord",
	38	+ "Shape",
	39	+ "MmadAtlasA2",
	40	+ "MmadAtlasA2Async",
	41	+ "MmadAtlasA2Pingpong",
	42	+ "MmadAtlasA2PingpongSliceKWithPrologue",
	43	+ "MmadAtlasA2PingPongWithPrologue",
	44	+ "MmadAtlasA2Preload",
	45	+ "MmadAtlasA2PreloadAsync",
	46	+ "MmadAtlasA2PreloadAsyncWithCallback",
	47	+ "GemmAtlasA2",
	48	+ "GemvAtlasA2",
	49	+ "MmadAtlasA2PingpongBias",
	50	+ "MmadAtlasA2FullLoadA",
	51	+ "MmadAtlasA2W8A16",
	52	+ "MmadAtlasA2DynamicCommon",
	53	+ "MmadAtlasA2Small",
	54	+ "MmadPingpong",
	55	+ "MmadPreloadAsyncWithCallback",
	56	+ "MmadMultiBatch",
	57	+]

Apython/catlass_cppgen/catlass_cppgen/catlass/arch/arch.py+16-0

@@ -0,0 +1,16 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from enum import Enum
	11	+
	12	+
	13	+class Arch(Enum):
	14	+ AtlasA2 = "Arch::AtlasA2"
	15	+ AtlasA3 = "Arch::AtlasA2" # Commonly, Atlas A2/Atlas A3 use 'Arch::AtlasA2'
	16	+ Ascend950 = "Arch::Ascend950"

Apython/catlass_cppgen/catlass_cppgen/catlass/evg/__init__.py+33-0

@@ -0,0 +1,33 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from catlass_cppgen.catlass.evg.evg_definition import EVGArg, EVGDef
	CheaterAbec7月7日确认一下EVG代码是否支持tanh，如果支持是否同步了主仓对于EVG_tanh的Nan修复 I init__zhb__ 7月8日评论： 7月8日评论：
	11	+from catlass_cppgen.catlass.evg.node import (
	12	+ CastNode,
	13	+ ComputeNode,
	14	+ ConstantNode,
	15	+ LoadNode,
	16	+ NodeBase,
	17	+ NodeMetadata,
	18	+ StoreNode,
	19	+ TopoVisitorNode,
	20	+)
	21	+
	22	+__all__ = [
	23	+ "EVGArg",
	24	+ "EVGDef",
	25	+ "CastNode",
	26	+ "ComputeNode",
	27	+ "ConstantNode",
	28	+ "LoadNode",
	29	+ "NodeBase",
	30	+ "NodeMetadata",
	31	+ "StoreNode",
	32	+ "TopoVisitorNode",
	33	+]

Apython/catlass_cppgen/catlass_cppgen/catlass/evg/evg_definition.py+712-0

@@ -0,0 +1,712 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from __future__ import annotations
	11	+
	12	+from copy import deepcopy
	13	+from collections import Counter
	14	+from typing import TYPE_CHECKING, Any, Callable, List, Set, Union
	15	+
	16	+import networkx as nx
	17	+from sympy import Expr, Symbol
	18	+
	19	+from catlass_cppgen.common.data_type import DataType
	20	+from ..library import BroadcastType, EpilogueOpVectorToScalar
	21	+from .node import (CastNode, ComputeNode, ConstantNode, LoadNode, NodeBase,
	22	+ StoreNode, TopoVisitorNode)
	23	+from .node_impl import NoOpImpl
	24	+
	25	+if TYPE_CHECKING:
	26	+ from ..evg_extension import EpilogueVisitorGraph
	27	+
	28	+
	29	+# ============ Graph Passes ============
	30	+def EliminateDupAndDeadNodes(dag: EpilogueVisitorGraph):
	31	+ """
	32	+ 1) Eliminate duplicated nodes:
	33	+ - Two nodes are considered duplicates if:
	34	+ * they are the same node kind (ComputeNode, CastNode, ...)
	35	+ * they have the same operator identity (for ComputeNode: `fn`;
	36	+ for CastNode: from/to element types, etc.)
	37	+ * they have identical ordered inputs (same input nodes in the same edge positions)
	38	+ - On duplicate detection, we keep the first seen node as canonical and replace
	39	+ all outgoing edges from the duplicate node to point from the canonical node.
	40	+ Then we remove the duplicate node.
	41	+
	42	+ 2) Remove dead nodes:
	43	+ - Compute all nodes that are ancestors (via predecessors) of the output store node(s).
	44	+ - Any node not in that ancestor set is dead and will be removed.
	45	+
	46	+ This modifies the graph in-place via the existing `add_node`, `remove_node`,
	47	+ `add_edge`, `remove_edge`, `get_sorted_inputs`, `get_outputs`, and similar helpers.
	48	+ """
	49	+ # Helper to build a canonical signature for a node.
	50	+ def node_signature(node):
	51	+ # Gather ordered input ids (by edge pos)
	52	+ try:
	53	+ inputs = dag.get_sorted_inputs(node)
	54	+ except Exception:
	55	+ inputs = []
	56	+ input_ids = tuple(inp.name for inp in inputs)
	57	+
	58	+ if isinstance(node, ComputeNode):
	59	+ # include the op identity (fn) and possibly element type
	60	+ op_id = getattr(node, "fn", None)
	61	+ return ("ComputeNode", op_id, input_ids)
	62	+ if isinstance(node, CastNode):
	63	+ # include cast types
	64	+ from_e = getattr(node, "from_element", None)
	65	+ to_e = getattr(node, "to_element", None)
	66	+ return ("CastNode", from_e, to_e, input_ids)
	67	+ if isinstance(node, LoadNode):
	68	+ # Loads are unique by name (and marked inputs) — normally we do not dedupe loads.
	69	+ return ("LoadNode", node.name)
	70	+ if isinstance(node, StoreNode):
	71	+ # Stores are unique by name
	72	+ return ("StoreNode", node.name)
	73	+ # Fallback: use op name and inputs
	74	+ return (type(node).__name__, node.metadata.op, input_ids)
	75	+
	76	+ # 1) Remove duplicated nodes
	77	+ # Work on a snapshot of nodes in topological order to avoid concurrent modification issues.
	78	+ topo_nodes = list(dag.topological_nodes())
	79	+
	80	+ for node in topo_nodes:
	81	+ # Node may have been removed during previous iterations
	82	+ if not dag.has_node(node):
	83	+ continue
	84	+
	85	+ # Remove the StoreNode after the ConstantNode
	86	+ if isinstance(node, StoreNode):
	87	+ input_node = dag.get_inputs(node)[0]
	88	+ if isinstance(input_node, ConstantNode):
	89	+ output_node = dag.get_outputs(node)[0]
	90	+ dag.add_edge(input_node, output_node)
	91	+ dag.remove_edge(node, output_node)
	92	+ dag.remove_node(node)
	93	+ continue
	94	+
	95	+ signature_map: dict = {}
	96	+
	97	+ for node in dag.topological_nodes():
	98	+ node_info = tuple(dag.get_inputs(node))
	99	+ signature_map[node] = node_info
	100	+
	101	+ def simplify_graph(signature_map):
	102	+ modified = False
	103	+ for node in signature_map:
	104	+ if isinstance(node, LoadNode) or not signature_map[node]:
	105	+ continue
	106	+ for node_compare in signature_map:
	107	+ if node == node_compare:
	108	+ continue
	109	+ if Counter(signature_map[node]) == Counter(signature_map[node_compare]):
	110	+ for input_node in signature_map[node_compare]:
	111	+ dag.remove_edge(input_node, node_compare)
	112	+ for output_node in dag.get_outputs(node_compare):
	113	+ dag.remove_edge(node_compare, output_node)
	114	+ dag.add_edge(node, output_node)
	115	+ dag.remove_node(node_compare)
	116	+ modified = True
	117	+ signature_map = {}
	118	+ for node_tmp in list(dag.topological_nodes()):
	119	+ node_info = tuple(dag.get_inputs(node_tmp))
	120	+ signature_map[node_tmp] = node_info
	121	+ break
	122	+ if modified:
	123	+ break
	124	+ if modified:
	125	+ break
	126	+ if modified:
	127	+ simplify_graph(signature_map)
	128	+
	129	+ simplify_graph(signature_map)
	130	+
	131	+ # 2) Remove dead nodes
	132	+ # Find output store nodes.
	133	+ output_nodes = [
	134	+ n
	135	+ for n in dag.topological_nodes()
	136	+ if isinstance(n, StoreNode) and n.metadata.is_output
	137	+ ]
	138	+ if not output_nodes:
	139	+ raise RuntimeError("The parsed DAG has no output store node.")
	140	+
	141	+ reachable = set()
	142	+ # We'll perform reverse DFS from each output to collect all ancestors
	143	+ graph_rev = dag._graph.reverse(copy=False)
	144	+ for out in output_nodes:
	145	+ for n in nx.dfs_preorder_nodes(graph_rev, source=out):
	146	+ reachable.add(n)
	147	+
	148	+ # Any node not in reachable is dead -> remove
	149	+ all_nodes = list(dag.topological_nodes())
	150	+ dead_nodes = [n for n in all_nodes if n not in reachable]
	151	+ # Remove dead nodes except maybe LoadNodes that you want to keep (we remove them as dead too)
	152	+ for dn in dead_nodes:
	153	+ # skip outputs (should not be in dead_nodes) and skip canonical check
	154	+ if isinstance(dn, StoreNode) and getattr(dn.metadata, "is_output", False):
	155	+ continue
	156	+ if not dag.has_node(dn):
	157	+ continue
	158	+ dag.remove_node(dn)
	159	+
	160	+
	161	+def ScalarOpsIdentification(dag: EpilogueVisitorGraph):
	162	+ """
	163	+ Identity Scalar Mulitplication and Add
	164	+ """
	165	+ all_nodes = dag.topological_nodes()
	166	+ for node in all_nodes:
	167	+ if isinstance(node, ComputeNode) and node.fn in EpilogueOpVectorToScalar:
	168	+ input_nodes = dag.get_inputs(node)
	169	+ for input_node in input_nodes:
	170	+ if isinstance(input_node, ConstantNode):
	171	+ node.fn = EpilogueOpVectorToScalar[node.fn]
	172	+ node._scalar_values[f"{node.name}_scalar_{len(node._scalar_values)}"] = (input_node.value, input_node.metadata.element)
	173	+ dag.remove_edge(input_node, node)
	174	+ dag.remove_node(input_node)
	175	+
	176	+
	177	+def InferShape(dag: EpilogueVisitorGraph):
	178	+ """
	179	+ Infer each node's output shape from its input nodes.
	180	+ - Uses lexicographical topological order so inputs are inferred before users.
	181	+ - Broadcasts shapes using left-padding with 1s, then matching dimensions from the end.
	182	+ """
	183	+ def pad_left(shape_list, target_len):
	184	+ # returns a new list padded on the left with 1s to target_len
	185	+ if shape_list is None:
	186	+ return [1] * target_len
	187	+ if len(shape_list) >= target_len:
	188	+ return list(shape_list)
	189	+ return [1] * (target_len - len(shape_list)) + list(shape_list)
	190	+
	191	+ # Walk nodes in topological order so inputs are available
	192	+ for node in dag.topological_nodes():
	193	+ # Get ordered inputs (by edge pos)
	194	+ inputs = dag.get_sorted_inputs(node)
	195	+
	196	+ # No inputs: for many nodes this is a scalar or pre-filled metadata; leave as-is
	197	+ if not inputs:
	198	+ continue
	199	+
	200	+ # Collect source shapes
	201	+ src_shapes = []
	202	+ for src in inputs:
	203	+ meta = src.metadata
	204	+ src_shapes.append(meta.shape)
	205	+
	206	+ # Infer broadcasted shape
	207	+ shape = None
	208	+ for src_shape in src_shapes:
	209	+ if shape is None:
	210	+ # start with first shape
	211	+ shape = list(src_shape)
	212	+ continue
	213	+
	214	+ # pad shapes to same length on the left with 1s
	215	+ max_len = max(len(shape), len(src_shape))
	216	+ a = pad_left(shape, max_len)
	217	+ b = pad_left(list(src_shape), max_len)
	218	+
	219	+ # build result from right-to-left
	220	+ result_rev = []
	221	+ for dim_a, dim_b in zip(reversed(a), reversed(b)):
	222	+ if dim_a == 1:
	223	+ result_rev.append(dim_b)
	224	+ elif dim_b == 1:
	225	+ result_rev.append(dim_a)
	226	+ elif dim_a == dim_b:
	227	+ result_rev.append(dim_a)
	228	+ else:
	229	+ # construct helpful error message listing input shapes
	230	+ shapes_msg = ", ".join(f"{inp.name}{tuple(inp.metadata.shape)}" for inp in inputs)
	231	+ raise RuntimeError(f"Dimension mismatch between {shapes_msg}.")
	232	+ # reverse result_rev to get normal order
	233	+ shape = list(reversed(result_rev))
	234	+
	235	+ final_shape = tuple(shape) if shape else ()
	236	+
	237	+ node.metadata.shape = final_shape
	238	+
	239	+
	240	+def BroadcastPropagation(dag: EpilogueVisitorGraph):
	241	+ """
	242	+ Decide if any node needs broadcast and whether row- or col-broadcast is required.
	243	+ """
	244	+
	245	+ def check_broadcast(src_shape, dst_shape):
	246	+ """
	247	+ Check if src_shape can broadcast to dst_shape according to Numpy's broadcast rule.
	248	+
	249	+ Returns (can_broadcast: bool, row_broadcast: bool, col_broadcast: bool)
	250	+
	251	+ - We align shapes on the right and compare dimensions from right to left.
	252	+ - A dimension is compatible if either equal or one of them is 1.
	253	+ - If src dim == 1 and dst dim != 1, src is being broadcast along that dst axis.
	254	+ - If that dst axis is the last axis (rightmost), treat it as a column broadcast.
	255	+ - Otherwise treat it as a row broadcast.
	256	+ - If dst dim == 1 and src dim != 1 (including when dst lacks the axis), src cannot be broadcast
	257	+ to dst (return False).
	258	+ """
	259	+ a = list(src_shape)
	260	+ b = list(dst_shape)
	261	+
	262	+ i = len(a) - 1
	263	+ j = len(b) - 1
	264	+ row_broadcast = False
	265	+ col_broadcast = False
	266	+
	267	+ # Compare from right to left (align shapes to the right)
	268	+ while i >= 0 or j >= 0:
	269	+ dim_a = a[i] if i >= 0 else 1
	270	+ dim_b = b[j] if j >= 0 else 1
	271	+
	272	+ if dim_a == dim_b:
	273	+ # compatible, nothing to record
	274	+ pass
	275	+ elif dim_a == 1 and dim_b != 1:
	276	+ # src is broadcast along this dst axis
	277	+ # if this is the last (rightmost) dst axis -> column broadcast,
	278	+ # otherwise -> row broadcast
	279	+ if j == len(b) - 1:
	280	+ col_broadcast = True
	281	+ else:
	282	+ row_broadcast = True
	283	+ else:
	284	+ return False, row_broadcast, col_broadcast
	285	+
	286	+ i -= 1
	287	+ j -= 1
	288	+
	289	+ return True, row_broadcast, col_broadcast
	290	+
	291	+ for node in dag.topological_nodes():
	292	+ if isinstance(node, StoreNode):
	293	+ continue
	294	+ outputs = dag.get_outputs(node)
	295	+ if not outputs:
	296	+ continue
	297	+
	298	+ out_shapes = [out.metadata.shape for out in outputs]
	299	+ # Note: for one input, it may have different shape of outputs,
	300	+ # e.g. one computation is (m, n) = (m, n) + (n,); one computation is (n,) = (n,) + constant
	301	+ # In this case, we should broadcast it to the largest shape. And set all outputs' shape to the largest shape,
	302	+ # which means, in the example, we need to set (n,) = (n,) + constant -> (m, n) = (m, n) + constant
	303	+ dst_shape = max(out_shapes, key=len)
	C C Chen_HaoWen7月10日 [正确性 / 广播语义错误] `dst_shape = max(out_shapes, key=len)` 仅按维度数量（`len`）选择广播目标形状，而非按广播语义选择最大形状。当多个输出形状具有相同维度数但尺寸不同时（例如 `(2, 3)` 和 `(4, 5)`），`max` 会返回列表中第一个出现的形状，而非真正的最大形状。建议改为按广播规则逐维取最大值来确定 `dst_shape`： `max_len = max(len(s) for s in out_shapes) padded = [tuple([1] * (max_len - len(s)) + list(s)) for s in out_shapes] dst_shape = tuple(max(dims) for dims in zip(*padded))` 这样可以正确处理维度数相同但尺寸不同的场景。 I init__zhb__ 7月10日评论： 7月10日评论：
	304	+ for out_node in outputs:
	305	+ out_node.metadata.shape = dst_shape
	306	+
	307	+ # May modified, if reduce ops is implemented in future.
	308	+ src_shape = node.metadata.shape
	309	+ if len(dst_shape) < len(src_shape):
	310	+ dst_shape = src_shape
	311	+ for out in outputs:
	312	+ out.metadata.shape = src_shape
	313	+
	314	+ ok, row_broadcast, col_broadcast = check_broadcast(src_shape, dst_shape)
	315	+ if not ok:
	316	+ raise RuntimeError(
	317	+ f"Node '{node.name}' shape {src_shape} cannot be broadcast to output shape {dst_shape}"
	318	+ )
	319	+ if col_broadcast:
	320	+ raise RuntimeError(
	321	+ "EVG does not support ColumnBroadcast yet!"
	322	+ )
	323	+ if row_broadcast:
	324	+ node.metadata.broadcast = BroadcastType.RowBroadcast
	325	+ if len(node.metadata.shape) == 1:
	326	+ node.metadata.shape = (1, node.metadata.shape[0])
	327	+
	328	+
	329	+def DynamicShapeTransfer(dag: EpilogueVisitorGraph):
	330	+ """
	331	+ Subtitutes all shape infos into a experssion of m & n, which aims to support dynamic shape.
	332	+ """
	333	+ for node in dag.topological_nodes():
	334	+ # Within dynamic shape mode, the symbolic shape should be a sympy.Symbol or
	335	+ # a sympy expression containing Symbol(s) (e.g. 2*s20). Skip nodes whose shape
	336	+ # is fully static (no free symbols at all).
	337	+ if all(
	338	+ not (isinstance(shape_i, Expr) and shape_i.free_symbols)
	339	+ for shape_i in node.metadata.shape
	340	+ ):
	341	+ continue
	342	+ symbolic_shape = []
	343	+ # Replace longer keys first so that a compound expression like "2*s20" is
	344	+ # substituted as a whole before its substring "s20" gets matched, e.g.
	345	+ # dict {"s20": "m", "2s20": "n"} should turn "2s20" into "n", not "2*m".
	346	+ sorted_substitution = sorted(
	347	+ dag.symbol_shape_substitution_dict.items(), key=lambda kv: len(kv[0]), reverse=True
	348	+ )
	349	+ for shape_i in node.metadata.shape:
	350	+ shape_i = str(shape_i)
	351	+ for key, item in sorted_substitution:
	352	+ shape_i = shape_i.replace(key, item)
	353	+ symbolic_shape.append(shape_i)
	354	+ node.metadata.shape = tuple(symbolic_shape) if symbolic_shape else ()
	355	+
	356	+
	357	+def SetNodeImpl(dag: EpilogueVisitorGraph):
	358	+ """
	359	+ Map each node of the EVG to the underlying node impl.
	360	+ """
	361	+ for node in dag.topological_nodes():
	362	+ if isinstance(node, TopoVisitorNode):
	363	+ # TopoVisitorNode's impl is already set at its initialization
	364	+ continue
	365	+ node.get_impl()
	366	+
	367	+ # Eliminate node with NoOpImpl
	368	+ for node in dag.topological_nodes():
	369	+ if isinstance(node.impl, NoOpImpl):
	370	+ input_nodes = dag.get_inputs(node)
	371	+ if len(input_nodes) != 1:
	372	+ raise ValueError(f"Node {node.name} with NoOpImpl must have exactly one input node")
	373	+ in_node = input_nodes[0]
	374	+ for out_node in dag.get_outputs(node):
	375	+ pos = dag.get_edge_pos(node, out_node)
	376	+ dag.add_edge(in_node, out_node, pos)
	377	+ dag.remove_edge(node, out_node)
	378	+ dag.remove_node(node)
	379	+
	380	+
	381	+def DAG2Tree(dag: EpilogueVisitorGraph):
	382	+ """
	383	+ Transform a DAG to Tree by fusing subgraphs containing nodes with multiple outputs.
	384	+ """
	385	+ def _find_lca(node: NodeBase):
	386	+ output_nodes = list(dag.get_outputs(node))
	387	+ if not output_nodes:
	388	+ return None, None
	389	+
	390	+ reachable_nodes = []
	391	+ for s_node in output_nodes:
	392	+ reachable_nodes.append(set(dag.all_reachable_nodes(s_node)))
	393	+ common_nodes = set.intersection(*reachable_nodes)
	394	+ if not common_nodes:
	395	+ return None, None
	396	+
	397	+ topo_nodes = dag.topological_nodes()
	398	+ lca = min(common_nodes, key=lambda node: topo_nodes.index(node))
	399	+ nodes_to_fuse = set.union(*reachable_nodes).difference(common_nodes)
	400	+ nodes_to_fuse.add(lca)
	401	+ return lca, nodes_to_fuse
	402	+
	403	+ def _fuse_subgraph(nodes_to_fuse: Set[NodeBase], lca: NodeBase):
	404	+ """
	405	+ Fuse the nodes between node and lca into a single TopoVisitorNode.
	406	+ """
	407	+ from ..evg_extension import EpilogueVisitorGraph
	408	+
	409	+ # Get all the immediate inputs & outputs of the nodes_to_fuse
	410	+ all_input_nodes = set()
	411	+ all_output_nodes = set()
	412	+ for node in nodes_to_fuse:
	413	+ all_input_nodes.update(dag.get_inputs(node))
	414	+ all_output_nodes.update(dag.get_outputs(node))
	415	+ new_subgraph_nodes = set.union(nodes_to_fuse, all_input_nodes, all_output_nodes)
	416	+
	417	+ lca_output_count = len(dag.get_outputs(lca))
	418	+
	419	+ subgraph_ = dag._graph.subgraph(new_subgraph_nodes)
	420	+ subgraph = EpilogueVisitorGraph()
	421	+ for node in subgraph_.nodes:
	422	+ new_node = deepcopy(node)
	423	+ if node not in nodes_to_fuse:
	424	+ new_node.disabled = True
	425	+ subgraph.add_node(new_node)
	426	+ for edge in subgraph_.edges:
	427	+ subgraph.add_edge(
	428	+ edge[0].name, edge[1].name, dag.get_edge_pos(edge[0], edge[1])
	429	+ )
	430	+
	431	+ tv_node = TopoVisitorNode(
	432	+ name=f"tv_{lca.name}",
	433	+ metadata=lca.metadata,
	434	+ subgraph=subgraph,
	435	+ output_node=lca,
	436	+ lca_output_count=lca_output_count,
	437	+ )
	438	+ dag.add_node(tv_node)
	439	+
	440	+ # Add input edges
	441	+ for idx, node in enumerate(all_input_nodes):
	442	+ dag.add_edge(node, tv_node, pos=idx)
	443	+
	444	+ # Replace all uses of lca with TopoVisitorNode
	445	+ for node in dag.get_outputs(lca):
	446	+ pos = dag.get_edge_pos(lca, node)
	447	+ dag.add_edge(tv_node, node, pos)
	448	+ dag.remove_edge(lca, node)
	449	+ dag.remove_node(lca)
	450	+
	451	+ # Replace all fused nodes
	452	+ nodes_to_fuse.remove(lca)
	453	+ for node in nodes_to_fuse:
	454	+ dag.remove_node(node)
	455	+
	456	+
	457	+ def topo_fuse_graph(dag):
	458	+ multiple_output_nodes = [node for node in dag.topological_nodes() if dag.out_degree(node) > 1]
	459	+ for node in multiple_output_nodes:
	460	+ if not (dag.has_node(node) and dag.out_degree(node) > 1):
	461	+ continue
	462	+
	463	+ # Find LCA, nodes_to_fuse
	464	+ lca, nodes_to_fuse = _find_lca(node)
	465	+
	466	+ if not lca:
	467	+ raise NotImplementedError("No LCA found.")
	468	+
	469	+ _fuse_subgraph(nodes_to_fuse, lca)
	470	+ new_multiple_output_nodes = [node for node in dag.topological_nodes() if dag.out_degree(node) > 1]
	471	+ if tuple(new_multiple_output_nodes) != tuple(multiple_output_nodes):
	472	+ topo_fuse_graph(dag)
	473	+
	474	+ topo_fuse_graph(dag)
	475	+
	476	+
	477	+class EVGDef:
	478	+ def __init__(self, dag: EpilogueVisitorGraph):
	479	+ self.dag = dag
	480	+ EliminateDupAndDeadNodes(self.dag)
	481	+ ScalarOpsIdentification(self.dag)
	482	+ InferShape(self.dag)
	483	+ BroadcastPropagation(self.dag)
	484	+ DynamicShapeTransfer(self.dag)
	485	+ SetNodeImpl(self.dag)
	486	+ DAG2Tree(self.dag)
	487	+
	488	+ def get_visitor_name(self, node: NodeBase) -> str:
	489	+ if not isinstance(node, TopoVisitorNode) and self.dag.in_degree(node) > 0:
	490	+ return f"EVG{node.type_name}"
	491	+ return node.type_name
	492	+
	493	+ def definition(self):
	494	+ nodes = self.dag.topological_nodes()
	495	+ evg_str = ""
	496	+ # Define 1. individual node type decl
	497	+ # 2. epilogue tree node
	498	+ # 3. topovisitor node
	499	+ for node in nodes:
	500	+ if not node.disabled:
	501	+ evg_str += self.def_node(node)
	502	+ if isinstance(node, TopoVisitorNode):
	503	+ evg_str += self.def_subgraph_node(node)
	504	+ else:
	505	+ # Tree visitor node
	506	+ evg_str += self.def_tree_node(node)
	507	+
	508	+ callback_name = self.get_visitor_name(nodes[-1])
	509	+ return evg_str, callback_name
	510	+
	511	+ def def_node(self, node: NodeBase):
	512	+ if isinstance(node, TopoVisitorNode):
	513	+ node_str = ""
	514	+ for inner_node in node.subgraph.topological_nodes():
	515	+ if not inner_node.disabled:
	516	+ node_str += self.def_node(inner_node)
	517	+ return node_str
	518	+ return node.impl.type_decl
	519	+
	520	+ def def_tree_node(self, node: NodeBase):
	521	+ if self.dag.in_degree(node) == 0:
	522	+ return ""
	523	+
	524	+ sorted_input_nodes = self.dag.get_sorted_inputs(node)
	525	+ inputs_str = ",\n".join(
	526	+ f" {self.get_visitor_name(i_node)}" for i_node in sorted_input_nodes
	527	+ )
	528	+ tree_node_str = f"""
	529	+using EVG{node.type_name} = Catlass::Epilogue::Fusion::TreeVisitor<
	530	+ {node.type_name},
	531	+{inputs_str}
	532	+>;
	533	+"""
	534	+ return tree_node_str
	535	+
	536	+ def def_subgraph_node(self, node: NodeBase):
	537	+ subgraph = node.subgraph
	538	+ subgraph_nodes = subgraph.topological_nodes()
	539	+
	540	+ # define the edge tuple
	541	+ edges_str = "tla::tuple<\n"
	542	+ for snode in subgraph_nodes[:-node.lca_output_count]:
	543	+ sorted_input_nodes = subgraph.get_sorted_inputs(snode)
	544	+ sorted_input_ids = [str(subgraph_nodes.index(i_node)) for i_node in sorted_input_nodes]
	545	+ edge_str = " tla::seq<" + ", ".join(sorted_input_ids) + ">,\n"
	546	+ edges_str += edge_str
	547	+ parts = edges_str.rsplit(",", 1)
	548	+ edges_str = parts[0] + "\n"
	549	+ edges_str += " >"
	550	+
	551	+ # define the nodes list
	552	+ tv_nodes_str_list = []
	553	+ for snode in subgraph_nodes[:-node.lca_output_count]:
	554	+ if snode.disabled:
	555	+ tv_nodes_str_list.append(f" {self.get_visitor_name(snode)}")
	556	+ else:
	557	+ tv_nodes_str_list.append(f" {snode.type_name}")
	558	+
	559	+ tv_nodes_str = ",\n".join(tv_nodes_str_list)
	560	+
	561	+ subgraph_str = f"""
	562	+using {node.type_name} = Catlass::Epilogue::Fusion::TopologicalVisitor<
	563	+ {edges_str},
	564	+{tv_nodes_str}
	565	+>;
	566	+"""
	567	+ return subgraph_str
	568	+
	569	+
	570	+class EVGArg:
	571	+ def __init__(self, dag: EpilogueVisitorGraph):
	572	+ self.dag = dag
	573	+
	574	+ def generate_graph_args(self) -> str:
	575	+ """
	576	+ Return Arg Infos and ArgRename for the EVG
	577	+ """
	578	+ # find output nodes
	579	+ output_nodes = [
	580	+ node
	581	+ for node in self.dag.topological_nodes()
	582	+ if self.dag.out_degree(node) == 0
	583	+ ]
	584	+ output_nodes = [node for node in output_nodes if isinstance(node, StoreNode)]
	585	+
	586	+ if not output_nodes:
	587	+ raise ValueError("Cannot find output node in EVG")
	588	+
	589	+ if len(output_nodes) > 1:
	590	+ raise ValueError(
	591	+ "Find more than one node in EVG, currently, only support one output"
	592	+ )
	593	+
	594	+ output_node = output_nodes[0]
	595	+
	596	+ # generate arguments for the given EVG
	597	+ final_args = f"typename EVG{output_node.type_name}::Arguments evg_args"
	598	+ final_args += self.generate_node_args(output_node, self.dag)
	599	+ final_args = self._replace_last_comma_with_semicolon(final_args)
	600	+
	601	+ return (
	602	+ final_args,
	603	+ None,
	604	+ )
	605	+
	606	+ def generate_compute_length(self) -> str:
	607	+ """
	608	+ Return the compute length arg for the given DAG, which is used as tiling infos
	609	+ """
	610	+ nodes_element_dict = {}
	611	+ self.dag.get_storage_nodes(nodes_element_dict)
	612	+ nodes_length_str = "(" + " + ".join([f"({num} * sizeof({element.value}))" for element, num in nodes_element_dict.items()]) + ")"
	613	+ return f"""
	614	+constexpr uint32_t computeLength = 216 * 1024 / {nodes_length_str} / 2 / 32 * 32;
	615	+"""
	616	+
	617	+
	618	+ def generate_node_args(
	619	+ self, node: NodeBase, graph: EpilogueVisitorGraph, hierarchical_count=0
	620	+ ) -> str:
	621	+ """
	622	+ Genrate arg infos for the given node
	623	+
	624	+ Args:
	625	+ node: input node
	626	+ graph: Input EVG graph
	627	+
	628	+ Returns:
	629	+ Arg infos
	630	+ """
	631	+ if isinstance(node, TopoVisitorNode):
	632	+ topo_inputs = graph.get_sorted_inputs(node)
	633	+ tree_result_list = []
	634	+ for node_i in topo_inputs:
	635	+ tree_result_list.append(
	636	+ self.generate_node_args(node_i, graph, hierarchical_count + 1)
	637	+ )
	638	+ topo_result = self.generate_topo_visitor_args(
	639	+ node, topo_inputs, hierarchical_count + 1
	640	+ )
	641	+ for i, value in enumerate(tree_result_list):
	642	+ topo_result = topo_result.replace(
	643	+ f"{topo_inputs[i].name}", tree_result_list[i]
	644	+ )
	645	+ return topo_result
	646	+
	647	+ return self.generate_regular_node_args(node, graph, hierarchical_count)
	648	+
	649	+ def generate_regular_node_args(
	650	+ self, node: NodeBase, graph: EpilogueVisitorGraph, hierarchical_count
	651	+ ) -> str:
	652	+ """generate tree node args"""
	653	+ input_nodes = graph.get_sorted_inputs(node)
	654	+ result = ""
	655	+ if input_nodes:
	656	+ result += self._generate_indent(hierarchical_count) + "{\n"
	657	+
	658	+ for i, input_node in enumerate(input_nodes):
	659	+ input_arg = self.generate_node_args(
	660	+ input_node, graph, hierarchical_count + 1
	661	+ )
	662	+ result += input_arg
	663	+ result += (
	664	+ self._generate_indent(hierarchical_count + 1)
	665	+ + node.impl.args_decl
	666	+ + ",\n"
	667	+ )
	668	+ if input_nodes:
	669	+ result = self._remove_last_comma(result)
	670	+ result += self._generate_indent(hierarchical_count) + "},\n"
	671	+ return result
	672	+
	673	+ def generate_topo_visitor_args(
	674	+ self, node: TopoVisitorNode, topo_inputs, hierarchical_count
	675	+ ) -> str:
	676	+ """generating the topo graph's arguments"""
	677	+ subgraph = node.subgraph
	678	+ output_node = node.output_node
	679	+
	680	+ # find all reachable nodes in the topo subgraphs
	681	+ reachable_nodes = subgraph.topological_nodes()[:-node.lca_output_count]
	682	+ result = self._generate_indent(hierarchical_count - 1) + "{\n"
	683	+ for node in reachable_nodes:
	684	+ if node not in topo_inputs:
	685	+ result += (
	686	+ self._generate_indent(hierarchical_count)
	687	+ + node.impl.args_decl
	688	+ + ",\n"
	689	+ )
	690	+ else:
	691	+ result += f"{node.name}"
	692	+ result = self._remove_last_comma(result)
	693	+ result += self._generate_indent(hierarchical_count - 1) + "},\n"
	694	+
	695	+ return result
	696	+
	697	+ def _replace_last_comma_with_semicolon(self, s):
	698	+ parts = s.rsplit(",", 1)
	699	+ if len(parts) == 1:
	700	+ return s
	701	+
	702	+ return ";".join(parts)
	703	+
	704	+ def _remove_last_comma(self, s):
	705	+ parts = s.rsplit(",", 1)
	706	+ if len(parts) == 1:
	707	+ return s
	708	+
	709	+ return parts[0] + parts[1]
	710	+
	711	+ def _generate_indent(self, hierarchical_count: int) -> str:
	712	+ return " " * hierarchical_count * 4

Apython/catlass_cppgen/catlass_cppgen/catlass/evg/node.py+182-0

@@ -0,0 +1,182 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from __future__ import annotations
	11	+
	12	+from dataclasses import dataclass
	13	+from typing import TYPE_CHECKING, Tuple, Dict
	14	+
	15	+from catlass_cppgen.common.data_type import DataType
	16	+from ..library import (
	17	+ BroadcastType,
	18	+ BroadcastTag,
	19	+ CastType,
	20	+ EpilogueOp,
	21	+ EpilogueScalarOp,
	22	+ LayoutType,
	23	+)
	24	+from .node_impl import (
	25	+ AccLoadImpl,
	26	+ AuxLoadImpl,
	27	+ AuxStoreImpl,
	28	+ ComputeImpl,
	29	+ CastImpl,
	30	+ NoOpImpl,
	31	+ RowBroadcastImpl,
	32	+ ScalarComputeImpl,
	33	+ TopoVisitorImpl,
	34	+)
	35	+
	36	+if TYPE_CHECKING:
	37	+ from ..evg_extension import EpilogueVisitorGraph
	38	+ from ..library import EpilogueOp
	39	+
	40	+
	41	+@dataclass
	42	+class NodeMetadata:
	43	+ op: str = ""
	44	+ element: DataType = DataType.UNDEFINED
	45	+ layout: LayoutType = LayoutType.RowMajor
	46	+ broadcast: BroadcastType = BroadcastType.NoBroadcast
	47	+ shape: Tuple[int, ...] = ()
	48	+ stride: Tuple[int, ...] = ()
	49	+ is_output: bool = False
	50	+
	51	+
	52	+class NodeBase:
	53	+ def __init__(self, name: str, metadata: NodeMetadata):
	54	+ self.name = name # unique name for the node
	55	+ self.metadata = metadata
	56	+ self.disabled = False
	57	+ self.impl = None # the underlying impl of this node
	58	+
	59	+ @property
	60	+ def type_name(self) -> str:
	61	+ return self.impl.type_name
	62	+
	63	+ def signature_name(self) -> str:
	64	+ res = str(self)[:-1]
	65	+ extra_attrs = [" "]
	66	+ if self.metadata.element != DataType.UNDEFINED:
	67	+ extra_attrs.append(self.metadata.element.value)
	68	+ if self.metadata.shape:
	69	+ extra_attrs.append("x".join(map(str, self.metadata.shape)))
	70	+ extra_str = " \| ".join(extra_attrs) + ">"
	71	+ res += extra_str
	72	+ return res
	73	+
	74	+ def get_impl(self):
	75	+ raise NotImplementedError(
	76	+ f"Function `get_impl` is not overloaded in {self.__class__.__name__}"
	77	+ )
	78	+
	79	+ def __repr__(self) -> str:
	80	+ return f"<Node {self.name} \| {self.metadata.op}>"
	81	+
	82	+ def __hash__(self) -> int:
	83	+ return hash(self.name)
	84	+
	85	+ def __eq__(self, other: NodeBase) -> bool:
	86	+ return isinstance(other, NodeBase) and self.name == other.name
	87	+
	88	+ def __lt__(self, other: NodeBase) -> bool:
	89	+ return self.name < other.name
	90	+
	91	+
	92	+class ComputeNode(NodeBase):
	93	+ def __init__(self, name: str, metadata: NodeMetadata, fn: EpilogueOp):
	94	+ super().__init__(name, metadata)
	95	+ self.fn = fn
	96	+ self._scalar_values: Dict[str, Tuple[str, DataType]] = {}
	97	+
	98	+ def get_impl(self):
	99	+ if self.fn in EpilogueScalarOp:
	100	+ self.impl = ScalarComputeImpl(self, self._scalar_values)
	101	+ else:
	102	+ self.impl = ComputeImpl(self)
	103	+
	104	+
	105	+class CastNode(NodeBase):
	106	+ def __init__(
	107	+ self, name: str, metadata: NodeMetadata, from_element, to_element, round_type=CastType.NONE
	108	+ ):
	109	+ super().__init__(name, metadata)
	110	+ self.from_element = from_element
	111	+ self.to_element = to_element
	112	+ self.round_type = round_type
	113	+
	114	+ def get_impl(self):
	115	+ self.impl = CastImpl(self)
	116	+
	117	+
	118	+class ReduceNode(NodeBase):
	119	+ def __init__(self, name: str, metadata: NodeMetadata, reduce_fn):
	120	+ super().__init__(name, metadata)
	121	+ self.reduce_fn = reduce_fn
	122	+
	123	+
	124	+class ConstantNode(NodeBase):
	125	+ def __init__(self, name: str, metadata: NodeMetadata, value):
	126	+ super().__init__(name, metadata)
	127	+ self.value = value
	128	+
	129	+
	130	+class LoadNode(NodeBase):
	131	+ def __init__(self, name: str, metadata: NodeMetadata):
	132	+ super().__init__(name, metadata)
	133	+
	134	+ # Possible impls:
	135	+ # 1. AccLoadImpl
	136	+ # 2. AuxLoadImpl
	137	+ # 3. RowBroadcastImpl
	138	+ def get_impl(self):
	139	+ if self.metadata.broadcast != BroadcastType.NoBroadcast:
	140	+ if self.metadata.op.lower() != "auxload":
	141	+ raise ValueError("For broadcast mode, the evg op must be 'auxload'")
	142	+ if self.metadata.broadcast == BroadcastType.RowBroadcast:
	143	+ self.impl = RowBroadcastImpl(self)
	144	+ return
	145	+ raise RuntimeError(
	146	+ f"Node `{self.name}` does not support {BroadcastTag[self.metadata.broadcast]}"
	147	+ )
	148	+ if self.metadata.op.lower() == "accload":
	149	+ self.impl = AccLoadImpl(self)
	150	+ elif self.metadata.op.lower() == "auxload":
	151	+ self.impl = AuxLoadImpl(self)
	152	+
	153	+
	154	+class StoreNode(NodeBase):
	155	+ def __init__(self, name: str, metadata: NodeMetadata):
	156	+ super().__init__(name, metadata)
	157	+
	158	+ # Possible impls:
	159	+ # 1. AuxStoreImpl
	160	+ # 2. NoOpImpl
	161	+ def get_impl(self):
	162	+ if self.metadata.is_output:
	163	+ self.impl = AuxStoreImpl(self)
	164	+ else:
	165	+ self.impl = NoOpImpl(self)
	166	+
	167	+
	168	+class TopoVisitorNode(NodeBase):
	169	+ def __init__(
	170	+ self,
	171	+ name: str,
	172	+ metadata: NodeMetadata,
	173	+ subgraph: EpilogueVisitorGraph,
	174	+ output_node: NodeBase,
	175	+ lca_output_count: int,
	176	+ ):
	177	+ super().__init__(name, metadata)
	178	+ self.metadata.op = "topo_visitor"
	179	+ self.subgraph = subgraph
	180	+ self.output_node = output_node
	181	+ self.impl = TopoVisitorImpl(self)
	182	+ self.lca_output_count = lca_output_count

Apython/catlass_cppgen/catlass_cppgen/catlass/evg/node_impl.py+263-0

@@ -0,0 +1,263 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from re import sub
	11	+
	12	+from catlass_cppgen.common.data_type import DataType
	13	+from ..library import LayoutTag, CastTypeTag, EpilogueOpTag
	14	+
	15	+
	16	+class ImplBase:
	17	+ def __init__(self, node):
	18	+ self.node = node
	19	+ self.name = node.name
	20	+ self.element = node.metadata.element
	21	+ self.shape = node.metadata.shape
	22	+ self.layout = node.metadata.layout
	23	+ self._type_decl = None
	24	+
	25	+ @property
	26	+ def type_name(self):
	27	+ return sub(r"(_\|-)+", " ", self.name).title().replace(" ", "")
	28	+
	29	+ @property
	30	+ def args_decl(self):
	31	+ return "{}"
	32	+
	33	+ def make_layout(self):
	34	+ shape_str = ", ".join(str(x) for x in self.shape)
	35	+ t_name = self.type_name
	36	+ layout_name = f"layout{t_name}"
	37	+ layout_str = f"""
	38	+using LayoutTag{t_name} = {LayoutTag[self.layout]};
	39	+LayoutTag{t_name} tag{t_name}{{{shape_str}}};
	40	+auto {layout_name} = tla::MakeLayoutFromTag(tag{t_name});
	41	+"""
	42	+ return layout_name, layout_str
	43	+
	44	+
	45	+class ComputeImplBase(ImplBase):
	46	+ """
	47	+ Base class for compute node implementations
	48	+ """
	49	+
	50	+ def __init__(self, node):
	51	+ super().__init__(node)
	52	+ self.fn = self.node.fn
	53	+ self.compute_element = node.metadata.element
	54	+
	55	+
	56	+class CastImplBase(ImplBase):
	57	+ """
	58	+ Base class for cast node implementations
	59	+ """
	60	+
	61	+ def __init__(self, node):
	62	+ super().__init__(node)
	63	+ self.to_element = node.to_element
	64	+ self.from_element = node.from_element
	65	+ self.round_type = node.round_type
	66	+
	67	+
	68	+class ReductionImplBase(ImplBase):
	69	+ """
	70	+ Base class for reduction node implementations
	71	+ """
	72	+
	73	+ def __init__(self, node):
	74	+ super().__init__(node)
	75	+ self.reduce_fn = self.node.reduce_fn
	76	+
	77	+
	78	+class NoOpImpl(ImplBase):
	79	+ """
	80	+ The NoOpImpl does nothing but forward its inputs to users.
	81	+ """
	82	+ def __init__(self, node):
	83	+ super().__init__(node)
	84	+
	85	+
	86	+class AccLoadImpl(ImplBase):
	87	+ @property
	88	+ def type_decl(self):
	89	+ """
	90	+ Return the string defining the type
	91	+ """
	92	+ if self._type_decl is not None:
	93	+ return self._type_decl
	94	+
	95	+# self._type_decl = f"""
	96	+# using {self.type_name} = Catlass::Epilogue::Fusion::VisitorAccLoad<{self.element.value}, EpilogueDispatchPolicy::USE_UB_WORKSPACE>;
	97	+# """
	98	+ self._type_decl = f"""
	99	+using {self.type_name} = Catlass::Epilogue::Fusion::VisitorAccLoad<{self.element.value}>;
	100	+"""
	101	+ return self._type_decl
	102	+
	103	+
	104	+class AuxLoadImpl(ImplBase):
	105	+ def __init__(self, node):
	106	+ super().__init__(node)
	107	+ self.layout_name = None
	108	+
	109	+ @property
	110	+ def type_decl(self):
	111	+ """
	112	+ Return the string defining the type
	113	+ """
	114	+ if self._type_decl is not None:
	115	+ return self._type_decl
	116	+
	117	+ self.layout_name, layout_str = self.make_layout()
	118	+ self._type_decl = layout_str
	119	+ self._type_decl += f"""
	120	+using {self.type_name} = Catlass::Epilogue::Fusion::VisitorAuxLoad<
	121	+ {self.element.value}, decltype({self.layout_name})
	122	+>;
	123	+"""
	124	+ return self._type_decl
	125	+
	126	+ @property
	127	+ def args_decl(self):
	128	+ if not self.layout_name:
	129	+ self.layout_name, _ = self.make_layout()
	130	+ return f"{{{self.name}_ptr, {self.layout_name}}}"
	131	+
	132	+
	133	+class AuxStoreImpl(ImplBase):
	134	+ def __init__(self, node):
	135	+ super().__init__(node)
	136	+ self.layout_name = None
	137	+
	138	+ @property
	139	+ def type_decl(self):
	140	+ """
	141	+ Return the string defining the type
	142	+ """
	143	+ if self._type_decl is not None:
	144	+ return self._type_decl
	145	+
	146	+ self.layout_name, layout_str = self.make_layout()
	147	+ self._type_decl = layout_str
	148	+ self._type_decl += f"""
	149	+using {self.type_name} = Catlass::Epilogue::Fusion::VisitorAuxStore<
	150	+ {self.element.value}, decltype({self.layout_name})
	151	+>;
	152	+"""
	153	+ return self._type_decl
	154	+
	155	+ @property
	156	+ def args_decl(self):
	157	+ if not self.layout_name:
	158	+ self.layout_name, _ = self.make_layout()
	159	+ return f"{{deviceC, {self.layout_name}}}"
	160	+
	161	+
	162	+class CastImpl(CastImplBase):
	163	+ @property
	164	+ def type_decl(self):
	165	+ """
	166	+ Return the string defining the type
	167	+ """
	168	+ if self._type_decl is not None:
	169	+ return self._type_decl
	170	+
	171	+ self._type_decl = f"""
	172	+using {self.type_name} = Catlass::Epilogue::Fusion::VisitorCast<
	173	+ {self.to_element.value}, {self.from_element.value},
	174	+ {CastTypeTag[self.round_type]}
	175	+>;
	176	+"""
	177	+ return self._type_decl
	178	+
	179	+
	180	+class ComputeImpl(ComputeImplBase):
	181	+ @property
	182	+ def type_decl(self):
	183	+ """
	184	+ Return the string defining the type
	185	+ """
	186	+ if self._type_decl is not None:
	187	+ return self._type_decl
	188	+
	189	+ self._type_decl = f"""
	190	+using {self.type_name} = Catlass::Epilogue::Fusion::VisitorCompute<
	191	+ {EpilogueOpTag[self.fn]}, {self.compute_element.value}
	192	+>;
	193	+"""
	194	+ return self._type_decl
	195	+
	196	+
	197	+class ScalarComputeImpl(ComputeImplBase):
	198	+ def __init__(self, node, values):
	199	+ super().__init__(node)
	200	+ self.scalar_values = values
	201	+
	202	+ @property
	203	+ def type_decl(self):
	204	+ """
	205	+ Return the string defining the type
	206	+ """
	207	+ if self._type_decl is not None:
	208	+ return self._type_decl
	209	+
	210	+ self._type_decl = ""
	211	+ for name, item in self.scalar_values.items():
	212	+ self._type_decl += f"""
	213	+{item[1].value} {name} = {item[0]};
	214	+"""
	215	+
	216	+ self._type_decl += f"""
	217	+using {self.type_name} = Catlass::Epilogue::Fusion::VisitorCompute<
	218	+ {EpilogueOpTag[self.fn]}, {self.compute_element.value},"""
	219	+ self._type_decl += ", ".join([item[1].value for _, item in self.scalar_values.items()])
	220	+ self._type_decl += """
	221	+>;
	222	+"""
	223	+ return self._type_decl
	224	+
	225	+ @property
	226	+ def args_decl(self):
	227	+ return f"{{{{{', '.join([name for name in self.scalar_values])}}}}}"
	228	+
	229	+
	230	+class RowBroadcastImpl(ImplBase):
	231	+ def __init__(self, node):
	232	+ super().__init__(node)
	233	+ self.layout_name = None
	234	+
	235	+ @property
	236	+ def type_decl(self):
	237	+ """
	238	+ Return the string defining the type
	239	+ """
	240	+ if self._type_decl is not None:
	241	+ return self._type_decl
	242	+
	243	+ self.layout_name, layout_str = self.make_layout()
	244	+ self._type_decl = layout_str
	245	+ self._type_decl += f"""
	246	+using {self.type_name} = Catlass::Epilogue::Fusion::VisitorRowBroadcast<
	247	+ {self.element.value}, decltype({self.layout_name})
	248	+>;
	249	+"""
	250	+ return self._type_decl
	251	+
	252	+ @property
	253	+ def args_decl(self):
	254	+ if not self.layout_name:
	255	+ self.layout_name, _ = self.make_layout()
	256	+ return f"{{{self.name}_ptr, {self.layout_name}}}"
	257	+
	258	+
	259	+class TopoVisitorImpl(ImplBase):
	260	+ def __init__(self, node):
	261	+ super().__init__(node.output_node)
	262	+ self.name = node.name
	263	+ self.element = node.output_node.metadata.element

Apython/catlass_cppgen/catlass_cppgen/catlass/evg_extension.py+448-0

@@ -0,0 +1,448 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+"""
	11	+Base class for Python EVG fronted
	12	+"""
	13	+
	14	+from __future__ import annotations
	15	+
	16	+import ast
	17	+import inspect
	18	+import itertools
	19	+import textwrap
	20	+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
	21	+
	22	+import networkx as nx
	23	+from sympy import Expr, Symbol
	24	+from contextlib import contextmanager
	25	+
	26	+from catlass_cppgen.common.data_type import DataType
	27	+from .evg.evg_definition import EVGArg, EVGDef
	28	+from .evg.node import (CastNode, ComputeNode, ConstantNode, LoadNode, NodeBase,
	29	+ NodeMetadata, StoreNode, TopoVisitorNode)
	30	+from .library import *
	31	+
	32	+_as_tuple = lambda x: x if isinstance(x, tuple) else (x,)
	33	+
	34	+def _get_tensor_element(tensor) -> DataType:
	35	+ """
	36	+ 统一获取 tensor 的 dtype 属性
	37	+ 仅支持 OpTensor
	38	+ """
	39	+ from catlass_cppgen.common.op_tensor import OpTensor
	40	+
	41	+ if not isinstance(tensor, OpTensor):
	42	+ raise TypeError(f"Unsupported tensor type: {type(tensor)}. Expected OpTensor")
	43	+ return tensor.dtype
	44	+
	45	+
	46	+class EpilogueVisitorGraph:
	47	+ """
	48	+ Helper class for constructing a DAG from Python EVG function
	49	+ """
	50	+ def __init__(self):
	51	+ self._graph = nx.DiGraph()
	52	+ self.nodes_map = {} # a little fragile
	53	+ self.compute_counter = itertools.count()
	54	+
	55	+ # used for dynamic shape, should be reimplemented with more mathematics if permute kind ops are supported
	56	+ self.symbol_shape_substitution_dict = {}
	57	+
	58	+ def add_node(self, node: NodeBase) -> None:
	59	+ self.check_not_exist(node)
	60	+ self._graph.add_node(node)
	61	+ self.nodes_map[node.name] = node
	62	+
	63	+ def get_node(self, node_or_name: Union[str, NodeBase]) -> NodeBase:
	64	+ if isinstance(node_or_name, str):
	65	+ return self.nodes_map[node_or_name]
	66	+ return node_or_name
	67	+
	68	+ def add_edge(
	69	+ self,
	70	+ src_node: Union[str, NodeBase],
	71	+ dst_node: Union[str, NodeBase],
	72	+ pos: int = 0,
	73	+ ):
	74	+ self.check_exist(src_node)
	75	+ self.check_exist(dst_node)
	76	+ src_node = self.get_node(src_node)
	77	+ dst_node = self.get_node(dst_node)
	78	+ self._graph.add_edge(src_node, dst_node, weight=pos)
	79	+
	80	+ def remove_node(self, node: NodeBase):
	81	+ self._graph.remove_node(node)
	82	+ del self.nodes_map[node.name]
	83	+
	84	+ def remove_edge(
	85	+ self, src_node: Union[str, NodeBase], dst_node: Union[str, NodeBase]
	86	+ ):
	87	+ src_node = self.get_node(src_node)
	88	+ dst_node = self.get_node(dst_node)
	89	+ self._graph.remove_edge(src_node, dst_node)
	90	+
	91	+ def has_node(self, node: NodeBase) -> bool:
	92	+ return self._graph.has_node(node)
	93	+
	94	+ def check_not_exist(self, node: Union[str, NodeBase]):
	95	+ if self.has_node(self.get_node(node)):
	96	+ raise SyntaxError(f"Variable '{str(node)}' is already defined before")
	97	+
	98	+ def check_exist(self, node: Union[str, NodeBase]):
	99	+ if not self.has_node(self.get_node(node)):
	100	+ raise SyntaxError(f"Variable '{str(node)}' is used before definiton")
	101	+
	102	+ def get_storage_nodes(self, nodes_element_dict):
	103	+ """
	104	+ Returns a dict, containing the infos of the total number of nodes for each type of element
	105	+ """
	106	+ for node in self.topological_nodes():
	107	+ if node.disabled or isinstance(node, StoreNode):
	108	+ continue
	109	+ if isinstance(node, TopoVisitorNode):
	110	+ node.subgraph.get_storage_nodes(nodes_element_dict)
	111	+ continue
	112	+ element = node.metadata.element
	113	+ if element in nodes_element_dict:
	114	+ nodes_element_dict[element] += 1
	115	+ else:
	116	+ nodes_element_dict[element] = 1
	117	+
	118	+ def to_networkx(self) -> nx.DiGraph:
	119	+ return self._graph
	120	+
	121	+ def topological_nodes(self) -> List[NodeBase]:
	122	+ return list(nx.lexicographical_topological_sort(self._graph))
	123	+
	124	+ def in_degree(self, node: NodeBase) -> int:
	125	+ return self._graph.in_degree(node)
	126	+
	127	+ def out_degree(self, node: NodeBase) -> int:
	128	+ return self._graph.out_degree(node)
	129	+
	130	+ def get_outputs(self, node: NodeBase) -> List[NodeBase]:
	131	+ return list(self._graph.successors(node))
	132	+
	133	+ def get_inputs(self, node: NodeBase) -> List[NodeBase]:
	134	+ return list(self._graph.predecessors(node))
	135	+
	136	+ def get_sorted_inputs(self, node: NodeBase) -> List[NodeBase]:
	137	+ input_nodes = {self.get_edge_pos(pnode, node): pnode for pnode in self.get_inputs(node)}
	138	+ return [input_nodes[key] for key in sorted(input_nodes.keys())]
	139	+
	140	+ def all_reachable_nodes(self, node: NodeBase) -> List[NodeBase]:
	141	+ return list(nx.dfs_preorder_nodes(self._graph, source=node))
	142	+
	143	+ def get_edge_pos(self, src_node: NodeBase, dst_node: NodeBase) -> int:
	144	+ return self._graph.get_edge_data(src_node, dst_node)["weight"]
	145	+
	146	+ def mark_output(self, name: str, traced_tensor):
	147	+ node = self.get_node(name)
	148	+ if not isinstance(node, StoreNode):
	149	+ raise ValueError(f"Only StoreNode can be marked as output. Got: {name}")
	150	+ node.metadata.is_output = True
	151	+ node.metadata.op = "auxstore"
	152	+ node.metadata.element = _get_tensor_element(traced_tensor)
	153	+ node.metadata.shape = traced_tensor.shape
	154	+ node.metadata.stride = traced_tensor.stride
	155	+
	156	+ def add_load_node(self, name: str, traced_tensor) -> str:
	157	+ if name is None:
	158	+ raise ValueError("Node name is not provided")
	159	+ if traced_tensor is None:
	160	+ raise ValueError(f"Input for {name} is not provided")
	161	+
	162	+ op = "accload" if name == "accum" else "auxload"
	163	+ if op == "accload":
	164	+ # Within dynamic shape mode, the symbolic shape shoule be sympy.Symbol type
	165	+ # Even in dynamic shape mode, the input shape might be (s1, 1) which is (sympy.Symbol, sympy.Int)
	166	+ # Obviously, we do not suppose to subtitute all 1s to "n"
	167	+ # NOTE: The whole shape expression maps to a single dynamic dimension tag (m/n),
	168	+ # e.g. shape '2s20' maps to 'm' (not '2m')
	169	+ if isinstance(traced_tensor.shape[0], Expr) and traced_tensor.shape[0].free_symbols:
	170	+ self.symbol_shape_substitution_dict[str(traced_tensor.shape[0])] = "m"
	171	+ if isinstance(traced_tensor.shape[1], Expr) and traced_tensor.shape[1].free_symbols:
	172	+ self.symbol_shape_substitution_dict[str(traced_tensor.shape[1])] = "n"
	173	+
	174	+ metadata = NodeMetadata(
	175	+ op=op,
	176	+ shape=traced_tensor.shape,
	177	+ element=_get_tensor_element(traced_tensor),
	178	+ )
	179	+ load_node = LoadNode(name, metadata)
	180	+ self.add_node(load_node)
	181	+ return name
	182	+
	183	+ def add_store_node(self, element: DataType, name: str):
	184	+ metadata = NodeMetadata(
	185	+ op="store",
	186	+ element=element,
	187	+ )
	188	+ node = StoreNode(name, metadata)
	189	+ self.add_node(node)
	190	+
	191	+ def add_compute_node(self, op, element: DataType, name=None):
	192	+ if name is None:
	193	+ name = f"compute_{next(self.compute_counter)}"
	194	+ metadata = NodeMetadata(
	195	+ op="compute",
	196	+ element=element,
	197	+ )
	198	+ compute_node = ComputeNode(
	199	+ name=name,
	200	+ metadata=metadata,
	201	+ fn=op,
	202	+ )
	203	+ self.add_node(compute_node)
	204	+ return name
	205	+
	206	+ def add_constant_node(self, value, dtype):
	207	+ if isinstance(dtype, DataType):
	208	+ element = dtype
	209	+ else:
	210	+ element = DataType.from_dtype(dtype)
	211	+ name = f"constant_{value}_{next(self.compute_counter)}"
	212	+ metadata = NodeMetadata(
	213	+ op="constant",
	214	+ element=element,
	215	+ )
	216	+ constant_node = ConstantNode(
	217	+ name=name,
	218	+ metadata=metadata,
	219	+ value=value,
	220	+ )
	221	+ self.add_node(constant_node)
	222	+ return name
	223	+
	224	+ def add_cast_node(self, dst_type, src_type):
	225	+ if isinstance(dst_type, DataType):
	226	+ dst_element = dst_type
	227	+ else:
	228	+ dst_element = DataType.from_dtype(dst_type)
	229	+ if isinstance(src_type, DataType):
	230	+ src_element = src_type
	231	+ else:
	232	+ src_element = DataType.from_dtype(src_type)
	233	+ name = f"cast_{next(self.compute_counter)}"
	234	+ metadata = NodeMetadata(
	235	+ op="cast",
	236	+ element=dst_element,
	237	+ )
	238	+ cast_node = CastNode(
	239	+ name=name,
	240	+ metadata=metadata,
	241	+ from_element=src_element,
	242	+ to_element=dst_element,
	243	+ )
	244	+ self.add_node(cast_node)
	245	+ return name
	246	+
	247	+
	248	+
	249	+class PythonEVGParser(EpilogueVisitorGraph, ast.NodeVisitor):
	250	+ """
	251	+ Transform a Python EVG function to DAG
	252	+ """
	253	+ def __init__(self):
	254	+ """
	255	+ e.g. example_inputs = {
	256	+ "accum": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	257	+ ...
	258	+ }
	259	+ """
	260	+ super().__init__()
	261	+ self.identity_inplace_dict: Dict[str, str] = {}
	262	+ self.visiting_return = False
	263	+
	264	+ @staticmethod
	265	+ def ast_op_bindings(op):
	266	+ mapping = {
	267	+ ast.Add: EpilogueOp.Add,
	268	+ ast.Sub: EpilogueOp.Sub,
	269	+ ast.Mult: EpilogueOp.Mul,
	270	+ ast.Div: EpilogueOp.Div,
	271	+ "relu": EpilogueOp.Relu,
	272	+ "leakyRelu": EpilogueOp.LeakyRelu,
	273	+ "Prelu": EpilogueOp.Prelu,
	274	+ "max": EpilogueOp.Max,
	275	+ "min": EpilogueOp.Min,
	276	+ "sigmoid": EpilogueOp.Sigmoid,
	277	+ "silu": EpilogueOp.Silu,
	278	+ }
	279	+ return mapping[op]
	280	+
	281	+ def parse(self, fn_src, example_inputs):
	282	+ self.example_inputs = example_inputs
	283	+ self.source = textwrap.dedent(fn_src)
	284	+ self.ast = ast.parse(self.source)
	285	+ self.visit(self.ast)
	286	+
	287	+ @contextmanager
	288	+ def _return_context(self):
	289	+ self.visiting_return = True # Set visit Flag to be True
	290	+ try:
	291	+ yield
	292	+ finally:
	293	+ self.visiting_return = False
	294	+
	295	+ def visit_FunctionDef(self, node: ast.FunctionDef):
	296	+ """Visit FunctionDef in ast.NodeVisiter"""
	297	+ # processs args
	298	+ for arg in node.args.args:
	299	+ self.visit(arg)
	300	+
	301	+ # process expression
	302	+ for expr in node.body:
	303	+ self.visit(expr)
	304	+
	305	+ def visit_arg(self, node: ast.arg):
	306	+ arg_name = node.arg
	307	+ try:
	308	+ input_tensor = self.example_inputs[arg_name]
	309	+ except KeyError as e:
	310	+ raise RuntimeError(f"Input for {arg_name} is not provided") from e
	311	+
	312	+ self.add_load_node(arg_name, input_tensor)
	313	+
	314	+ def visit_Name(self, node: ast.Name):
	315	+ return node.id
	316	+
	317	+ def visit_Attribute(self, node: ast.Attribute):
	318	+ """处理属性访问，如 DataType.FLOAT16"""
	319	+ value = self.visit(node.value)
	320	+ if value == "DataType":
	321	+ return getattr(DataType, node.attr, None)
	322	+ return f"{value}.{node.attr}"
	323	+
	324	+ def visit_Constant(self, node: ast.Constant):
	325	+ return node.value
	326	+
	327	+ def visit_Tuple(self, node: ast.Tuple):
	328	+ return tuple(self.visit(elt) for elt in node.elts)
	329	+
	330	+ def visit_keyword(self, node: ast.keyword):
	331	+ return {node.arg: self.visit(node.value)}
	332	+
	333	+ def visit_Return(self, node: ast.Return):
	334	+ with self._return_context():
	335	+ results = self.visit(node.value)
	336	+
	337	+ for res in _as_tuple(results):
	338	+ try:
	339	+ traced_tensor = self.example_inputs[res]
	340	+ except KeyError as e:
	341	+ raise RuntimeError(f"Input for {res} is not provided") from e
	342	+ self.mark_output(res, traced_tensor)
	343	+
	344	+ def visit_BinOp(self, node: ast.BinOp):
	345	+ if self.visiting_return:
	346	+ raise SyntaxError("Return value cannot be an expression")
	347	+ op = self.ast_op_bindings(type(node.op))
	348	+
	349	+ # all elements of args should be same for a single compute node.
	350	+ lhs = self.visit(node.left)
	351	+ rhs = self.visit(node.right)
	352	+ input_element = self.get_node(lhs).metadata.element
	353	+ name = self.add_compute_node(op, input_element)
	354	+ self.add_edge(lhs, name, pos=0)
	355	+ self.add_edge(rhs, name, pos=1)
	356	+ return name
	357	+
	358	+ def visit_Assign(self, node: ast.BinOp):
	359	+ target = self.visit(node.targets[0])
	360	+ value = self.visit(node.value)
	361	+ # Create store node
	362	+ input_element = self.get_node(value).metadata.element
	363	+ self.add_store_node(input_element, target)
	364	+ self.add_edge(value, target)
	365	+ return target
	366	+
	367	+ def visit_Call(self, node: ast.Call):
	368	+ if self.visiting_return:
	369	+ raise SyntaxError("Return value cannot be an expression")
	370	+ func = self.visit(node.func)
	371	+ args = [self.visit(arg) for arg in node.args]
	372	+
	373	+ if func == "constant":
	374	+ # constant ops look like constant(value, dtype)
	375	+ # dtype can be DataType enum or string like "FLOAT16", "FLOAT"
	376	+ dtype = DataType(args[1])
	377	+ name = self.add_constant_node(args[0], dtype)
	378	+ return name
	379	+
	380	+ if func == "cast":
	381	+ # cast ops look like cast(input, dst_type, src_type)
	382	+ # dst_type and src_type can be DataType enum or string like "FLOAT16", "FLOAT"
	383	+ dst_type = DataType(args[1])
	384	+ src_type = DataType(args[2])
	385	+ # convert string to DataType if needed
	386	+ name = self.add_cast_node(dst_type, src_type)
	387	+ self.add_edge(args[0], name, pos=0)
	388	+ return name
	389	+
	390	+ op = self.ast_op_bindings(func)
	391	+ # all elements of args should be same for a single compute node.
	392	+ input_element = self.get_node(args[0]).metadata.element
	393	+ name = self.add_compute_node(op, input_element)
	394	+
	395	+ # add edges
	396	+ for idx, arg in enumerate(args):
	397	+ self.add_edge(arg, name, pos=idx)
	398	+ return name
	399	+
	400	+
	401	+class EVGArgRenames:
	402	+ """Handles mapping buffer names to variable names in the cpp kernel signature and body"""
	403	+
	404	+ def __init__(self) -> None:
	405	+ self.buf_renames: dict[str, str] = {}
	406	+
	407	+ def new_name(self, name: str) -> str:
	408	+ if name in self.buf_renames:
	409	+ return self.buf_renames[name]
	410	+ else:
	411	+ new_name = f"ptr_{len(self.buf_renames)}"
	412	+ self.buf_renames[name] = new_name
	413	+ return new_name
	414	+
	415	+ def get(self, name: str) -> str:
	416	+ return self.buf_renames.get(name)
	417	+
	418	+
	419	+def evg(
	420	+ fn_src: str,
	421	+ example_inputs,
	422	+ # accum_type: DataType,
	423	+ # output_type: DataType,
	424	+ # tile_description: TileDescription,
	425	+ # name_to_buffer: dict[str, Buffer],
	426	+ # size_hint_fn: Callable[[Union[Expr, int]], int],
	427	+ # **kwargs: dict[str, Any],
	428	+) -> tuple[str, str, str, EVGArgRenames]:
	429	+ # Transfer python_func to a DAG
	430	+ parser = PythonEVGParser()
	431	+ parser.parse(fn_src, example_inputs)
	432	+
	433	+ # Generate Catlass EVG definition
	434	+ evg_def = EVGDef(parser)
	435	+ evg_str, callback_name = evg_def.definition()
	436	+
	437	+ # Generate Catlass EVG arguments
	438	+ evg_arg = EVGArg(parser)
	439	+ evg_args, arg_renames = evg_arg.generate_graph_args()
	440	+
	441	+ # FIXME: move to EVG args
	442	+ # Generate compute node length, which is used as "vector tiling"
	443	+ evg_compute_length = evg_arg.generate_compute_length()
	444	+
	445	+ # combine evg args and evg_compute_length as the entire args
	446	+ evg_args += evg_compute_length
	447	+
	448	+ return callback_name, evg_args, evg_str, arg_renames

Apython/catlass_cppgen/catlass_cppgen/catlass/gemm/dispatch_policy.py+292-0

@@ -0,0 +1,292 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+import inspect
	11	+from abc import ABC
	12	+from typing import List, Tuple, Union
	13	+from catlass_cppgen.catlass.arch.arch import Arch
	14	+from catlass_cppgen.common.utils import _get_cpp_value, _snake_to_camel, _get_cpp_type
	15	+
	16	+
	17	+class MmadBase(ABC):
	18	+ """Base class for MMAD policies."""
	19	+ def __init__(self, arch_tag: Arch, async_: bool):
	20	+ self.arch_tag = arch_tag
	21	+ self.async_ = async_
	22	+
	23	+ def to_cpp(self, const_mode: bool = False) -> Union[str, Tuple[List[str], str]]:
	24	+ """生成 C++ 代码字符串.
	25	+
	26	+ 根据类名和属性自动生成 C++ 模板代码：
	27	+ - 类名转换为 Gemm::ClassName 格式
	28	+ - 模板参数取自 __init__ 参数（排除self），按定义顺序排列
	29	+ - arch_tag 若存在则取其 value 作为模板参数
	30	+ - 布尔值自动转换为 C++ 风格的 true/false
	31	+
	32	+ :param const_mode: 当为 True 时，返回常量声明列表和变量名形式的模板字符串；当为 False 时，返回硬编码的模板字符串
	33	+ :return: 当 const_mode=False 时返回 C++ 模板代码字符串；当 const_mode=True 时返回 (常量声明列表, 模板字符串) 元组
	34	+ """
	35	+ cls_name = f"Gemm::{self.__class__.__name__}"
	36	+ template_params: List[str] = []
	37	+ if const_mode:
	38	+ const_declarations: List[str] = []
	39	+ try:
	40	+ init_params = list(inspect.signature(self.__init__).parameters.keys())
	41	+ for param_name in init_params:
	42	+ if param_name in {"self", "async_"} or not hasattr(self, param_name):
	43	+ continue
	44	+ if param_name == "arch_tag":
	45	+ template_params.append("ArchTag")
	46	+ continue
	47	+ param_value = getattr(self, param_name)
	48	+ var_name = _snake_to_camel(param_name)
	49	+ const_declarations.append(f"constexpr {_get_cpp_type(param_value)} {var_name} = {_get_cpp_value(param_value)};")
	50	+ template_params.append(var_name)
	51	+ except (ValueError, TypeError, AttributeError):
	52	+ pass
	53	+ params_str = ", ".join(template_params)
	54	+ return (const_declarations, f"{cls_name}<{params_str}>" if params_str else f"{cls_name}<>")
	55	+ else:
	56	+ try:
	57	+ init_params = list(inspect.signature(self.__init__).parameters.keys())
	58	+ for param_name in init_params:
	59	+ if param_name in {"self", "async_"} or not hasattr(self, param_name):
	60	+ continue
	61	+ if param_name == "arch_tag":
	62	+ param_value = getattr(self, param_name).value
	63	+ else:
	64	+ param_value = getattr(self, param_name)
	65	+ template_params.append(_get_cpp_value(param_value))
	66	+ except (ValueError, TypeError, AttributeError):
	67	+ pass
	68	+ params_str = ", ".join(template_params)
	69	+ return f"{cls_name}<{params_str}>" if params_str else f"{cls_name}<>"
	70	+
	71	+
	72	+# Block Mmad Policies
	73	+
	74	+class MmadAtlasA2(MmadBase):
	75	+ """MMAD policy for AtlasA2 architecture, synchronous."""
	76	+ def __init__(self):
	77	+ super().__init__(Arch.AtlasA2, False)
	78	+
	79	+
	80	+class MmadAtlasA2Async(MmadBase):
	81	+ """MMAD policy for AtlasA2 architecture, asynchronous."""
	82	+ def __init__(self):
	83	+ super().__init__(Arch.AtlasA2, True)
	84	+
	85	+
	86	+class MmadAtlasA2Pingpong(MmadAtlasA2):
	87	+ """MMAD policy with pingpong staging."""
	88	+ def __init__(self, enable_unit_flag: bool = False):
	89	+ super().__init__()
	90	+ self.stages = 2
	91	+ self.enable_unit_flag = enable_unit_flag
	92	+
	93	+
	94	+class MmadAtlasA2PingpongSliceKWithPrologue(MmadAtlasA2):
	95	+ """MMAD policy with pingpong staging and sliced K dimension."""
	96	+ def __init__(self, enable_unit_flag: bool = False):
	97	+ super().__init__()
	98	+ self.stages = 2
	99	+ self.enable_unit_flag = enable_unit_flag
	100	+
	101	+
	102	+class MmadAtlasA2PingPongWithPrologue(MmadAtlasA2):
	103	+ """MMAD policy with pingpong staging and prologue."""
	104	+ def __init__(self, enable_unit_flag: bool = False):
	105	+ super().__init__()
	106	+ self.stages = 2
	107	+ self.enable_unit_flag = enable_unit_flag
	108	+
	109	+
	110	+class MmadAtlasA2Preload(MmadAtlasA2):
	111	+ """MMAD policy with preload capability."""
	112	+ def __init__(self, enable_unit_flag: bool = False, enable_shuffle_k: bool = False):
	113	+ super().__init__()
	114	+ self.stages = 2
	115	+ self.enable_unit_flag = enable_unit_flag
	116	+ self.enable_shuffle_k = enable_shuffle_k
	117	+
	118	+
	119	+class MmadAtlasA2PreloadAsync(MmadAtlasA2Async):
	120	+ """MMAD policy with async preload capability."""
	121	+ def __init__(
	122	+ self,
	123	+ preload_stages: int,
	124	+ l1_stages: int,
	125	+ l0a_stages: int,
	126	+ l0b_stages: int,
	127	+ l0c_stages: int,
	128	+ enable_unit_flag: bool = False,
	129	+ enable_shuffle_k: bool = False
	130	+ ):
	131	+ super().__init__()
	132	+ self.preload_stages = preload_stages
	133	+ self.l1_stages = l1_stages
	134	+ self.l0a_stages = l0a_stages
	135	+ self.l0b_stages = l0b_stages
	136	+ self.l0c_stages = l0c_stages
	137	+ self.enable_unit_flag = enable_unit_flag
	138	+ self.enable_shuffle_k = enable_shuffle_k
	139	+
	140	+
	141	+class MmadAtlasA2PreloadAsyncWithCallback(MmadAtlasA2PreloadAsync):
	142	+ """MMAD policy with async preload and callback capability."""
	143	+ def __init__(
	144	+ self,
	145	+ preload_stages: int,
	146	+ l1_stages: int,
	147	+ l0a_stages: int,
	148	+ l0b_stages: int,
	149	+ l0c_stages: int,
	150	+ enable_unit_flag: bool = False,
	151	+ enable_shuffle_k: bool = False
	152	+ ):
	153	+ super().__init__(
	154	+ preload_stages,
	155	+ l1_stages,
	156	+ l0a_stages,
	157	+ l0b_stages,
	158	+ l0c_stages,
	159	+ enable_unit_flag,
	160	+ enable_shuffle_k
	161	+ )
	162	+
	163	+
	164	+class GemmAtlasA2(MmadAtlasA2):
	165	+ """GEMM policy for AtlasA2 architecture."""
	166	+ def __init__(self, enable_unit_flag: bool = False, enable_shuffle_k: bool = False, enable_abba: bool = False):
	167	+ super().__init__()
	168	+ self.stages = 2
	169	+ self.enable_unit_flag = enable_unit_flag
	170	+ self.enable_shuffle_k = enable_shuffle_k
	171	+ self.enable_abba = enable_abba
	172	+
	173	+
	174	+class GemvAtlasA2(MmadAtlasA2):
	175	+ """GEMV policy for AtlasA2 architecture."""
	176	+ def __init__(self):
	177	+ super().__init__()
	178	+ self.stages = 2
	179	+
	180	+
	181	+class MmadAtlasA2PingpongBias(MmadAtlasA2):
	182	+ """MMAD policy with pingpong staging and bias support."""
	183	+ def __init__(self, enable_unit_flag: bool = False):
	184	+ super().__init__()
	185	+ self.stages = 2
	186	+ self.enable_unit_flag = enable_unit_flag
	187	+
	188	+
	189	+class MmadAtlasA2FullLoadA(MmadAtlasA2):
	190	+ """MMAD policy with full load of matrix A."""
	191	+ def __init__(self, enable_unit_flag: bool = False):
	192	+ super().__init__()
	193	+ self.stages = 2
	194	+ self.enable_unit_flag = enable_unit_flag
	195	+
	196	+
	197	+class MmadAtlasA2W8A16(MmadAtlasA2):
	198	+ """MMAD policy with W8A16 configuration."""
	199	+ def __init__(self, enable_unit_flag: bool = False, enable_shuffle_k: bool = False):
	200	+ super().__init__()
	201	+ self.stages = 2
	202	+ self.enable_unit_flag = enable_unit_flag
	203	+ self.enable_shuffle_k = enable_shuffle_k
	204	+
	205	+
	206	+class MmadAtlasA2DynamicCommon(MmadAtlasA2):
	207	+ """MMAD policy with dynamic common configuration."""
	208	+ def __init__(self, enable_unit_flag: bool = False, enable_shuffle_k: bool = False):
	209	+ super().__init__()
	210	+ self.stages = 2
	211	+ self.enable_unit_flag = enable_unit_flag
	212	+ self.enable_shuffle_k = enable_shuffle_k
	213	+
	214	+
	215	+class MmadAtlasA2Small(MmadAtlasA2):
	216	+ """MMAD policy for small problem sizes."""
	217	+ def __init__(self, stages: int, enable_unit_flag: bool = False, enable_shuffle_k: bool = False):
	218	+ super().__init__()
	219	+ self.stages = stages
	220	+ self.enable_unit_flag = enable_unit_flag
	221	+ self.enable_shuffle_k = enable_shuffle_k
	222	+
	223	+
	224	+# Generic MMAD policies that work with different architectures
	225	+
	226	+class MmadPingpong(MmadBase):
	227	+ """Generic MMAD policy with pingpong staging."""
	228	+ def __init__(
	229	+ self,
	230	+ arch_tag: Arch,
	231	+ enable_unit_flag: bool = False,
	232	+ use_hf32_mode: bool = False,
	233	+ l0c_stages: int = 1,
	234	+ enable_l1_resident: bool = False,
	235	+ l1a_stages: int = 2,
	236	+ l1b_stages: int = 2,
	237	+ l0a_stages: int = 2,
	238	+ l0b_stages: int = 2
	239	+ ):
	240	+ super().__init__(arch_tag, False)
	241	+ self.stages = 2 # May be removed
	242	+ self.enable_unit_flag = enable_unit_flag
	243	+ self.use_hf32_mode = use_hf32_mode
	244	+ self.l0c_stages = l0c_stages
	245	+ self.enable_l1_resident = enable_l1_resident
	246	+ self.l1a_stages = l1a_stages
	247	+ self.l1b_stages = l1b_stages
	248	+ self.l0a_stages = l0a_stages
	249	+ self.l0b_stages = l0b_stages
	250	+
	251	+
	252	+class MmadPreloadAsyncWithCallback(MmadBase):
	253	+ """Generic MMAD policy with async preload and callback."""
	254	+ def __init__(
	255	+ self,
	256	+ arch_tag: Arch,
	257	+ preload_stages: int,
	258	+ l1a_stages: int,
	259	+ l1b_stages: int,
	260	+ l0a_stages: int,
	261	+ l0b_stages: int,
	262	+ l0c_stages: int,
	263	+ enable_unit_flag: bool,
	264	+ enable_shuffle_k: bool,
	265	+ use_hf32_mode: bool = False,
	266	+ enable_l1_resident: bool = False
	267	+ ):
	268	+ super().__init__(arch_tag, True)
	269	+ self.preload_stages = preload_stages
	270	+ self.l1a_stages = l1a_stages
	271	+ self.l1b_stages = l1b_stages
	272	+ self.l0a_stages = l0a_stages
	273	+ self.l0b_stages = l0b_stages
	274	+ self.l0c_stages = l0c_stages
	275	+ self.enable_unit_flag = enable_unit_flag
	276	+ self.enable_shuffle_k = enable_shuffle_k
	277	+ self.use_hf32_mode = use_hf32_mode
	278	+ self.enable_l1_resident = enable_l1_resident
	279	+
	280	+
	281	+class MmadMultiBatch(MmadBase):
	282	+ """Generic MMAD policy for multi-batch operations."""
	283	+ def __init__(
	284	+ self,
	285	+ arch_tag: Arch,
	286	+ use_hf32_mode: bool = False,
	287	+ l0c_stages: int = 2
	288	+ ):
	289	+ super().__init__(arch_tag, False)
	290	+ self.stages = 2
	291	+ self.use_hf32_mode = use_hf32_mode
	292	+ self.l0c_stages = l0c_stages

Apython/catlass_cppgen/catlass_cppgen/catlass/gemm_coord.py+48-0

@@ -0,0 +1,48 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from dataclasses import dataclass
	11	+
	12	+
	13	+@dataclass
	14	+class Shape:
	15	+ m: int
	16	+ n: int
	17	+ k: int
	18	+
	19	+ def __str__(self):
	20	+ return "Shape<Int<{m}>, Int<{n}>, Int<{k}>>".format(
	21	+ m=self.m, n=self.n, k=self.k
	22	+ )
	23	+
	24	+
	25	+@dataclass
	26	+class GemmShape:
	27	+ m: int
	28	+ n: int
	29	+ k: int
	30	+
	31	+ def __str__(self):
	32	+ return "GemmShape<{m}, {n}, {k}>".format(m=self.m, n=self.n, k=self.k)
	33	+
	34	+ def tla(self):
	35	+ return Shape(self.m, self.n, self.k)
	36	+
	37	+
	38	+@dataclass
	39	+class GemmCoord:
	40	+ m: int
	41	+ n: int
	42	+ k: int
	43	+
	44	+ def __str__(self):
	45	+ return "GemmCoord"
	46	+
	47	+ def tla(self):
	48	+ return NotImplementedError("GemmCoord tla is not implemented")

Apython/catlass_cppgen/catlass_cppgen/catlass/layout/__init__.py+35-0

@@ -0,0 +1,35 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from catlass_cppgen.catlass.layout.layout import (
	11	+ Layout,
	12	+ Coord,
	13	+ RowMajor,
	14	+ ColumnMajor,
	15	+ PaddingRowMajor,
	16	+ PaddingColumnMajor,
	17	+ VectorLayout,
	18	+ nZ,
	19	+ zN,
	20	+ zZ,
	21	+ nN,
	22	+)
	23	+
	24	+__all__ = [
	25	+ "Coord",
	26	+ "RowMajor",
	27	+ "ColumnMajor",
	28	+ "PaddingRowMajor",
	29	+ "PaddingColumnMajor",
	30	+ "VectorLayout",
	31	+ "nZ",
	32	+ "zN",
	33	+ "zZ",
	34	+ "nN",
	35	+]

Apython/catlass_cppgen/catlass_cppgen/catlass/layout/coord.py+14-0

@@ -0,0 +1,14 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from typing import Iterable, Union
	11	+
	12	+class Coord:
	13	+ def __init__(self, value: Iterable[int]):
	14	+ pass

Apython/catlass_cppgen/catlass_cppgen/catlass/layout/layout.py+157-0

@@ -0,0 +1,157 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from abc import ABC, abstractmethod
	11	+from math import prod
	12	+from typing import Iterable
	13	+
	14	+class Coord:
	15	+ def __init__(self, value: Iterable[int]):
	16	+ self.idx = tuple(value)
	17	+
	18	+class Layout(ABC):
	19	+ value: str = ""
	20	+
	21	+ def __init__(self, shape: Iterable[int], stride: Iterable[int] = None):
	22	+ self.shape = tuple(shape)
	23	+ self.stride = tuple(stride) if stride else tuple(1 for _ in range(len(shape)))
	24	+
	25	+ @property
	26	+ def capacity(self) -> int:
	27	+ return prod(self.shape)
	28	+
	29	+ def get_offset(self, coord: Iterable[int]) -> int:
	30	+ pass
	31	+
	32	+ @abstractmethod
	33	+ def is_need_padding(self, align: int) -> bool:
	34	+ pass
	35	+
	36	+
	37	+class RowMajor(Layout):
	38	+ value = "layout::RowMajor"
	39	+
	40	+ def __init__(self, shape: tuple[int, int]):
	41	+ super().__init__(shape, (shape[1], 1))
	42	+
	43	+ def is_need_padding(self, align: int) -> bool:
	44	+ if self.stride[0] < 65536:
	45	+ return self.stride[0] % align != 0
	46	+ else:
	47	+ return True
	48	+
	49	+ def get_padding_layout(self, align: int) -> Layout:
	50	+ if self.is_need_padding(align):
	51	+ return PaddingRowMajor(
	52	+ self.shape[0],
	53	+ self.shape[1],
	54	+ (self.shape[1] + align - 1) // align * align,
	55	+ )
	56	+ else:
	57	+ return self
	58	+
	59	+
	60	+class ColumnMajor(Layout):
	61	+ value = "layout::ColumnMajor"
	62	+
	63	+ def __init__(self, shape: tuple[int, int]):
	64	+ super().__init__(shape, (1, shape[0]))
	65	+
	66	+ def is_need_padding(self, align: int) -> bool:
	67	+ if self.stride[0] < 65536:
	68	+ return self.stride[0] % align != 0
	69	+ else:
	70	+ return True
	71	+
	72	+ def get_padding_layout(self, align: int) -> Layout:
	73	+ if self.is_need_padding(align):
	74	+ return PaddingColumnMajor(
	75	+ self.shape[0],
	76	+ self.shape[1],
	77	+ (self.shape[0] + align - 1) // align * align,
	78	+ )
	79	+ else:
	80	+ return self
	81	+
	82	+
	83	+class PaddingRowMajor(Layout):
	84	+ value = "layout::PaddingRowMajor"
	85	+
	86	+ def __init__(self, shape: tuple[int, int], block_shape: tuple[int, int]):
	87	+ super().__init__(
	88	+ (
	89	+ block_shape[0],
	90	+ (shape[0] + block_shape[0] - 1)//(block_shape[0]),
	91	+ block_shape[1],
	92	+ (shape[1] + block_shape[1] - 1)//(block_shape[1])
	93	+ ),
	94	+ (
	95	+ block_shape[1],
	96	+ block_shape[0] * block_shape[1] * (shape[1] + block_shape[1] - 1)//(block_shape[1]),
	97	+ 1,
	98	+ block_shape[0] * block_shape[1]
	99	+ )
	100	+ )
	101	+
	102	+ def is_need_padding(self, align: int) -> bool:
	103	+ return False
	104	+
	105	+
	106	+class PaddingColumnMajor(Layout):
	107	+ value = "layout::PaddingColumnMajor"
	108	+
	109	+ def __init__(self, shape: tuple[int, int], block_shape: tuple[int, int]):
	110	+ super().__init__(
	111	+ (
	112	+ block_shape[0],
	113	+ (shape[0] + block_shape[0] - 1)//(block_shape[0]),
	114	+ block_shape[1],
	115	+ (shape[1] + block_shape[1] - 1)//(block_shape[1])
	116	+ ),
	117	+ (
	118	+ 1,
	119	+ block_shape[0] * block_shape[1],
	120	+ block_shape[1],
	121	+ block_shape[1] * block_shape[0] * (shape[0] + block_shape[0] - 1)//(block_shape[0])
	122	+ )
	123	+ )
	124	+
	125	+ def is_need_padding(self, align: int) -> bool:
	126	+ return False
	127	+
	128	+
	129	+class VectorLayout(Layout):
	130	+ value = "layout::VectorLayout"
	131	+
	132	+ def __init__(self, shape: int, stride: int = 1):
	133	+ super().__init__((shape,), (stride,))
	134	+
	135	+ def is_need_padding(self, align: int) -> bool:
	136	+ """检查向量布局是否需要 padding"""
	137	+ return False
	138	+
	139	+class PrivateLayout(Layout):
	140	+ def is_need_padding(self, align: int) -> bool:
	141	+ return False
	142	+
	143	+
	144	+class nZ(PrivateLayout):
	145	+ value = "layout::nZ"
	146	+
	147	+
	148	+class zN(PrivateLayout):
	149	+ value = "layout::zN"
	150	+
	151	+
	152	+class zZ(PrivateLayout):
	153	+ value = "layout::zZ"
	154	+
	155	+
	156	+class nN(PrivateLayout):
	157	+ value = "layout::nN"

Apython/catlass_cppgen/catlass_cppgen/catlass/library.py+213-0

@@ -0,0 +1,213 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+import enum
	11	+
	12	+from enum import auto as enum_auto
	13	+
	14	+
	15	+class LayoutType(enum.Enum):
	16	+ ColumnMajor = enum_auto()
	17	+ RowMajor = enum_auto()
	18	+ VectorLayout = enum_auto()
	19	+ nZ = enum_auto()
	20	+ zN = enum_auto()
	21	+ zZ = enum_auto()
	22	+ nN = enum_auto()
	23	+
	24	+
	25	+#
	26	+LayoutTag = {
	27	+ LayoutType.ColumnMajor: "Catlass::layout::ColumnMajor",
	28	+ LayoutType.RowMajor: "Catlass::layout::RowMajor",
	29	+ LayoutType.VectorLayout: "Catlass::layout::VectorLayout",
	30	+ LayoutType.nZ: "Catlass::layout::nZ",
	31	+ LayoutType.zN: "Catlass::layout::zN",
	32	+ LayoutType.zZ: "Catlass::layout::zZ",
	33	+ LayoutType.nN: "Catlass::layout::nN",
	34	+}
	35	+
	36	+
	37	+class BroadcastType(enum.Enum):
	38	+ RowBroadcast = enum_auto()
	39	+ ColBroadcast = enum_auto()
	40	+ RowColBroadcast = enum_auto()
	41	+ NoBroadcast = enum_auto()
	42	+
	43	+
	44	+BroadcastTag = {
	45	+ BroadcastType.RowBroadcast: "RowBroadcast",
	46	+ BroadcastType.ColBroadcast: "ColBroadcast",
	47	+ BroadcastType.RowColBroadcast: "RowAndColBroadcast",
	48	+ BroadcastType.NoBroadcast: "NoBroadcast",
	49	+}
	50	+
	51	+
	52	+class EpilogueOp(enum.Enum):
	53	+ # unary op
	54	+ Cast = enum_auto()
	55	+ Exp = enum_auto()
	56	+ Reciprocal = enum_auto()
	57	+ Sqrt = enum_auto()
	58	+
	59	+ # binary op
	60	+ Add = enum_auto()
	61	+ Adds = enum_auto() # scalar ver
	62	+ Div = enum_auto()
	63	+ Max = enum_auto()
	64	+ Min = enum_auto()
	65	+ Mul = enum_auto()
	66	+ Muls = enum_auto() # scalar ver
	67	+ Sub = enum_auto()
	68	+
	69	+ # activation op
	70	+ LeakyRelu = enum_auto()
	71	+ Prelu = enum_auto()
	72	+ Relu = enum_auto()
	73	+ Rsqrt = enum_auto()
	74	+ Sigmoid = enum_auto()
	75	+ Silu = enum_auto()
	76	+
	77	+ # Binary-tensor op
	78	+ Maxs = enum_auto()
	79	+ Mins = enum_auto()
	80	+ AddRelu = enum_auto()
	81	+
	82	+
	83	+EpilogueOpTag = {
	84	+ # unary op
	85	+ EpilogueOp.Cast: "Catlass::Epilogue::Fusion::Cast",
	86	+ EpilogueOp.Exp: "Catlass::Epilogue::Fusion::Exp",
	87	+ EpilogueOp.Sqrt: "Catlass::Epilogue::Fusion::Sqrt",
	88	+ EpilogueOp.Rsqrt: "Catlass::Epilogue::Fusion::Rsqrt",
	89	+ EpilogueOp.Reciprocal: "Catlass::Epilogue::Fusion::Reciprocal",
	90	+ # binary op
	91	+ EpilogueOp.Add: "Catlass::Epilogue::Fusion::Add",
	92	+ EpilogueOp.Adds: "Catlass::Epilogue::Fusion::Adds",
	93	+ EpilogueOp.Div: "Catlass::Epilogue::Fusion::Div",
	94	+ EpilogueOp.Max: "Catlass::Epilogue::Fusion::Max",
	95	+ EpilogueOp.Min: "Catlass::Epilogue::Fusion::Min",
	96	+ EpilogueOp.Mul: "Catlass::Epilogue::Fusion::Mul",
	97	+ EpilogueOp.Muls: "Catlass::Epilogue::Fusion::Muls",
	98	+ EpilogueOp.Sub: "Catlass::Epilogue::Fusion::Sub",
	99	+
	100	+ # activation op
	101	+ EpilogueOp.LeakyRelu: "Catlass::Epilogue::Fusion::LeakyRelu",
	102	+ EpilogueOp.Prelu: "Catlass::Epilogue::Fusion::Prelu",
	103	+ EpilogueOp.Relu: "Catlass::Epilogue::Fusion::Relu",
	104	+ EpilogueOp.Sigmoid: "Catlass::Epilogue::Fusion::Sigmoid",
	105	+ EpilogueOp.Silu: "Catlass::Epilogue::Fusion::Silu",
	106	+
	107	+ # tensor
	108	+ EpilogueOp.Maxs: "Catlass::Epilogue::Fusion::Maxs",
	109	+ EpilogueOp.Mins: "Catlass::Epilogue::Fusion::Mins",
	110	+ EpilogueOp.AddRelu: "Catlass::Epilogue::Fusion::AddRelu",
	111	+}
	112	+
	113	+
	114	+EpilogueOpVectorToScalar = {
	115	+ EpilogueOp.Add: EpilogueOp.Adds,
	116	+ EpilogueOp.Mul: EpilogueOp.Muls,
	117	+}
	118	+
	119	+
	120	+EpilogueScalarOp = {
	121	+ EpilogueOp.Adds,
	122	+ EpilogueOp.Muls,
	123	+}
	124	+
	125	+
	126	+class CastType(enum.Enum):
	127	+ NONE = enum_auto() # When there is precision loss in conversion, it means RINT mode; when there is no precision loss, it means no rounding
	128	+ RINT = enum_auto() # round to nearest even (bankers' rounding)
	129	+ FLOOR = enum_auto() # round towards negative infinity
	130	+ CEIL = enum_auto() # round towards positive infinity
	131	+ ROUND = enum_auto() # round half away from zero
	132	+ TRUNC = enum_auto() # round half away from zero
	133	+ ODD = enum_auto() # Von Neumann rounding, round to nearest odd
	134	+
	135	+CastTypeTag = {
	136	+ CastType.NONE: "AscendC::RoundMode::CAST_NONE",
	137	+ CastType.RINT: "AscendC::RoundMode::CAST_RINT",
	138	+ CastType.FLOOR: "AscendC::RoundMode::CAST_FLOOR",
	139	+ CastType.CEIL: "AscendC::RoundMode::CAST_CEIL",
	140	+ CastType.ROUND: "AscendC::RoundMode::CAST_ROUND",
	141	+ CastType.TRUNC: "AscendC::RoundMode::CAST_TRUNC",
	142	+ CastType.ODD: "AscendC::RoundMode::CAST_ODD",
	143	+}
	144	+
	145	+
	146	+class TileDescription:
	147	+ def __init__(self, L1TileShape, L0TileShape):
	148	+ self.l1_tile_shape = list(L1TileShape)
	149	+ self.l0_tile_shape = list(L0TileShape)
	150	+
	151	+ @property
	152	+ def l1_m(self):
	153	+ return self.l1_tile_shape[0]
	154	+
	155	+ @property
	156	+ def l1_n(self):
	157	+ return self.l1_tile_shape[1]
	158	+
	159	+ @property
	160	+ def l1_k(self):
	161	+ return self.l1_tile_shape[2]
	162	+
	163	+ @property
	164	+ def l0_m(self):
	165	+ return self.l0_tile_shape[0]
	166	+
	167	+ @property
	168	+ def l0_n(self):
	169	+ return self.l0_tile_shape[1]
	170	+
	171	+ @property
	172	+ def l0_k(self):
	173	+ return self.l0_tile_shape[2]
	174	+
	175	+ def set_l1_tile(self, new_l1_tile):
	176	+ self.l1_tile_shape = list(new_l1_tile)
	177	+ self.l0_tile_shape[0] = self.l1_m
	178	+ self.l0_tile_shape[1] = self.l1_n
	179	+ # the new l1k may be less than l0k
	180	+ if self.l0_k > self.l1_k:
	181	+ self.l0_tile_shape[2] = self.l1_k
	182	+
	183	+ def procedural_name(self):
	184	+ return "l1_{l1m}x{l1n}x{l1k}_l0_{l0m}x{l0n}x{l0k}".format(
	185	+ l1m=self.l1_tile_shape[0],
	186	+ l1n=self.l1_tile_shape[1],
	187	+ l1k=self.l1_tile_shape[2],
	188	+ l0m=self.l0_tile_shape[0],
	189	+ l0n=self.l0_tile_shape[1],
	190	+ l0k=self.l0_tile_shape[2],
	191	+ )
	192	+
	193	+ def l1_tile_typename(self, is_tla=False):
	194	+ if is_tla:
	195	+ tile_fmt = "tla::Shape<Int<{l1m}>, Int<{l1n}>, Int<{l1k}>>"
	196	+ else:
	197	+ tile_fmt = "GemmShape<{l1m}, {l1n}, {l1k}>"
	198	+ return tile_fmt.format(
	199	+ l1m=self.l1_tile_shape[0],
	200	+ l1n=self.l1_tile_shape[1],
	201	+ l1k=self.l1_tile_shape[2],
	202	+ )
	203	+
	204	+ def l0_tile_typename(self, is_tla=False):
	205	+ if is_tla:
	206	+ tile_fmt = "tla::Shape<Int<{l0m}>, Int<{l0n}>, Int<{l0k}>>"
	207	+ else:
	208	+ tile_fmt = "GemmShape<{l0m}, {l0n}, {l0k}>"
	209	+ return tile_fmt.format(
	210	+ l0m=self.l0_tile_shape[0],
	211	+ l0n=self.l0_tile_shape[1],
	212	+ l0k=self.l0_tile_shape[2],
	213	+ )

Apython/catlass_cppgen/catlass_cppgen/common/__init__.py+12-0

@@ -0,0 +1,12 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from catlass_cppgen.common.data_type import DataType, get_default_accumulator
	11	+from catlass_cppgen.common.typing import SupportedDataType, SupportedTensor, GM_ADDR
	12	+from catlass_cppgen.common.op_tensor import OpTensor

Apython/catlass_cppgen/catlass_cppgen/common/data_type.py+149-0

@@ -0,0 +1,149 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+import warnings
	11	+from enum import Enum
	12	+from functools import lru_cache
	13	+
	14	+from catlass_cppgen.common.typing import SupportedDataType
	15	+import torch
	16	+
	17	+
	18	+class DataType(Enum):
	19	+ """数据类型枚举，仅允许通过 from_dtype 构造"""
	20	+
	21	+ AUTO = "auto"
	22	+ UNDEFINED = "void"
	23	+ FLOAT = "float"
	24	+ FLOAT16 = "half"
	25	+ INT8 = "int8_t"
	26	+ INT32 = "int32_t"
	27	+ UINT8 = "uint8_t"
	28	+ INT16 = "int16_t"
	29	+ UINT16 = "uint16_t"
	30	+ UINT32 = "uint32_t"
	31	+ INT64 = "int64_t"
	32	+ UINT64 = "uint64_t"
	33	+ DOUBLE = "double"
	34	+ BOOL = "bool"
	35	+ STRING = "string"
	36	+ COMPLEX64 = "complex64"
	37	+ COMPLEX128 = "complex128"
	38	+ BF16 = "bfloat16_t"
	39	+ INT4 = "AscendC::int4_t"
	40	+ UINT1 = "uint1"
	41	+ COMPLEX32 = "complex32"
	42	+ HIFLOAT8 = "hi_float8"
	43	+ FLOAT8_E5M2 = "float8_e5m2_t"
	44	+ FLOAT8_E4M3FN = "float8_e4m3_t"
	45	+ FLOAT8_E8M0 = "float8_e8m0_t"
	46	+ FLOAT6_E3M2 = "float6_e3m2"
	47	+ FLOAT6_E2M3 = "float6_e2m3"
	48	+ FLOAT4_E2M1 = "float4_e2m1x2_t"
	49	+ FLOAT4_E1M2 = "float4_e1m2x2_t"
	50	+
	51	+ @classmethod
	52	+ @lru_cache(maxsize=10)
	C C Chen_HaoWen7月10日 [代码质量 / 前向兼容性] Python 3.12 起，在 `@classmethod` 上叠加 `@lru_cache` 已被标记为 `DeprecationWarning`，并在 Python 3.14 中计划移除该用法。当前写法虽然在 3.8–3.11 能正常工作，但存在前向兼容风险。建议方案：改为在 `@staticmethod` 上使用 `@lru_cache`，或使用 `functools.cache` 包装一个独立的纯函数（接收 `raw_dtype` 参数），`from_dtype` 内部调用该缓存函数。 I init__zhb__ 7月10日评论： 7月10日评论：
	53	+ def from_dtype(cls, raw_dtype: SupportedDataType) -> "DataType":
	54	+ """仅通过from_dtype接口进行构造, 其余转换接口全部移除"""
	55	+ # 定义映射（不保留为类字段，简化为这里局部）
	56	+ torch_map = {
	57	+ torch.float32: cls.FLOAT,
	58	+ torch.float: cls.FLOAT,
	59	+ torch.float16: cls.FLOAT16,
	60	+ torch.half: cls.FLOAT16,
	61	+ torch.int8: cls.INT8,
	62	+ torch.int32: cls.INT32,
	63	+ torch.int: cls.INT32,
	64	+ torch.uint8: cls.UINT8,
	65	+ torch.int16: cls.INT16,
	66	+ torch.short: cls.INT16,
	67	+ torch.int64: cls.INT64,
	68	+ torch.long: cls.INT64,
	69	+ torch.float64: cls.DOUBLE,
	70	+ torch.double: cls.DOUBLE,
	71	+ torch.bool: cls.BOOL,
	72	+ torch.complex64: cls.COMPLEX64,
	73	+ torch.complex128: cls.COMPLEX128,
	74	+ }
	75	+
	76	+ torch_map_optional = {
	77	+ "bfloat16": cls.BF16,
	78	+ "float8_e5m2": cls.FLOAT8_E5M2,
	79	+ "float8_e4m3fn": cls.FLOAT8_E4M3FN,
	80	+ "float8_e8m0fnu": cls.FLOAT8_E8M0,
	81	+ "float4_e2m1fn_x2": cls.FLOAT4_E2M1,
	82	+ "float4_e1m2fn_x2": cls.FLOAT4_E1M2,
	83	+ }
	84	+ for dtype_name, dtype in torch_map_optional.items():
	85	+ if hasattr(torch, dtype_name):
	86	+ torch_map[getattr(torch, dtype_name)] = dtype
	87	+
	88	+ # direct torch match
	89	+ if isinstance(raw_dtype, torch.dtype):
	90	+ if raw_dtype in torch_map:
	91	+ return torch_map[raw_dtype]
	92	+ if str(raw_dtype) == "torch.float8_e5m2":
	93	+ return cls.FLOAT8_E5M2
	94	+ if str(raw_dtype) == "torch.float8_e4m3fn":
	95	+ return cls.FLOAT8_E4M3FN
	96	+ # 如果传入的是 torch.dtype 但不在映射中，返回 UNDEFINED
	97	+ return cls.UNDEFINED
	98	+
	99	+ # 如果传入的是其他类型（非 torch.dtype），返回 UNDEFINED
	100	+ return cls.UNDEFINED
	101	+
	102	+ def data_size(self) -> int:
	103	+ """获取数据类型的字节大小
	104	+
	105	+ Returns:
	106	+ int: 数据类型的字节大小
	107	+
	108	+ Raises:
	109	+ ValueError: 如果数据类型的大小未定义
	110	+ """
	111	+ size_map = {
	C C Chen_HaoWen7月10日 [正确性 / 健壮性] `data_size()` 的 `size_map` 仅覆盖了 6 种数据类型（FLOAT, FLOAT16, BF16, INT8, INT32, INT64），而 `DataType` 枚举定义了 25+ 种类型。缺失的类型包括： DOUBLE (8 bytes), UINT8 (1), INT16 (2), UINT16 (2), UINT32 (4), UINT64 (8), BOOL (1) INT4 (1), 以及所有 FLOAT8 / FLOAT4 / FLOAT6 类型当用户使用这些合法的 `DataType` 调用 `data_size()` 时会抛出 `ValueError`。建议补全所有枚举类型对应的字节大小，或在枚举定义中标注哪些类型不支持 `data_size()`。 I init__zhb__ 7月10日评论： 7月10日评论：
	112	+ self.FLOAT: 4,
	113	+ self.FLOAT16: 2,
	114	+ self.BF16: 2,
	115	+ self.INT8: 1,
	116	+ self.INT32: 4,
	117	+ self.INT64: 8,
	118	+ }
	119	+ if self not in size_map:
	120	+ raise ValueError(f"Data size not defined for {self}")
	121	+ return size_map[self]
	122	+
	123	+
	124	+def get_default_accumulator(data_type_A: DataType, data_type_B: DataType) -> DataType:
	125	+ """获取默认的累加器数据类型"""
	126	+ if data_type_A == DataType.UNDEFINED or data_type_B == DataType.UNDEFINED:
	127	+ raise ValueError("accumulator dtype cannot be derived when A or B is DataType.UNDEFINED")
	128	+ if data_type_A == DataType.AUTO and data_type_B == DataType.AUTO:
	129	+ raise ValueError("accumulator dtype cannot be derived when A and B are both DataType.AUTO")
	130	+
	131	+ if data_type_A == DataType.AUTO:
	132	+ warnings.warn("The dtype of A is auto-derived from B since A is DataType.AUTO", UserWarning, stacklevel=2)
	133	+ data_type_A = data_type_B
	134	+ if data_type_B == DataType.AUTO:
	135	+ warnings.warn("The dtype of B is auto-derived from A since B is DataType.AUTO", UserWarning, stacklevel=2)
	136	+ data_type_B = data_type_A
	137	+
	138	+ if data_type_A != data_type_B:
	139	+ raise ValueError(f"Accumulator type cannot be derived when the dtype of A and B are not the same")
	140	+
	141	+ accumulator_map = {
	142	+ (DataType.FLOAT16, DataType.FLOAT16): DataType.FLOAT,
	143	+ (DataType.FLOAT, DataType.FLOAT16): DataType.FLOAT,
	144	+ (DataType.BF16, DataType.BF16): DataType.FLOAT,
	145	+ (DataType.INT8, DataType.INT8): DataType.INT32,
	146	+ (DataType.INT4, DataType.INT4): DataType.INT32,
	147	+ }
	148	+
	149	+ return accumulator_map.get((data_type_A, data_type_B), data_type_A)

Apython/catlass_cppgen/catlass_cppgen/common/op_tensor.py+109-0

@@ -0,0 +1,109 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+import ctypes
	11	+from typing import Optional
	12	+import torch
	13	+import numpy as np
	14	+
	15	+from catlass_cppgen.common.data_type import DataType
	16	+from catlass_cppgen.catlass.layout.layout import Layout
	17	+from catlass_cppgen.common.typing import SupportedTensor
	18	+from catlass_cppgen.common.utils import infer_layout_from_stride, get_tensor_data_ptr
	19	+
	20	+
	21	+class OpTensor:
	22	+ """OpTensor类是op的输入输出tensor的抽象，用于表示op的输入输出tensor
	23	+
	24	+ 该类提供了统一的tensor抽象，可以从torch.Tensor或np.ndarray创建，
	25	+ 或直接通过shape和stride创建（避免实例化），并自动推断数据类型和布局信息。
	26	+
	27	+ 参数:
	28	+ dtype: DataType, tensor的数据类型
	29	+ layout: Layout, tensor的布局
	30	+ shape: tuple[int, ...], tensor的完整形状
	31	+ data_ptr: Optional[ctypes.c_void_p], tensor的数据指针
	32	+ """
	33	+ def __init__(self, dtype: DataType, layout: Layout, shape: Optional[tuple[int, ...]] = None, data_ptr: Optional[ctypes.c_void_p] = None):
	34	+ self.dtype = dtype
	35	+ self.layout = layout
	36	+ self._shape = shape if shape is not None else layout.shape
	37	+ self.data_ptr = data_ptr
	38	+
	39	+ @classmethod
	40	+ def from_shape_stride(
	41	+ cls,
	42	+ shape: tuple[int, ...],
	43	+ stride: tuple[int, ...],
	44	+ dtype: DataType,
	45	+ ) -> "OpTensor":
	46	+ """直接从 shape 和 stride 创建 OpTensor（避免实例化 tensor）
	47	+
	48	+ 参数:
	49	+ shape: tensor 的形状（可以是 2D 或 3D，支持 batched）
	50	+ stride: tensor 的步长
	51	+ dtype: tensor 的数据类型
	52	+
	53	+ 返回:
	54	+ OpTensor 对象
	55	+ """
	56	+ # 从 shape 和 stride 推断 Layout（对于 batched，只推断内层矩阵的布局）
	57	+ layout = infer_layout_from_stride(shape, stride)
	58	+ # 保存完整的 shape（包括 batch 维度）
	59	+ return cls(dtype=dtype, layout=layout, shape=shape, data_ptr=None)
	60	+
	61	+ @classmethod
	62	+ def from_tensor(
	63	+ cls,
	64	+ tensor: SupportedTensor,
	65	+ layout: Optional[Layout] = None,
	66	+ dtype: Optional[DataType] = None
	67	+ ) -> "OpTensor":
	68	+ """从 torch.Tensor 或 np.ndarray 创建 OpTensor
	69	+
	70	+ 参数:
	71	+ tensor: torch.Tensor 或 np.ndarray
	72	+ layout: 可选的 Layout，如果不提供则从 tensor 的 stride 推断
	73	+ dtype: 可选的 DataType，如果不提供则从 tensor 的 dtype 推断
	74	+
	75	+ 返回:
	76	+ OpTensor 对象
	77	+ """
	78	+ # 获取 shape 和 stride
	79	+ if isinstance(tensor, torch.Tensor):
	80	+ shape = tuple(tensor.shape)
	81	+ stride = tuple(tensor.stride())
	82	+ raw_dtype = tensor.dtype
	83	+ elif isinstance(tensor, np.ndarray):
	84	+ shape = tuple(tensor.shape)
	85	+ stride = tuple(tensor.strides)
	86	+ if tensor.dtype.itemsize > 0:
	87	+ stride = tuple(s // tensor.dtype.itemsize for s in stride)
	88	+ raw_dtype = tensor.dtype
	89	+ else:
	90	+ raise TypeError(f"Unsupported tensor type: {type(tensor)}")
	91	+ if dtype is None:
	92	+ dtype = DataType.from_dtype(raw_dtype)
	93	+ if layout is None:
	94	+ layout = infer_layout_from_stride(shape, stride)
	95	+ data_ptr = get_tensor_data_ptr(tensor)
	96	+ return cls(dtype=dtype, layout=layout, shape=shape, data_ptr=data_ptr)
	97	+
	98	+ @property
	99	+ def shape(self) -> tuple[int, ...]:
	100	+ """返回 tensor 的完整形状"""
	101	+ return self._shape
	102	+
	103	+ @property
	104	+ def stride(self) -> tuple[int, ...]:
	105	+ return self.layout.stride
	106	+
	107	+ @property
	108	+ def capacity(self) -> int:
	109	+ return self.layout.capacity

Apython/catlass_cppgen/catlass_cppgen/common/typing.py+31-0

@@ -0,0 +1,31 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from __future__ import annotations
	11	+
	12	+import ctypes
	13	+from typing import Union
	14	+import torch
	15	+import numpy as np
	16	+
	17	+"""
	18	+类型预留
	19	+"""
	20	+
	21	+# 使用字符串形式避免循环导入，from __future__ import annotations 会自动处理
	22	+SupportedTensor = Union[torch.Tensor, np.ndarray, "OpTensor"] # type: ignore[name-defined]
	23	+SupportedDataType = Union[torch.dtype, np.dtype]
	CANN-robot7月6日此条代码评论区间+21至+23 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：mypy，请Committer检视其合理性。
	24	+
	25	+
	26	+class GM_ADDR(ctypes.c_void_p):
	27	+ pass
	28	+
	29	+# 未实现
	30	+class EpilogueParams(ctypes.c_void_p):
	31	+ pass

Apython/catlass_cppgen/catlass_cppgen/common/utils.py+149-0

@@ -0,0 +1,149 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+import ctypes
	11	+from typing import Optional
	12	+import torch
	13	+import numpy as np
	14	+
	15	+from catlass_cppgen.common.data_type import DataType
	16	+from catlass_cppgen.catlass.layout.layout import Layout, RowMajor, ColumnMajor, VectorLayout
	17	+from catlass_cppgen.common.typing import SupportedTensor
	18	+
	19	+
	20	+def infer_layout_from_stride(shape: tuple[int, ...], stride: tuple[int, ...]) -> Layout:
	21	+ """从 shape 和 stride 推断 Layout 类型
	22	+
	23	+ 参数:
	24	+ shape: tensor 的形状
	25	+ stride: tensor 的步长
	26	+
	27	+ 返回:
	28	+ Layout 对象
	29	+ """
	30	+ if len(shape) == 1:
	31	+ # 一维向量
	32	+ return VectorLayout(shape[0])
	33	+ if len(shape) == 2:
	34	+ # 二维输入，目前只支持RowMajor和ColumnMajor
	35	+ m, n = shape
	36	+ stride_m, stride_n = stride
	37	+ if stride_m == n and stride_n == 1:
	38	+ return RowMajor((m, n))
	39	+ if stride_m == 1 and stride_n == m:
	40	+ return ColumnMajor((m, n))
	41	+ # 非标准布局，默认使用 RowMajor，但保留原始 stride
	42	+ return RowMajor((m, n))
	43	+ else:
	44	+ # 三维输入
	45	+ if len(shape) == 3:
	46	+ batch, m, n = shape
	47	+ stride_batch, stride_m, stride_n = stride
	48	+ if stride_m == n and stride_n == 1:
	49	+ return RowMajor((m, n))
	50	+ if stride_m == 1 and stride_n == m:
	51	+ return ColumnMajor((m, n))
	52	+
	53	+ return RowMajor((m, n))
	54	+ else:
	55	+ return RowMajor((shape[-2], shape[-1]))
	56	+
	57	+
	58	+def get_tensor_data_ptr(tensor: SupportedTensor) -> Optional[ctypes.c_void_p]:
	59	+ """从 torch.Tensor 或 np.ndarray 获取数据指针
	60	+
	61	+ 参数:
	62	+ tensor: torch.Tensor 或 np.ndarray
	63	+
	64	+ 返回:
	65	+ ctypes.c_void_p 数据指针，如果无法获取则返回 None
	66	+ """
	67	+ if isinstance(tensor, torch.Tensor):
	68	+ return ctypes.c_void_p(tensor.data_ptr())
	69	+ elif isinstance(tensor, np.ndarray):
	70	+ return ctypes.c_void_p(tensor.ctypes.data)
	71	+ return None
	72	+
	73	+def extract_info(tensor, default_element, default_layout):
	74	+ """从 OpTensor 或实际 tensor 中提取信息
	75	+
	76	+ 参数:
	77	+ tensor: OpTensor、torch.Tensor、np.ndarray 或 None
	78	+ default_element: 默认的数据类型
	79	+ default_layout: 默认的布局类型
	80	+
	81	+ 返回:
	82	+ tuple: (shape, element, layout, tensor_obj)
	83	+ - shape: tensor 的形状
	84	+ - element: 数据类型
	85	+ - layout: 布局类型
	86	+ - tensor_obj: tensor 对象（如果是 OpTensor 则为 None，否则为原 tensor）
	87	+ """
	88	+ # 延迟导入以避免循环导入
	89	+ from catlass_cppgen.common.op_tensor import OpTensor
	90	+
	91	+ if tensor is None:
	92	+ return None, None, None, None
	93	+ if isinstance(tensor, OpTensor):
	94	+ # 使用 OpTensor 的信息，避免实例化
	95	+ shape = tensor.shape
	96	+ element = tensor.dtype
	97	+ layout = tensor.layout
	98	+ return shape, element, layout, None # 不传递 tensor 对象
	99	+ else:
	100	+ # 从实际 tensor 中提取信息（向后兼容）
	101	+ shape = tuple(tensor.shape)
	102	+ if isinstance(tensor, torch.Tensor):
	103	+ element = DataType.from_dtype(tensor.dtype)
	104	+ elif isinstance(tensor, np.ndarray):
	105	+ element = DataType.from_dtype(tensor.dtype)
	106	+ else:
	107	+ element = default_element
	108	+ # 对于实际 tensor，layout 使用默认值
	109	+ layout = default_layout
	110	+ return shape, element, layout, tensor # 传递 tensor 对象以保持兼容
	111	+
	112	+def _get_cpp_value(value: any) -> str:
	113	+ """将 Python 值转换为 C++ 代码字符串."""
	114	+ if isinstance(value, bool):
	115	+ return "true" if value else "false"
	116	+ return str(value)
	117	+
	118	+
	119	+def _snake_to_camel(snake_str: str) -> str:
	120	+ """将蛇形命名转换为驼峰命名."""
	121	+ components = snake_str.split('_')
	122	+ return components[0] + ''.join(word.capitalize() for word in components[1:])
	123	+
	124	+
	125	+def _get_cpp_type(value: any) -> str:
	126	+ """根据 Python 值推断对应的 C++ 类型."""
	127	+ if isinstance(value, bool):
	128	+ return "bool"
	129	+ elif isinstance(value, int):
	130	+ return "uint32_t"
	131	+ else:
	132	+ return "auto"
	133	+
	134	+
	135	+def get_type_name(type_str) -> str:
	136	+ """获取类型的名称字符串.
	137	+
	138	+ 对于 DataType 枚举，返回其 value；对于字符串，直接返回；对于其他类型，返回其 __name__.
	139	+
	140	+ :param type_str: 类型对象，可以是 DataType 枚举、字符串或其他类型.
	141	+ :return: 类型的名称字符串.
	142	+ :rtype: str
	143	+ """
	144	+ if isinstance(type_str, DataType):
	145	+ return type_str.value
	146	+ elif isinstance(type_str, str):
	147	+ return type_str
	148	+ else:
	149	+ return type_str.__name__

Apython/catlass_cppgen/catlass_cppgen/kernel/__init__.py+11-0

@@ -0,0 +1,11 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from catlass_cppgen.kernel.kernel_base import KernelBase
	11	+from catlass_cppgen.kernel.visitor_kernel_base import VisitorKernelBase

Apython/catlass_cppgen/catlass_cppgen/kernel/gemm/__init__.py+26-0

@@ -0,0 +1,26 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from catlass_cppgen.kernel.gemm.gemm_base import GemmKernelBase
	11	+from catlass_cppgen.kernel.gemm.basic_matmul import BasicMatmulKernel
	12	+from catlass_cppgen.kernel.gemm.batched_matmul import BatchedMatmulKernel
	13	+from catlass_cppgen.kernel.gemm.basic_matmul_tla_visitor import BasicMatmulTlaVisitorKernel
	14	+from catlass_cppgen.kernel.gemm.multi_core_splitk_matmul import MultiCoreSplitkMatmulKernel
	15	+from catlass_cppgen.kernel.gemm.streamk_matmul import StreamkMatmulKernel
	16	+from catlass_cppgen.kernel.gemm.tail_multi_core_splitk_matmul import TailMultiCoreSplitkMatmulKernel
	17	+
	18	+__all__ = [
	19	+ "GemmKernelBase",
	20	+ "BasicMatmulKernel",
	21	+ "BatchedMatmulKernel",
	22	+ "BasicMatmulTlaVisitorKernel",
	23	+ "MultiCoreSplitkMatmulKernel",
	24	+ "StreamkMatmulKernel",
	25	+ "TailMultiCoreSplitkMatmulKernel",
	26	+]

Apython/catlass_cppgen/catlass_cppgen/kernel/gemm/basic_matmul.py+126-0

@@ -0,0 +1,126 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from typing import Any, Dict, List, Tuple
	11	+from catlass_cppgen.kernel.gemm.gemm_base import GemmKernelBase
	12	+from catlass_cppgen.catlass.gemm_coord import GemmCoord, GemmShape
	13	+from catlass_cppgen.catlass.layout.layout import Layout
	14	+from catlass_cppgen.common.typing import GM_ADDR
	15	+from catlass_cppgen.catlass.gemm.dispatch_policy import (
	16	+ MmadPingpong,
	17	+)
	18	+from catlass_cppgen.catlass.arch.arch import Arch
	19	+
	20	+
	21	+class BasicMatmulKernel(GemmKernelBase):
	22	+ _KERNEL_NAME_BASE = "BasicMatmulTla"
	23	+ _FEATURES = {"is_support_evg": True, "is_support_relu": True, "slice_axis": None, "is_mix": False}
	24	+
	25	+ def __init__(self, **kwargs):
	26	+ """初始化 BasicMatmulKernel.
	27	+
	28	+ relu_enable 现在通过 tune() 方法设置，默认为 False。
	29	+ """
	30	+ super().__init__(**kwargs)
	31	+
	32	+ _INCLUDES = [
	33	+ "catlass/catlass.hpp",
	34	+ "catlass/arch/arch.hpp",
	35	+ "catlass/layout/layout.hpp",
	36	+ "catlass/status.hpp",
	37	+
	38	+ "catlass/gemm/block/block_mmad.hpp",
	39	+ "catlass/gemm/block/block_swizzle.hpp",
	40	+ "catlass/gemm/dispatch_policy.hpp",
	41	+ "catlass/gemm/gemm_type.hpp",
	42	+ "catlass/gemm/device/device_gemm.hpp",
	43	+ "catlass/gemm_coord.hpp",
	44	+ "catlass/matrix_coord.hpp",
	45	+ "tla/layout.hpp",
	46	+
	47	+ "catlass/gemm/kernel/basic_matmul_tla.hpp",
	48	+ ]
	49	+ _KERNEL_NAME = "{arch_name}_{kernel_name}_{dispatch_policy_name}_{swizzle_name}_{l1_tile_shape_str}_{l0_tile_shape_str}"
	50	+ _PARAMS_DEVICE = [
	51	+ (GemmCoord, "problemShape"),
	52	+ (GM_ADDR, "deviceA"),
	53	+ (Layout, "layoutA"),
	54	+ (GM_ADDR, "deviceB"),
	55	+ (Layout, "layoutB"),
	56	+ (GM_ADDR, "deviceC"),
	57	+ (Layout, "layoutC"),
	58	+ (GM_ADDR, "deviceBias"),
	59	+ ]
	60	+ _DISPATCH_POLICY = """\
	61	+ using ArchTag = {arch_tag};
	62	+{constexpr_declarations}
	63	+ using DispatchPolicy = {dispatch_policy_template};
	64	+"""
	65	+ _KERNEL_TEMPLATE = """\
	66	+ using L1TileShape = {l1_tile_shape_tla};
	67	+ using L0TileShape = {l0_tile_shape_tla};
	68	+
	69	+ using ElementA = {element_A};
	70	+ using ElementB = {element_B};
	71	+ using ElementC = {element_C};
	72	+ using LayoutTagA = {layout_A};
	73	+ using LayoutTagB = {layout_B};
	74	+ using LayoutTagC = layout::RowMajor;
	75	+ using ElementBias = {element_Bias};
	76	+
	77	+ using TileCopy = Gemm::Tile::PackedTileCopyTla<ArchTag, ElementA, LayoutTagA, ElementB, LayoutTagB, ElementC, LayoutTagC, ElementBias, {relu_enable}>;
	78	+ using BlockMmad = Gemm::Block::BlockMmadTla<DispatchPolicy, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, ElementBias, TileCopy>;
	79	+ using BlockEpilogue = void;
	80	+
	81	+ using BlockScheduler = typename Gemm::Block::GemmIdentityBlockSwizzle<3, 0>;
	82	+ using GemmKernel = Gemm::Kernel::BasicMatmulTla<BlockMmad, BlockEpilogue, BlockScheduler>;
	83	+"""
	84	+ _INPUT_TEMPLATE = """\
	85	+ uint32_t m = M;
	86	+ uint32_t k = K;
	87	+ uint32_t n = N;
	88	+"""
	89	+ _LAYOUT_TEMPLATE = """\
	90	+ GemmCoord problemShape{{m, n, k}};
	91	+ // Define the layout of each matrix
	92	+ LayoutTagA tagA{{m, k}};
	93	+ LayoutTagB tagB{{k, n}};
	94	+ LayoutTagC tagC{{m, n}};
	95	+ auto layoutA = tla::MakeLayoutFromTag(tagA);
	96	+ auto layoutB = tla::MakeLayoutFromTag(tagB);
	97	+ auto layoutC = tla::MakeLayoutFromTag(tagC);
	98	+"""
	99	+
	100	+ def get_default_tile_shape(self) -> Tuple[GemmShape, GemmShape]:
	101	+ element_max_size = max(
	102	+ self.element_A.data_size(),
	103	+ self.element_B.data_size(),
	104	+ self.element_C.data_size(),
	105	+ )
	106	+ if self.arch_tag == Arch.AtlasA2:
	107	+ l1_m, l1_n, l1_k, l0_k = 128, 256, 512//element_max_size, 128//element_max_size
	108	+ elif self.arch_tag == Arch.Ascend950:
	109	+ l1_m, l1_n, l1_k, l0_k = 256, 256, 512//element_max_size, 128//element_max_size
	110	+ if self.element_Bias is not None and self.element_Bias != "void":
	111	+ l1_m -= 16
	112	+ return GemmShape(l1_m, l1_n, l1_k), GemmShape(l1_m, l1_n, l0_k)
	113	+
	114	+ def get_default_dispatch_policy_list(self) -> List[MmadPingpong]:
	115	+ """获取 BasicMatmulKernel 的默认 dispatch_policy 列表.
	116	+
	117	+ :return: 包含默认 dispatch_policy 的列表，列表的第一个元素 [0] 是默认策略.
	118	+ :rtype: List[MmadPingpong]
	119	+ """
	120	+ return [MmadPingpong(arch_tag=self.arch_tag, enable_unit_flag=True)]
	121	+
	122	+ def get_render_params(self, use_constexpr: bool = True) -> Dict[str, Any]:
	123	+ """获取渲染参数，包括 kernel 名称格式化参数."""
	124	+ params = super().get_render_params(use_constexpr)
	125	+ params['relu_enable'] = 'true' if self.relu_enable else 'false'
	126	+ return self._add_kernel_name_params(params, self._KERNEL_NAME_BASE)

Apython/catlass_cppgen/catlass_cppgen/kernel/gemm/basic_matmul_tla_visitor.py+129-0

@@ -0,0 +1,129 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from typing import Any, Dict, List, Tuple
	11	+from catlass_cppgen.kernel.gemm.gemm_base import GemmKernelBase
	12	+from catlass_cppgen.kernel.visitor_kernel_base import VisitorKernelBase
	13	+from catlass_cppgen.catlass.gemm_coord import GemmCoord, GemmShape
	14	+from catlass_cppgen.catlass.layout.layout import Layout
	15	+from catlass_cppgen.common.typing import GM_ADDR
	16	+from catlass_cppgen.catlass.gemm.dispatch_policy import (
	17	+ MmadPingpong,
	18	+)
	19	+from catlass_cppgen.catlass.arch.arch import Arch
	20	+
	21	+
	22	+class BasicMatmulTlaVisitorKernel(GemmKernelBase, VisitorKernelBase):
	23	+ _KERNEL_NAME_BASE = "BasicMatmulTlaVisitor"
	24	+ _FEATURES = {"is_support_evg": True, "is_support_relu": False, "slice_axis": None, "is_mix": True}
	25	+
	26	+ _INCLUDES = [
	27	+ "catlass/catlass.hpp",
	28	+ "catlass/arch/arch.hpp",
	29	+ "catlass/layout/layout.hpp",
	30	+ "catlass/status.hpp",
	31	+
	32	+ "catlass/gemm/block/block_mmad.hpp",
	33	+ "catlass/gemm/block/block_swizzle.hpp",
	34	+ "catlass/gemm/dispatch_policy.hpp",
	35	+ "catlass/gemm/gemm_type.hpp",
	36	+ "catlass/gemm/device/device_gemm.hpp",
	37	+ "catlass/gemm_coord.hpp",
	38	+ "catlass/matrix_coord.hpp",
	39	+
	40	+ "tla/layout.hpp",
	41	+
	42	+ "catlass/epilogue/block/block_epilogue.hpp",
	43	+ "catlass/epilogue/fusion/fusion.hpp",
	44	+
	45	+ "catlass/gemm/kernel/basic_matmul_tla_visitor.hpp",
	46	+ ]
	47	+ _KERNEL_NAME = "{arch_name}_{kernel_name}_{dispatch_policy_name}_{swizzle_name}_{l1_tile_shape_str}_{l0_tile_shape_str}"
	48	+ _PARAMS_DEVICE = [
	49	+ (GemmCoord, "problemShape"),
	50	+ (GM_ADDR, "deviceA"),
	51	+ (Layout, "layoutA"),
	52	+ (GM_ADDR, "deviceB"),
	53	+ (Layout, "layoutB"),
	54	+ (GM_ADDR, "nullptr"),
	55	+ (Layout, "{}"), # layoutC 使用空初始化 {}
	56	+ (GM_ADDR, "nullptr"),
	57	+ ("typename EVG::Arguments", "evg_args"), # evg_args 参数
	58	+ ]
	59	+ _DISPATCH_POLICY = """\
	60	+ using ArchTag = {arch_tag};
	61	+{constexpr_declarations}
	62	+ using DispatchPolicy = {dispatch_policy_template};
	63	+"""
	64	+ _KERNEL_TEMPLATE = """\
	65	+ using L1TileShape = {l1_tile_shape_tla};
	66	+ using L0TileShape = {l0_tile_shape_tla};
	67	+
	68	+ using ElementA = {element_A};
	69	+ using ElementB = {element_B};
	70	+ using ElementC = {element_C};
	71	+ using LayoutTagA = {layout_A};
	72	+ using LayoutTagB = {layout_B};
	73	+ using LayoutTagC = layout::RowMajor;
	74	+
	75	+ using TileCopy = Gemm::Tile::PackedTileCopyTla<ArchTag, ElementA, LayoutTagA, ElementB, LayoutTagB, ElementC, LayoutTagC>;
	76	+ using BlockMmad = Gemm::Block::BlockMmadTla<DispatchPolicy, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, void, TileCopy>;
	77	+ using BlockEpilogue = {epilogue_str};
	78	+
	79	+ using BlockScheduler = typename Gemm::Block::GemmIdentityBlockSwizzle<3, 0>;
	80	+ using GemmKernel = Gemm::Kernel::BasicMatmulTlaVisitor<BlockMmad, BlockEpilogue, BlockScheduler>;
	81	+"""
	82	+ _INPUT_TEMPLATE = """\
	83	+ uint32_t m = M;
	84	+ uint32_t k = K;
	85	+ uint32_t n = N;
	86	+"""
	87	+ _LAYOUT_TEMPLATE = """\
	88	+ GemmCoord problemShape{{m, n, k}};
	89	+ // Define the layout of each matrix
	90	+ LayoutTagA tagA{{m, k}};
	91	+ LayoutTagB tagB{{k, n}};
	92	+ LayoutTagC tagC{{m, n}};
	93	+ auto layoutA = tla::MakeLayoutFromTag(tagA);
	94	+ auto layoutB = tla::MakeLayoutFromTag(tagB);
	95	+ auto layoutC = tla::MakeLayoutFromTag(tagC);
	96	+"""
	97	+ _EVG_TEMPLATE = """\
	98	+using EpilogueDispatchPolicy = Epilogue::EpilogueVisitor<false>;
	99	+{evg_str}
	100	+{evg_args}
	101	+"""
	102	+
	103	+ def get_default_tile_shape(self) -> Tuple[GemmShape, GemmShape]:
	104	+ element_max_size = max(
	105	+ self.element_A.data_size(),
	106	+ self.element_B.data_size(),
	107	+ self.element_C.data_size(),
	108	+ )
	109	+ return GemmShape(256, 256, 128), GemmShape(256, 256, 32)
	110	+
	111	+ def get_default_dispatch_policy_list(self) -> List:
	112	+ """获取 BasicMatmulKernel 的默认 dispatch_policy 列表.
	113	+
	114	+ :return: 包含默认 dispatch_policy 的列表，列表的第一个元素 [0] 是默认策略.
	115	+ :rtype: List
	116	+ """
	117	+ return [MmadPingpong(arch_tag=self.arch_tag, enable_unit_flag=True)]
	118	+
	119	+ def get_render_params(self, use_constexpr: bool = True) -> Dict[str, Any]:
	120	+ """获取渲染参数，包括 kernel 名称格式化参数."""
	121	+ params = super().get_render_params(use_constexpr)
	122	+
	123	+ # 根据 arch_tag 是否为 A5 设置 EpilogueDispatchPolicy 的值
	124	+ if self.arch_tag == Arch.Ascend950:
	125	+ params["epilogue_dispatch_policy_value"] = "true"
	126	+ else:
	127	+ params["epilogue_dispatch_policy_value"] = "false"
	128	+
	129	+ return self._add_kernel_name_params(params, self._KERNEL_NAME_BASE)

Apython/catlass_cppgen/catlass_cppgen/kernel/gemm/batched_matmul.py+148-0

@@ -0,0 +1,148 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from typing import Any, Dict, List, Optional, Tuple
	11	+from catlass_cppgen.kernel.gemm.gemm_base import GemmKernelBase
	12	+from catlass_cppgen.catlass.gemm_coord import GemmCoord, GemmShape
	13	+from catlass_cppgen.catlass.layout.layout import Layout
	14	+from catlass_cppgen.common.data_type import DataType
	15	+from catlass_cppgen.common.typing import GM_ADDR
	16	+from catlass_cppgen.catlass.gemm.dispatch_policy import (
	17	+ MmadPingpong,
	18	+)
	19	+
	20	+
	21	+class BatchedMatmulKernel(GemmKernelBase):
	22	+ slice_axis = None
	23	+ _KERNEL_NAME_BASE = "BatchedMatmulTla"
	24	+ _FEATURES = {"is_support_evg": False, "is_support_relu": False, "slice_axis": None, "is_mix": True}
	25	+
	26	+ _INCLUDES = [
	27	+ "catlass/catlass.hpp",
	28	+ "catlass/arch/arch.hpp",
	29	+ "catlass/layout/layout.hpp",
	30	+ "catlass/status.hpp",
	31	+
	32	+ "catlass/gemm/block/block_mmad.hpp",
	33	+ "catlass/gemm/block/block_swizzle.hpp",
	34	+ "catlass/gemm/dispatch_policy.hpp",
	35	+ "catlass/gemm/gemm_type.hpp",
	36	+ "catlass/gemm/device/device_gemm.hpp",
	37	+ "catlass/gemm_coord.hpp",
	38	+ "catlass/matrix_coord.hpp",
	39	+
	40	+ "tla/layout.hpp",
	41	+
	42	+ "catlass/gemm/kernel/batched_matmul_tla.hpp",
	43	+ ]
	44	+ _KERNEL_NAME = "{arch_name}_{kernel_name}_{dispatch_policy_name}_{swizzle_name}_{l1_tile_shape_str}_{l0_tile_shape_str}"
	45	+ _PARAMS_DEVICE = [
	46	+ (DataType.UINT32, "batchCount"),
	47	+ (GemmCoord, "problemShape"),
	48	+ (GM_ADDR, "deviceA"),
	49	+ (Layout, "layoutA"),
	50	+ (GM_ADDR, "deviceB"),
	51	+ (Layout, "layoutB"),
	52	+ (GM_ADDR, "deviceC"),
	53	+ (Layout, "layoutC"),
	54	+ ]
	55	+ _DISPATCH_POLICY = """\
	56	+ using ArchTag = {arch_tag};
	57	+{constexpr_declarations}
	58	+ using DispatchPolicy = {dispatch_policy_template};
	59	+"""
	60	+ _KERNEL_TEMPLATE = """\
	61	+ using L1TileShape = {l1_tile_shape_tla};
	62	+ using L0TileShape = {l0_tile_shape_tla};
	63	+
	64	+ using ElementA = {element_A};
	65	+ using ElementB = {element_B};
	66	+ using ElementC = {element_C};
	67	+ using LayoutTagA = {layout_A};
	68	+ using LayoutTagB = {layout_B};
	69	+ using LayoutTagC = layout::RowMajor;
	70	+
	71	+ using TileCopy = Gemm::Tile::PackedTileCopyTla<ArchTag, ElementA, LayoutTagA, ElementB, LayoutTagB, ElementC, LayoutTagC>;
	72	+ using BlockMmad = Gemm::Block::BlockMmadTla<DispatchPolicy, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, void, TileCopy>;
	73	+ using BlockEpilogue = void;
	74	+
	75	+ using BlockScheduler = typename Gemm::Block::GemmIdentityBlockSwizzle<3, 0>;
	76	+ using GemmKernel = Gemm::Kernel::BatchedMatmulTla<BlockMmad, BlockEpilogue, BlockScheduler>;
	77	+"""
	78	+ _INPUT_TEMPLATE = """\
	79	+ uint32_t m = M;
	80	+ uint32_t k = K;
	81	+ uint32_t n = N;
	82	+"""
	83	+ _LAYOUT_TEMPLATE = """\
	84	+ uint32_t batchCount = {batchCount};
	85	+ GemmCoord problemShape{{m, n, k}};
	86	+ // Define the layout of each matrix
	87	+ LayoutTagA tagA{{m, k}};
	88	+ LayoutTagB tagB{{k, n}};
	89	+ LayoutTagC tagC{{m, n}};
	90	+ auto layoutA = tla::MakeLayoutFromTag(tagA);
	91	+ auto layoutB = tla::MakeLayoutFromTag(tagB);
	92	+ auto layoutC = tla::MakeLayoutFromTag(tagC);
	93	+"""
	94	+ _ADDITIONAL_DEFINITIONS_TEMPLATE = """\
	95	+ int64_t strideA = m * k;
	96	+ int64_t strideB = k * n;
	97	+ int64_t strideC = m * n;
	98	+"""
	99	+ # 定义在哪些参数后插入什么参数
	100	+ _PARAMS_INSERTIONS = {
	101	+ "layoutA": "strideA",
	102	+ "layoutB": "strideB",
	103	+ "layoutC": "strideC",
	104	+ }
	105	+
	106	+ def __init__(self, batchCount: Optional[int] = None, **kwargs):
	107	+ """初始化 BatchedMatmulKernel.
	108	+
	109	+ :param batchCount: 批处理数量，如果为 None 则使用默认值 1
	110	+ :param kwargs: 传递给父类的其他参数，包括 M, K, N 等
	111	+ """
	112	+ super().__init__(**kwargs)
	113	+ self.batchCount = batchCount if batchCount is not None else 1
	114	+ self.strideA = self.M * self.K
	115	+ self.strideB = self.K * self.N
	116	+ self.strideC = self.M * self.N
	117	+
	118	+ def get_default_tile_shape(self) -> Tuple[GemmShape, GemmShape]:
	119	+ element_max_size = max(
	120	+ self.element_A.data_size(),
	121	+ self.element_B.data_size(),
	122	+ self.element_C.data_size(),
	123	+ )
	124	+ # MmadPingpong 路径与 BasicMatmul 共用 block_mmad_pingpong_tla，tile 约束相同
	125	+ if element_max_size <= 2:
	126	+ l1_m, l1_n, l1_k, l0_k = 256, 256, 256, 64
	127	+ else:
	128	+ l1_m, l1_n, l1_k, l0_k = 256, 256, 128, 32
	129	+ return GemmShape(l1_m, l1_n, l1_k), GemmShape(l1_m, l1_n, l0_k)
	130	+
	131	+ def get_default_dispatch_policy_list(self) -> List:
	132	+ """获取 BatchedMatmulKernel 的默认 dispatch_policy 列表.
	133	+
	134	+ :return: 包含默认 dispatch_policy 的列表，列表的第一个元素 [0] 是默认策略.
	135	+ :rtype: List
	136	+ """
	137	+ return [MmadPingpong(arch_tag=self.arch_tag, enable_unit_flag=True)]
	138	+
	139	+ def get_render_params(self, use_constexpr: bool = True) -> Dict[str, Any]:
	140	+ """获取渲染参数，包括动态生成的 dispatch_policy C++ 代码.
	141	+
	142	+ :param use_constexpr: 当为 True 时，生成包含常量声明的完整代码块；当为 False 时，只生成 using 语句（使用变量名）
	143	+ :return: 渲染参数字典.
	144	+ :rtype: Dict[str, Any]
	145	+ """
	146	+ params = super().get_render_params(use_constexpr)
	147	+ params['batchCount'] = self.batchCount
	148	+ return self._add_kernel_name_params(params, self._KERNEL_NAME_BASE)

Apython/catlass_cppgen/catlass_cppgen/kernel/gemm/gemm_base.py+195-0

@@ -0,0 +1,195 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
	11	+
	12	+from catlass_cppgen.catlass.layout.layout import Layout
	13	+from catlass_cppgen.kernel.kernel_base import KernelBase
	14	+from catlass_cppgen.common.data_type import DataType
	15	+from catlass_cppgen.catlass.arch.arch import Arch
	16	+from catlass_cppgen.catlass.gemm_coord import GemmCoord
	17	+from catlass_cppgen.catlass.gemm.dispatch_policy import (
	18	+ MmadAtlasA2Pingpong,
	19	+ MmadPingpong,
	20	+)
	21	+from catlass_cppgen.catlass.evg_extension import evg as generate_evg
	22	+
	23	+if TYPE_CHECKING:
	24	+ from catlass_cppgen.kernel.gemm.basic_matmul_tla_visitor import BasicMatmulTlaVisitorKernel
	25	+
	26	+
	27	+class GemmKernelBase(KernelBase):
	28	+ def __init__(
	29	+ self,
	30	+ element_accumulator: DataType,
	31	+ element_A: DataType,
	32	+ element_B: DataType,
	33	+ element_C: DataType,
	34	+ element_Bias: DataType,
	35	+ layout_A: Layout,
	36	+ layout_B: Layout,
	37	+ arch_tag: Arch,
	38	+ layout_Bias: Optional[Layout] = None,
	39	+ layout_C: Optional[Layout] = None,
	40	+ M: Optional[int] = None,
	41	+ K: Optional[int] = None,
	42	+ N: Optional[int] = None,
	43	+ evg: Optional[Dict[str, Any]] = None,
	44	+ args, *kwargs
	45	+ ):
	46	+ self.element_A = element_A
	47	+ self.element_B = element_B
	48	+ self.element_Bias = element_Bias
	49	+ self.element_C = element_C
	50	+ self.layout_A = layout_A
	51	+ self.layout_B = layout_B
	52	+ self.layout_Bias = layout_Bias
	53	+ self.layout_C = layout_C
	54	+ self.element_accumulator = element_accumulator
	55	+ self.arch_tag = arch_tag
	56	+ self.M = M
	57	+ self.K = K
	58	+ self.N = N
	59	+ self.evg = evg
	60	+
	61	+ super().__init__(args, *kwargs)
	62	+
	63	+ def get_render_params(self, use_constexpr: bool = True) -> Dict[str, Any]:
	64	+ """获取渲染参数，包括动态生成的 dispatch_policy C++ 代码.
	65	+
	66	+ :param use_constexpr: 当为 True 时，生成包含常量声明的完整代码块；当为 False 时，只生成 using 语句（使用变量名）
	67	+ :return: 渲染参数字典.
	68	+ :rtype: Dict[str, Any]
	69	+ """
	70	+ params = {
	71	+ "arch_tag": self.arch_tag,
	72	+ "l1_tile_shape": self.l1_tile_shape,
	73	+ "l0_tile_shape": self.l0_tile_shape,
	74	+ "l1_tile_shape_tla": self.l1_tile_shape.tla(),
	75	+ "l0_tile_shape_tla": self.l0_tile_shape.tla(),
	76	+ "element_A": self.element_A,
	77	+ "element_B": self.element_B,
	78	+ "element_Bias": self.element_Bias,
	79	+ "element_C": self.element_C,
	80	+ "layout_A": self.layout_A,
	81	+ "layout_B": self.layout_B,
	82	+ "layout_Bias": self.layout_Bias,
	83	+ "layout_C": self.layout_C,
	84	+ "M": self.M,
	85	+ "K": self.K,
	86	+ "N": self.N,
	87	+ "slice_axis": self.slice_axis,
	88	+ }
	89	+ if params.get('element_Bias') is None:
	90	+ params['element_Bias'] = "void"
	91	+
	92	+ if len(self.dispatch_policy) == 0:
	93	+ raise ValueError("dispatch_policy cannot be empty")
	94	+ dispatch_policy = self.dispatch_policy[0]
	95	+ result = dispatch_policy.to_cpp(const_mode=True)
	96	+
	97	+ if self.evg is not None:
	98	+ fn_src = self.evg["fn_src"]
	99	+ example_inputs = self.evg["example_inputs"]
	100	+ callback_name, evg_args, evg_str, arg_renames = generate_evg(
	101	+ fn_src=fn_src,
	102	+ example_inputs=example_inputs,
	103	+ )
	104	+ # 将生成的 EVG 信息添加到字典中，同时保留原先内容
	105	+ self.evg.update({
	106	+ "callback_name": callback_name,
	107	+ "evg_args": evg_args,
	108	+ "evg_str": evg_str,
	109	+ "arg_renames": arg_renames,
	110	+ })
	111	+ params["evg_args"] = evg_args
	112	+ params["evg_str"] = evg_str
	113	+ params["evg_callback_name"] = callback_name # 添加 callback_name 到渲染参数
	114	+ params["epilogue_str"] = f"""Epilogue::Block::BlockEpilogue<
	115	+ EpilogueDispatchPolicy,
	116	+ ArchTag,
	117	+ Int<computeLength>,
	118	+ {callback_name},
	119	+ ElementC
	120	+ >"""
	121	+ else:
	122	+ params["evg_args"] = ""
	123	+ params["evg_str"] = ""
	124	+ params["evg_callback_name"] = ""
	125	+ params["epilogue_str"] = "void"
	126	+
	127	+ if isinstance(result, tuple):
	128	+ const_decls, template_str = result
	129	+ params['constexpr_declarations'] = "\n".join([f" {decl}" for decl in const_decls])
	130	+ params['dispatch_policy_template'] = f"{template_str}"
	131	+ else:
	132	+ params['constexpr_declarations'] = ""
	133	+ params['dispatch_policy_template'] = result
	134	+
	135	+ return params
	136	+
	137	+ def _add_kernel_name_params(self, params: Dict[str, Any], kernel_name_base: str) -> Dict[str, Any]:
	138	+ """添加用于格式化 kernel 名称的参数.
	139	+
	140	+ :param params: 渲染参数字典.
	141	+ :param kernel_name_base: kernel 名称基础（如 "BasicMatmulTla", "GroupedMatmulSliceMTla"）.
	142	+ :return: 添加了格式化参数的参数字典.
	143	+ """
	144	+ params['arch_name'] = self.arch_tag.name
	145	+ params['kernel_name'] = kernel_name_base
	146	+
	147	+ # dispatch_policy 名称只使用类名
	148	+ params['dispatch_policy_name'] = self.dispatch_policy[0].__class__.__name__
	149	+ params['swizzle_name'] = "GemmIdentityBlockSwizzle_3_0"
	150	+ params['l1_tile_shape_str'] = f"{self.l1_tile_shape.m}_{self.l1_tile_shape.n}_{self.l1_tile_shape.k}"
	151	+ params['l0_tile_shape_str'] = f"{self.l0_tile_shape.m}_{self.l0_tile_shape.n}_{self.l0_tile_shape.k}"
	152	+
	153	+ return params
	154	+
	155	+ def to_evg(self, evg_config: Dict[str, Any]) -> Optional['BasicMatmulTlaVisitorKernel']:
	156	+ """将支持 EVG 的 kernel 转换为 EVG 版本.
	157	+
	158	+ :param evg_config: EVG 配置，包含 'fn_src' 和 'example_inputs'
	159	+ :type evg_config: Dict[str, Any]
	160	+ :return: 如果当前 kernel 支持 EVG，返回 BasicMatmulTlaVisitorKernel 实例；否则返回 None
	161	+ :rtype: Optional[BasicMatmulTlaVisitorKernel]
	162	+ """
	163	+ # 检查是否支持 EVG 特性
	164	+ if not self._features.get("is_support_evg", False):
	165	+ return None
	166	+
	167	+ # 延迟导入以避免循环导入
	168	+ from catlass_cppgen.kernel.gemm.basic_matmul_tla_visitor import BasicMatmulTlaVisitorKernel
	169	+
	170	+ # 创建 BasicMatmulTlaVisitorKernel 实例，传递相同的参数和 EVG 配置
	171	+ evg_kernel = BasicMatmulTlaVisitorKernel(
	172	+ element_accumulator=self.element_accumulator,
	173	+ element_A=self.element_A,
	174	+ element_B=self.element_B,
	175	+ element_C=self.element_C,
	176	+ element_Bias=self.element_Bias,
	177	+ layout_A=self.layout_A,
	178	+ layout_B=self.layout_B,
	179	+ layout_Bias=self.layout_Bias,
	180	+ layout_C=self.layout_C,
	181	+ arch_tag=self.arch_tag,
	182	+ M=self.M,
	183	+ K=self.K,
	184	+ N=self.N,
	185	+ evg=evg_config,
	186	+ )
	187	+
	188	+ # 设置 tile shape 和 dispatch_policy
	189	+ evg_kernel.tune(
	190	+ l1_tile_shape=self.l1_tile_shape,
	191	+ l0_tile_shape=self.l0_tile_shape,
	192	+ dispatch_policy=self.dispatch_policy,
	193	+ )
	194	+
	195	+ return evg_kernel

Apython/catlass_cppgen/catlass_cppgen/kernel/gemm/multi_core_splitk_matmul.py+119-0

@@ -0,0 +1,119 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from typing import Any, Dict, List, Optional, Tuple
	11	+from catlass_cppgen.kernel.gemm.gemm_base import GemmKernelBase
	12	+from catlass_cppgen.catlass.gemm_coord import GemmCoord, GemmShape
	13	+from catlass_cppgen.catlass.layout.layout import Layout
	14	+from catlass_cppgen.common.data_type import DataType
	15	+from catlass_cppgen.common.typing import GM_ADDR
	16	+from catlass_cppgen.catlass.gemm.dispatch_policy import (
	17	+ MmadPingpong,
	18	+)
	19	+
	20	+
	21	+class MultiCoreSplitkMatmulKernel(GemmKernelBase):
	22	+ _KERNEL_NAME_BASE = "MultiCoreSplitkMatmulTla"
	23	+ _FEATURES = {"is_support_evg": False, "is_support_relu": False, "slice_axis": "K", "is_mix": True}
	24	+
	25	+ _INCLUDES = [
	26	+ "catlass/catlass.hpp",
	27	+ "catlass/arch/arch.hpp",
	28	+ "catlass/layout/layout.hpp",
	29	+ "catlass/status.hpp",
	30	+
	31	+ "catlass/gemm/block/block_mmad.hpp",
	32	+ "catlass/gemm/block/block_swizzle.hpp",
	33	+ "catlass/gemm/dispatch_policy.hpp",
	34	+ "catlass/gemm/gemm_type.hpp",
	35	+ "catlass/gemm/device/device_gemm.hpp",
	36	+ "catlass/gemm_coord.hpp",
	37	+ "catlass/matrix_coord.hpp",
	38	+
	39	+ "tla/layout.hpp",
	40	+
	41	+ "catlass/gemm/kernel/multi_core_splitk_matmul_tla.hpp",
	42	+ ]
	43	+ _KERNEL_NAME = "{arch_name}_{kernel_name}_{dispatch_policy_name}_{swizzle_name}_{l1_tile_shape_str}_{l0_tile_shape_str}"
	44	+ _PARAMS_DEVICE = [
	45	+ (GemmCoord, "problemShape"),
	46	+ (GM_ADDR, "deviceA"),
	47	+ (Layout, "layoutA"),
	48	+ (GM_ADDR, "deviceB"),
	49	+ (Layout, "layoutB"),
	50	+ (GM_ADDR, "deviceC"),
	51	+ (Layout, "layoutC"),
	52	+ (DataType.AUTO, "aicCoreNum"),
	53	+ (GM_ADDR, "deviceBias"),
	54	+ ]
	55	+ _DISPATCH_POLICY = """\
	56	+ using ArchTag = {arch_tag};
	57	+{constexpr_declarations}
	58	+ using DispatchPolicy = {dispatch_policy_template};
	59	+"""
	60	+ _KERNEL_TEMPLATE = """\
	61	+ using L1TileShape = {l1_tile_shape_tla};
	62	+ using L0TileShape = {l0_tile_shape_tla};
	63	+
	64	+ using ElementA = {element_A};
	65	+ using ElementB = {element_B};
	66	+ using ElementC = {element_C};
	67	+ using ElementBias = {element_Bias};
	68	+ using ElementBiasType = std::conditional_t<std::is_void_v<ElementBias>, uint8_t, ElementBias>;
	69	+ using LayoutTagA = {layout_A};
	70	+ using LayoutTagB = {layout_B};
	71	+ using LayoutTagC = layout::RowMajor;
	72	+
	73	+ using TileCopy = Gemm::Tile::PackedTileCopyTla<ArchTag, ElementA, LayoutTagA, ElementB, LayoutTagB, ElementC, LayoutTagC, ElementBias>;
	74	+ using BlockMmad = Gemm::Block::BlockMmadTla<DispatchPolicy, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, ElementBias, TileCopy>;
	75	+ using BlockEpilogue = void;
	76	+
	77	+ using BlockScheduler = typename Gemm::Block::SplitkGemmIdentityBlockSwizzle<3, 0>;
	78	+ using GemmKernel = Gemm::Kernel::MultiCoreSplitkMatmulTla<BlockMmad, BlockEpilogue, BlockScheduler>;
	79	+"""
	80	+ _INPUT_TEMPLATE = """\
	81	+ uint32_t m = M;
	82	+ uint32_t k = K;
	83	+ uint32_t n = N;
	84	+"""
	85	+ _LAYOUT_TEMPLATE = """\
	86	+ GemmCoord problemShape{{m, n, k}};
	87	+ // Define the layout of each matrix
	88	+ LayoutTagA tagA{{m, k}};
	89	+ LayoutTagB tagB{{k, n}};
	90	+ LayoutTagC tagC{{m, n}};
	91	+ auto layoutA = tla::MakeLayoutFromTag(tagA);
	92	+ auto layoutB = tla::MakeLayoutFromTag(tagB);
	93	+ auto layoutC = tla::MakeLayoutFromTag(tagC);
	94	+"""
	95	+
	96	+ def get_default_tile_shape(self) -> Tuple[GemmShape, GemmShape]:
	97	+ element_max_size = max(
	98	+ self.element_A.data_size(),
	99	+ self.element_B.data_size(),
	100	+ self.element_C.data_size(),
	101	+ )
	102	+ # 根据是否传入bias返回不同的shape
	103	+ if self.element_Bias is not None and self.element_Bias != "void":
	104	+ return GemmShape(240, 256, 128), GemmShape(240, 256, 32)
	105	+ else:
	106	+ return GemmShape(256, 256, 128), GemmShape(256, 256, 32)
	107	+
	108	+ def get_default_dispatch_policy_list(self) -> List:
	109	+ """获取 MultiCoreSplitkMatmulKernel 的默认 dispatch_policy 列表.
	110	+
	111	+ :return: 包含默认 dispatch_policy 的列表，列表的第一个元素 [0] 是默认策略.
	112	+ :rtype: List
	113	+ """
	114	+ return [MmadPingpong(arch_tag=self.arch_tag, enable_unit_flag=True)]
	115	+
	116	+ def get_render_params(self, use_constexpr: bool = True) -> Dict[str, Any]:
	117	+ """获取渲染参数，包括 kernel 名称格式化参数."""
	118	+ params = super().get_render_params(use_constexpr)
	119	+ return self._add_kernel_name_params(params, self._KERNEL_NAME_BASE)

Apython/catlass_cppgen/catlass_cppgen/kernel/gemm/streamk_matmul.py+119-0

@@ -0,0 +1,119 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from typing import Any, Dict, List, Optional, Tuple
	11	+from catlass_cppgen.kernel.gemm.gemm_base import GemmKernelBase
	12	+from catlass_cppgen.catlass.gemm_coord import GemmCoord, GemmShape
	13	+from catlass_cppgen.catlass.layout.layout import Layout
	14	+from catlass_cppgen.common.data_type import DataType
	15	+from catlass_cppgen.common.typing import GM_ADDR
	16	+from catlass_cppgen.catlass.gemm.dispatch_policy import (
	17	+ MmadPingpong,
	18	+)
	19	+
	20	+
	21	+class StreamkMatmulKernel(GemmKernelBase):
	22	+ _KERNEL_NAME_BASE = "StreamkMatmulTla"
	23	+ _FEATURES = {"is_support_evg": False, "is_support_relu": False, "slice_axis": "K", "is_mix": True}
	24	+
	25	+ _INCLUDES = [
	26	+ "catlass/catlass.hpp",
	27	+ "catlass/arch/arch.hpp",
	28	+ "catlass/layout/layout.hpp",
	29	+ "catlass/status.hpp",
	30	+
	31	+ "catlass/gemm/block/block_mmad.hpp",
	32	+ "catlass/gemm/block/block_swizzle.hpp",
	33	+ "catlass/gemm/dispatch_policy.hpp",
	34	+ "catlass/gemm/gemm_type.hpp",
	35	+ "catlass/gemm/device/device_gemm.hpp",
	36	+ "catlass/gemm_coord.hpp",
	37	+ "catlass/matrix_coord.hpp",
	38	+
	39	+ "tla/layout.hpp",
	40	+
	41	+ "catlass/gemm/kernel/streamk_matmul_tla.hpp",
	42	+ ]
	43	+ _KERNEL_NAME = "{arch_name}_{kernel_name}_{dispatch_policy_name}_{swizzle_name}_{l1_tile_shape_str}_{l0_tile_shape_str}"
	44	+ _PARAMS_DEVICE = [
	45	+ (GemmCoord, "problemShape"),
	46	+ (GM_ADDR, "deviceA"),
	47	+ (Layout, "layoutA"),
	48	+ (GM_ADDR, "deviceB"),
	49	+ (Layout, "layoutB"),
	50	+ (GM_ADDR, "deviceC"),
	51	+ (Layout, "layoutC"),
	52	+ (DataType.AUTO, "aicCoreNum"),
	53	+ (GM_ADDR, "deviceBias"),
	54	+ ]
	55	+ _DISPATCH_POLICY = """\
	56	+ using ArchTag = {arch_tag};
	57	+{constexpr_declarations}
	58	+ using DispatchPolicy = {dispatch_policy_template};
	59	+"""
	60	+ _KERNEL_TEMPLATE = """\
	61	+ using L1TileShape = {l1_tile_shape_tla};
	62	+ using L0TileShape = {l0_tile_shape_tla};
	63	+
	64	+ using ElementA = {element_A};
	65	+ using ElementB = {element_B};
	66	+ using ElementC = {element_C};
	67	+ using ElementBias = {element_Bias};
	68	+ using ElementBiasType = std::conditional_t<std::is_void_v<ElementBias>, uint8_t, ElementBias>;
	69	+ using LayoutTagA = {layout_A};
	70	+ using LayoutTagB = {layout_B};
	71	+ using LayoutTagC = layout::RowMajor;
	72	+
	73	+ using TileCopy = Gemm::Tile::PackedTileCopyTla<ArchTag, ElementA, LayoutTagA, ElementB, LayoutTagB, ElementC, LayoutTagC, ElementBias>;
	74	+ using BlockMmad = Gemm::Block::BlockMmadTla<DispatchPolicy, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, ElementBias, TileCopy>;
	75	+ using BlockEpilogue = void;
	76	+
	77	+ using BlockScheduler = typename Gemm::Block::StreamkGemmIdentityBlockSwizzle<3, 0>;
	78	+ using GemmKernel = Gemm::Kernel::StreamkMatmulTla<BlockMmad, BlockEpilogue, BlockScheduler>;
	79	+"""
	80	+ _INPUT_TEMPLATE = """\
	81	+ uint32_t m = M;
	82	+ uint32_t k = K;
	83	+ uint32_t n = N;
	84	+"""
	85	+ _LAYOUT_TEMPLATE = """\
	86	+ GemmCoord problemShape{{m, n, k}};
	87	+ // Define the layout of each matrix
	88	+ LayoutTagA tagA{{m, k}};
	89	+ LayoutTagB tagB{{k, n}};
	90	+ LayoutTagC tagC{{m, n}};
	91	+ auto layoutA = tla::MakeLayoutFromTag(tagA);
	92	+ auto layoutB = tla::MakeLayoutFromTag(tagB);
	93	+ auto layoutC = tla::MakeLayoutFromTag(tagC);
	94	+"""
	95	+
	96	+ def get_default_tile_shape(self) -> Tuple[GemmShape, GemmShape]:
	97	+ element_max_size = max(
	98	+ self.element_A.data_size(),
	99	+ self.element_B.data_size(),
	100	+ self.element_C.data_size(),
	101	+ )
	102	+ # 根据是否传入bias返回不同的shape
	103	+ if self.element_Bias is not None and self.element_Bias != "void":
	104	+ return GemmShape(240, 256, 128), GemmShape(240, 256, 32)
	105	+ else:
	106	+ return GemmShape(256, 256, 128), GemmShape(256, 256, 32)
	107	+
	108	+ def get_default_dispatch_policy_list(self) -> List:
	109	+ """获取 StreamkMatmulKernel 的默认 dispatch_policy 列表.
	110	+
	111	+ :return: 包含默认 dispatch_policy 的列表，列表的第一个元素 [0] 是默认策略.
	112	+ :rtype: List
	113	+ """
	114	+ return [MmadPingpong(arch_tag=self.arch_tag, enable_unit_flag=True)]
	115	+
	116	+ def get_render_params(self, use_constexpr: bool = True) -> Dict[str, Any]:
	117	+ """获取渲染参数，包括 kernel 名称格式化参数."""
	118	+ params = super().get_render_params(use_constexpr)
	119	+ return self._add_kernel_name_params(params, self._KERNEL_NAME_BASE)

Apython/catlass_cppgen/catlass_cppgen/kernel/gemm/tail_multi_core_splitk_matmul.py+118-0

@@ -0,0 +1,118 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from typing import Any, Dict, List, Optional, Tuple
	11	+from catlass_cppgen.kernel.gemm.gemm_base import GemmKernelBase
	12	+from catlass_cppgen.catlass.gemm_coord import GemmCoord, GemmShape
	13	+from catlass_cppgen.catlass.layout.layout import Layout
	14	+from catlass_cppgen.common.data_type import DataType
	15	+from catlass_cppgen.common.typing import GM_ADDR
	16	+from catlass_cppgen.catlass.gemm.dispatch_policy import (
	17	+ MmadPingpong,
	18	+)
	19	+
	20	+
	21	+class TailMultiCoreSplitkMatmulKernel(GemmKernelBase):
	22	+ _KERNEL_NAME_BASE = "TailMultiCoreSplitkMatmulTla"
	23	+ _FEATURES = {"is_support_evg": False, "is_support_relu": False, "slice_axis": "K", "is_mix": True}
	24	+
	25	+ _INCLUDES = [
	26	+ "catlass/catlass.hpp",
	27	+ "catlass/arch/arch.hpp",
	28	+ "catlass/layout/layout.hpp",
	29	+ "catlass/status.hpp",
	30	+ "tla/layout.hpp",
	31	+
	32	+ "catlass/gemm/block/block_mmad.hpp",
	33	+ "catlass/gemm/block/block_swizzle.hpp",
	34	+ "catlass/gemm/dispatch_policy.hpp",
	35	+ "catlass/gemm/gemm_type.hpp",
	36	+ "catlass/gemm/device/device_gemm.hpp",
	37	+ "catlass/gemm_coord.hpp",
	38	+ "catlass/matrix_coord.hpp",
	39	+
	40	+ "catlass/gemm/kernel/tail_multi_core_splitk_matmul_tla.hpp",
	41	+ ]
	42	+ _KERNEL_NAME = "{arch_name}_{kernel_name}_{dispatch_policy_name}_{swizzle_name}_{l1_tile_shape_str}_{l0_tile_shape_str}"
	43	+ _PARAMS_DEVICE = [
	44	+ (GemmCoord, "problemShape"),
	45	+ (GM_ADDR, "deviceA"),
	46	+ (Layout, "layoutA"),
	47	+ (GM_ADDR, "deviceB"),
	48	+ (Layout, "layoutB"),
	49	+ (GM_ADDR, "deviceC"),
	50	+ (Layout, "layoutC"),
	51	+ (DataType.AUTO, "aicCoreNum"),
	52	+ (GM_ADDR, "deviceBias"),
	53	+ ]
	54	+ _DISPATCH_POLICY = """\
	55	+ using ArchTag = {arch_tag};
	56	+{constexpr_declarations}
	57	+ using DispatchPolicy = {dispatch_policy_template};
	58	+"""
	59	+ _KERNEL_TEMPLATE = """\
	60	+ using L1TileShape = {l1_tile_shape_tla};
	61	+ using L0TileShape = {l0_tile_shape_tla};
	62	+
	63	+ using ElementA = {element_A};
	64	+ using ElementB = {element_B};
	65	+ using ElementC = {element_C};
	66	+ using ElementBias = {element_Bias};
	67	+ using ElementBiasType = std::conditional_t<std::is_void_v<ElementBias>, uint8_t, ElementBias>;
	68	+ using LayoutTagA = {layout_A};
	69	+ using LayoutTagB = {layout_B};
	70	+ using LayoutTagC = layout::RowMajor;
	71	+
	72	+ using TileCopy = Gemm::Tile::PackedTileCopyTla<ArchTag, ElementA, LayoutTagA, ElementB, LayoutTagB, ElementC, LayoutTagC, ElementBias>;
	73	+ using BlockMmad = Gemm::Block::BlockMmadTla<DispatchPolicy, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, ElementBias, TileCopy>;
	74	+ using BlockEpilogue = void;
	75	+
	76	+ using BlockScheduler = typename Gemm::Block::TailSplitkGemmIdentityBlockSwizzle<3, 0>;
	77	+ using GemmKernel = Gemm::Kernel::TailMultiCoreSplitkMatmulTla<BlockMmad, BlockEpilogue, BlockScheduler>;
	78	+"""
	79	+ _INPUT_TEMPLATE = """\
	80	+ uint32_t m = M;
	81	+ uint32_t k = K;
	82	+ uint32_t n = N;
	83	+"""
	84	+ _LAYOUT_TEMPLATE = """\
	85	+ GemmCoord problemShape{{m, n, k}};
	86	+ // Define the layout of each matrix
	87	+ LayoutTagA tagA{{m, k}};
	88	+ LayoutTagB tagB{{k, n}};
	89	+ LayoutTagC tagC{{m, n}};
	90	+ auto layoutA = tla::MakeLayoutFromTag(tagA);
	91	+ auto layoutB = tla::MakeLayoutFromTag(tagB);
	92	+ auto layoutC = tla::MakeLayoutFromTag(tagC);
	93	+"""
	94	+
	95	+ def get_default_tile_shape(self) -> Tuple[GemmShape, GemmShape]:
	96	+ element_max_size = max(
	97	+ self.element_A.data_size(),
	98	+ self.element_B.data_size(),
	99	+ self.element_C.data_size(),
	100	+ )
	101	+ # 根据是否传入bias返回不同的shape
	102	+ if self.element_Bias is not None and self.element_Bias != "void":
	103	+ return GemmShape(240, 256, 128), GemmShape(240, 256, 32)
	104	+ else:
	105	+ return GemmShape(256, 256, 128), GemmShape(256, 256, 32)
	106	+
	107	+ def get_default_dispatch_policy_list(self) -> List:
	108	+ """获取 TailMultiCoreSplitkMatmulKernel 的默认 dispatch_policy 列表.
	109	+
	110	+ :return: 包含默认 dispatch_policy 的列表，列表的第一个元素 [0] 是默认策略.
	111	+ :rtype: List
	112	+ """
	113	+ return [MmadPingpong(arch_tag=self.arch_tag, enable_unit_flag=True)]
	114	+
	115	+ def get_render_params(self, use_constexpr: bool = True) -> Dict[str, Any]:
	116	+ """获取渲染参数，包括 kernel 名称格式化参数."""
	117	+ params = super().get_render_params(use_constexpr)
	118	+ return self._add_kernel_name_params(params, self._KERNEL_NAME_BASE)

Apython/catlass_cppgen/catlass_cppgen/kernel/group_gemm/__init__.py+14-0

@@ -0,0 +1,14 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from catlass_cppgen.kernel.group_gemm.grouped_matmul_slice_m import GroupedMatmulSliceMKernel
	11	+
	12	+__all__ = [
	13	+ "GroupedMatmulSliceMKernel",
	14	+]

Apython/catlass_cppgen/catlass_cppgen/kernel/group_gemm/grouped_matmul_slice_m.py+136-0

@@ -0,0 +1,136 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from typing import Any, Dict, List, Optional, Tuple
	11	+from catlass_cppgen.kernel.gemm.gemm_base import GemmKernelBase
	12	+from catlass_cppgen.catlass.gemm_coord import GemmCoord, GemmShape
	13	+from catlass_cppgen.catlass.layout.layout import Layout
	14	+from catlass_cppgen.common.typing import GM_ADDR
	15	+from catlass_cppgen.common.data_type import DataType
	16	+from catlass_cppgen.common.utils import get_type_name
	17	+from catlass_cppgen.catlass.gemm.dispatch_policy import (
	18	+ MmadPingpong,
	19	+)
	20	+
	21	+
	22	+class GroupedMatmulSliceMKernel(GemmKernelBase):
	23	+ _KERNEL_NAME_BASE = "GroupedMatmulSliceMTla"
	24	+ _FEATURES = {"is_support_evg": False, "is_support_relu": True, "slice_axis": "M", "is_mix": True}
	25	+
	26	+ _INCLUDES = [
	27	+ "catlass/catlass.hpp",
	28	+ "catlass/arch/arch.hpp",
	29	+ "catlass/layout/layout.hpp",
	30	+ "catlass/status.hpp",
	31	+
	32	+ "catlass/gemm/block/block_mmad.hpp",
	33	+ "catlass/gemm/block/block_swizzle.hpp",
	34	+ "catlass/gemm/dispatch_policy.hpp",
	35	+ "catlass/gemm/gemm_type.hpp",
	36	+ "catlass/gemm/device/device_gemm.hpp",
	37	+ "catlass/gemm_coord.hpp",
	38	+ "catlass/matrix_coord.hpp",
	39	+
	40	+ "catlass/gemm/kernel/grouped_matmul_slice_m_tla.hpp",
	41	+ "tla/layout.hpp",
	42	+ ]
	43	+ _KERNEL_NAME = "{arch_name}_{kernel_name}_{dispatch_policy_name}_{swizzle_name}_{l1_tile_shape_str}_{l0_tile_shape_str}"
	44	+ _PARAMS_DEVICE = [
	45	+ (GemmCoord, "problemShape"),
	46	+ (DataType.UINT32, "problemCount"),
	47	+ (GM_ADDR, "deviceGroupList"),
	48	+ (GM_ADDR, "deviceA"),
	49	+ (Layout, "layoutA"),
	50	+ (GM_ADDR, "deviceB"),
	51	+ (Layout, "layoutB"),
	52	+ (GM_ADDR, "deviceC"),
	53	+ (Layout, "layoutC"),
	54	+ ]
	55	+ _DISPATCH_POLICY = """\
	56	+ using ArchTag = {arch_tag};
	57	+{constexpr_declarations}
	58	+ using DispatchPolicy = {dispatch_policy_template};
	59	+"""
	60	+ _KERNEL_TEMPLATE = """\
	61	+ using L1TileShape = {l1_tile_shape_tla};
	62	+ using L0TileShape = {l0_tile_shape_tla};
	63	+
	64	+ using ElementA = {element_A};
	65	+ using ElementB = {element_B};
	66	+ using ElementC = {element_C};
	67	+ using LayoutTagA = {layout_A};
	68	+ using LayoutTagB = {layout_B};
	69	+ using LayoutTagC = layout::RowMajor;
	70	+
	71	+ using TileCopy = Gemm::Tile::PackedTileCopyTla<ArchTag, ElementA, LayoutTagA, ElementB, LayoutTagB, ElementC, LayoutTagC, void, {relu_enable}>;
	72	+ using BlockMmadTla = Gemm::Block::BlockMmadTla<DispatchPolicy, L1TileShape, L0TileShape, ElementA, ElementB, ElementC, void, TileCopy>;
	73	+ using BlockEpilogue = void;
	74	+
	75	+ using BlockScheduler = typename Gemm::Block::GemmIdentityBlockSwizzle<3, 0>;
	76	+ using GemmKernel = Gemm::Kernel::GroupedMatmulSliceMTla<BlockMmadTla, BlockEpilogue, BlockScheduler, {groupList_element_type}>;
	77	+"""
	78	+ _INPUT_TEMPLATE = """\
	79	+ uint32_t problemCount = {problemCount};
	80	+ uint32_t m = M;
	81	+ uint32_t k = K;
	82	+ uint32_t n = N;
	83	+"""
	84	+ _LAYOUT_TEMPLATE = """\
	85	+ GemmCoord problemShape{{m, n, k}};
	86	+ // Define the layout of each matrix
	87	+ LayoutTagA tagA{{m, k}};
	88	+ LayoutTagB tagB{{k, n}};
	89	+ LayoutTagC tagC{{m, n}};
	90	+ auto layoutA = tla::MakeLayoutFromTag(tagA);
	91	+ auto layoutB = tla::MakeLayoutFromTag(tagB);
	92	+ auto layoutC = tla::MakeLayoutFromTag(tagC);
	93	+"""
	94	+
	95	+ def __init__(self, problemCount: Optional[int] = None, groupList_element: Optional[DataType] = None, **kwargs):
	96	+ """初始化 GroupedMatmulSliceMKernel.
	97	+
	98	+ :param problemCount: 问题数量，如果为 None 则使用默认值 1
	99	+ :param groupList_element: groupList 的数据类型，如果为 None 则使用默认值 int64_t
	100	+ :param kwargs: 传递给父类的其他参数，包括 M, K, N 等
	101	+ """
	102	+ super().__init__(**kwargs)
	103	+ self.problemCount = problemCount if problemCount is not None else 1
	104	+ self.groupList_element = groupList_element if groupList_element is not None else DataType.INT64
	105	+
	106	+ def get_default_tile_shape(self) -> Tuple[GemmShape, GemmShape]:
	107	+ """获取默认的 tile shape.
	108	+
	109	+ 根据 grouped_matmul_slice_m.cpp 中的设置：
	110	+ L1TileShape = Shape<Int<256>, Int<256>, Int<256>>
	111	+ L0TileShape = Shape<Int<256>, Int<256>, Int<64>>
	112	+ """
	113	+ return GemmShape(256, 256, 256), GemmShape(256, 256, 64)
	114	+
	115	+ def get_default_dispatch_policy_list(self) -> List:
	116	+ """获取 GroupedMatmulSliceMKernel 的默认 dispatch_policy 列表.
	117	+
	118	+ :return: 包含默认 dispatch_policy 的列表，列表的第一个元素 [0] 是默认策略.
	119	+ :rtype: List
	120	+ """
	121	+ return [MmadPingpong(arch_tag=self.arch_tag, enable_unit_flag=True)]
	122	+
	123	+ def get_render_params(self, use_constexpr: bool = True) -> Dict[str, Any]:
	124	+ """获取渲染参数，包括动态生成的 dispatch_policy C++ 代码.
	125	+
	126	+ :param use_constexpr: 当为 True 时，生成包含常量声明的完整代码块；当为 False 时，只生成 using 语句（使用变量名）
	127	+ :return: 渲染参数字典.
	128	+ :rtype: Dict[str, Any]
	129	+ """
	130	+ params = super().get_render_params(use_constexpr)
	131	+ params['problemCount'] = self.problemCount
	132	+ # 将 groupList_element 转换为 C++ 类型字符串
	133	+ params['groupList_element_type'] = get_type_name(self.groupList_element)
	134	+ # 添加 relu_enable 参数
	135	+ params['relu_enable'] = 'true' if self.relu_enable else 'false'
	136	+ return self._add_kernel_name_params(params, self._KERNEL_NAME_BASE)

Apython/catlass_cppgen/catlass_cppgen/kernel/kernel_base.py+366-0

@@ -0,0 +1,366 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from abc import abstractmethod
	11	+import warnings
	12	+from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union
	13	+
	14	+from catlass_cppgen.catlass.gemm_coord import GemmCoord, GemmShape
	15	+from catlass_cppgen.common.typing import GM_ADDR
	16	+from catlass_cppgen.common.data_type import DataType
	17	+from catlass_cppgen.common.utils import get_type_name
	18	+
	19	+DispatchPolicy = TypeVar("DispatchPolicy")
	20	+BlockScheduler = TypeVar("BlockScheduler")
	21	+
	22	+def _get_deprecated_arg(kwargs: Dict, key: Any, sub_key: Optional[str] = None) -> Any:
	23	+ if key in kwargs:
	24	+ warnings.warn(
	25	+ f"{key!r} is deprecated." + (f" Use {sub_key!r} instead." if sub_key else ""),
	26	+ DeprecationWarning,
	27	+ stacklevel=3
	28	+ )
	29	+ v = kwargs.pop(key)
	30	+ else:
	31	+ v = None
	32	+ return v
	33	+
	34	+class KernelBase:
	35	+ """Kernel Base Class"""
	36	+ _INCLUDES: List[str] = [] # 头文件引入
	37	+ _PARAMS_DEVICE: List[Tuple[str, Type]] = [] # 核函数参数
	38	+ _KERNEL_NAME: str = "" # 核函数名
	39	+ _DISPATCH_POLICY: str = "" # dispatch policy模板
	40	+ _KERNEL_TEMPLATE: str = "" # 核函数模板
	41	+ _INPUT_TEMPLATE: str = "" # 输入模板（用于定义 m, k, n 等变量）
	42	+ _LAYOUT_TEMPLATE: str = "" # layout模板
	43	+ _ADDITIONAL_DEFINITIONS_TEMPLATE: str = "" # 额外定义模板
	44	+ _FEATURES: Dict[str, Any] = {} # 核函数特性字典
	45	+ # 参数插入映射：在哪些参数名之后插入什么参数
	46	+ # 例如：{"layoutA": "strideA", "layoutB": "strideB"} 表示在 layoutA 后插入 strideA
	47	+ _PARAMS_INSERTIONS: Dict[str, str] = {} # 参数插入映射
	48	+
	49	+
	50	+ def __init__(self, args, *kwargs):
	51	+ self.l1_tile_shape, self.l0_tile_shape = self.get_default_tile_shape()
	52	+ self.block_scheduler = None
	53	+ self.dispatch_policy = self.get_default_dispatch_policy_list()
	54	+ self._features = {}
	55	+ # 从类属性中读取特性值
	56	+ if hasattr(self.__class__, "_FEATURES"):
	57	+ self._features.update(self.__class__._FEATURES)
	58	+ # 初始化 relu_enable，默认为 False
	59	+ self.relu_enable = False
	60	+
	61	+ @abstractmethod
	62	+ def get_default_tile_shape(self) -> Tuple[GemmShape, GemmShape]:
	63	+ pass
	64	+
	65	+ def get_default_dispatch_policy_list(self) -> List[DispatchPolicy]:
	66	+ """获取默认的 dispatch_policy 列表.
	67	+
	68	+ 子类可以重写此方法以定义自己的默认 dispatch_policy 列表。
	69	+ 如果子类不重写，默认返回空列表。
	70	+
	71	+ :return: 默认的 dispatch_policy 列表.
	72	+ :rtype: List[DispatchPolicy]
	73	+ """
	74	+ return []
	75	+
	76	+ """op interface,尚未实现"""
	77	+
	78	+ def get_workspace_size(self) -> int:
	79	+ return 0
	80	+
	81	+ def need_workspace(self) -> bool:
	82	+ return False
	83	+
	84	+ def get_core_num(self) -> int:
	85	+ return 0
	86	+
	87	+ """flag interface"""
	88	+
	89	+ def is_support(self, feature: str) -> bool:
	90	+ """检查是否支持某个特性"""
	91	+ pass
	92	+
	93	+ def __getattr__(self, name: str) -> Any:
	94	+ """通过属性方式访问特性值"""
	95	+ # 对于特殊属性（如 __setstate__, __getstate__ 等），直接抛出 AttributeError
	96	+ # 避免干扰 deepcopy、pickle 等内部操作
	97	+ if name.startswith('__') and name.endswith('__'):
	98	+ raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
	99	+
	100	+ # 直接检查 __dict__ 而不是使用 hasattr，避免递归
	101	+ # _features 在 __init__ 中总是会被初始化，所以可以直接检查
	102	+ if "_features" in self.__dict__ and name in self._features:
	103	+ return self._features[name]
	104	+ raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
	105	+
	106	+
	107	+ """tune interface"""
	108	+
	109	+ def set_l1_tile_shape(self, l1_tile_shape: GemmShape):
	110	+ """设置L1 Tile Shape.
	111	+
	112	+ :param l1_tile_shape: L1 Tile Shape.
	113	+ :type l1_tile_shape: GemmShape
	114	+ """
	115	+ self.l1_tile_shape = l1_tile_shape
	116	+
	117	+ def set_l0_tile_shape(self, l0_tile_shape: GemmShape):
	118	+ """设置L0 Tile Shape.
	119	+
	120	+ :param l0_tile_shape: L0 Tile Shape.
	121	+ :type l0_tile_shape: GemmShape
	122	+ """
	123	+ self.l0_tile_shape = l0_tile_shape
	124	+
	125	+ def set_dispatch_policy(self, dispatch_policy: Union[DispatchPolicy, List[DispatchPolicy]]):
	126	+ """设置Dispatch Policy.
	127	+
	128	+ :param dispatch_policy: Dispatch Policy 或 Dispatch Policy 列表。如果传入单个 policy，会转换为包含该 policy 的列表。
	129	+ :type dispatch_policy: Union[DispatchPolicy, List[DispatchPolicy]]
	130	+ """
	131	+ if isinstance(dispatch_policy, list):
	132	+ self.dispatch_policy = dispatch_policy
	133	+ else:
	134	+
	135	+ self.dispatch_policy = [dispatch_policy] if dispatch_policy is not None else []
	136	+
	137	+ def get_dispatch_policy(self) -> List[DispatchPolicy]:
	138	+ """获取Dispatch Policy 列表.
	139	+
	140	+ :return: Dispatch Policy 列表。默认情况下返回包含默认 policy 的列表。
	141	+ :rtype: List[DispatchPolicy]
	142	+ """
	143	+ return self.dispatch_policy
	144	+
	145	+ def set_block_scheduler(self, block_scheduler: BlockScheduler):
	146	+ """设置Block Scheduler.
	147	+
	148	+ :param block_scheduler: Block Scheduler.
	149	+ :type block_scheduler: BlockScheduler
	150	+ """
	151	+ self.block_scheduler = block_scheduler
	152	+
	153	+ def set_relu_enable(self, relu_enable: bool):
	154	+ """设置是否启用随路 Relu.
	155	+
	156	+ :param relu_enable: 是否启用随路 Relu.
	157	+ :type relu_enable: bool
	158	+ """
	159	+ self.relu_enable = relu_enable
	160	+
	161	+ def set_use_hf32_mode(self, is_hf32: bool):
	162	+ """设置是否启用 HF32 模式.
	163	+
	164	+ 此方法会检查当前 dispatch_policy 列表中是否包含 use_hf32_mode 属性，
	165	+ 如果包含则设置为指定值，如果不包含则抛出提示信息。
	166	+
	167	+ :param is_hf32: 是否启用 HF32 模式.
	168	+ :type is_hf32: bool
	169	+ :raises ValueError: 如果 dispatch_policy 中没有任何 policy 支持 use_hf32_mode
	170	+ """
	171	+ if not self.dispatch_policy:
	172	+ raise ValueError("dispatch_policy 列表为空，无法设置 use_hf32_mode")
	173	+
	174	+ policies_with_hf32 = []
	175	+ policies_without_hf32 = []
	176	+
	177	+ for policy in self.dispatch_policy:
	178	+ if hasattr(policy, 'use_hf32_mode'):
	179	+ policy.use_hf32_mode = is_hf32
	180	+ policies_with_hf32.append(policy.__class__.__name__)
	181	+ else:
	182	+ policies_without_hf32.append(policy.__class__.__name__)
	183	+
	184	+ if not policies_with_hf32:
	185	+ policy_names = ", ".join(policies_without_hf32)
	186	+ raise ValueError(
	187	+ f"当前 dispatch_policy 中的 policy ({policy_names}) 不支持 use_hf32_mode。"
	188	+ f"支持 use_hf32_mode 的 policy 包括: MmadPingpong, MmadPreloadAsyncWithCallback, MmadMultiBatch"
	189	+ )
	190	+
	191	+ def tune(
	192	+ self,
	193	+ l1_tile_shape: Optional[GemmShape] = None,
	194	+ l0_tile_shape: Optional[GemmShape] = None,
	195	+ dispatch_policy: Optional[Union[DispatchPolicy, List[DispatchPolicy]]] = None, # reserved
	196	+ block_scheduler: Optional[BlockScheduler] = None, # reserved
	197	+ relu_enable: Optional[bool] = None,
	198	+ is_hf32: Optional[bool] = None,
	199	+ **kwargs
	200	+ ):
	201	+ if_hf32 = _get_deprecated_arg(kwargs, "if_hf32", "is_hf32")
	202	+ if is_hf32 is None and if_hf32 is not None:
	203	+ is_hf32 = if_hf32
	204	+ elif if_hf32 is not None and if_hf32 != is_hf32:
	205	+ raise ValueError("There is a conflict between suggested 'is_hf32' and deprecated 'if_hf32'.")
	206	+
	207	+ self.set_l1_tile_shape(l1_tile_shape or self.l1_tile_shape)
	208	+ self.set_l0_tile_shape(l0_tile_shape or self.l0_tile_shape)
	209	+ if dispatch_policy is not None:
	210	+ self.set_dispatch_policy(dispatch_policy)
	211	+ self.set_block_scheduler(block_scheduler or self.block_scheduler)
	212	+ if relu_enable is not None:
	213	+ self.set_relu_enable(relu_enable)
	214	+ if is_hf32 is not None:
	215	+ self.set_use_hf32_mode(is_hf32)
	216	+
	217	+ """codegen interface"""
	218	+ @abstractmethod
	219	+ def get_render_params(self) -> Dict[str, Any]:
	220	+ pass
	221	+
	222	+ def gen_includes(self) -> str:
	223	+ """生成头文件引入部分源码.
	224	+
	225	+ :return: 头文件.
	226	+ :rtype: str
	227	+ """
	228	+ return "\n".join(
	229	+ ["#include <{}>".format(include) for include in self._INCLUDES]
	230	+ )
	231	+
	232	+ def gen_kernel_name(self) -> str:
	233	+ """生成核函数名.
	234	+
	235	+ :return: 渲染后的核函数名.
	236	+ :rtype: str
	237	+ """
	238	+ return self._KERNEL_NAME.format(**self.get_render_params())
	239	+
	240	+ def gen_params_device(self, def_mode: bool = False) -> str:
	241	+ """生成核函数参数.
	242	+
	243	+ :param def_mode: 是否为定义模式. 在定义模式下，会生成函数定义中的形式，如`int a, int b`.
	244	+ 否则，生成函数调用中的形式，如`a, b`.
	245	+ :type def_mode: bool
	246	+ :return: 核函数参数.
	247	+ :rtype: str
	248	+ """
	249	+ if def_mode:
	250	+ generated = ", ".join(
	251	+ [
	252	+ "{} {}".format(get_type_name(type_str), var_name)
	253	+ for type_str, var_name in self._PARAMS_DEVICE
	254	+ ]
	255	+ )
	256	+ else:
	257	+ generated = ", ".join(
	258	+ ["{}".format(var_name) for _, var_name in self._PARAMS_DEVICE]
	259	+ )
	260	+ return generated
	261	+
	262	+ def gen_kernel_template(self) -> str:
	263	+ """生成核函数组装部分模板.
	264	+ 比如，从`using Arch=ArchTag::AtlasA2`到`kernel(params);`之间的部分.
	265	+
	266	+ :return: 渲染后的核函数模板.
	267	+ :rtype: str
	268	+ """
	269	+ render_params = self.get_render_params()
	270	+ for key, value in render_params.items():
	271	+ if hasattr(value, "value"):
	272	+ value = value.value
	273	+ render_params[key] = str(value)
	274	+ result = self._DISPATCH_POLICY.format(render_params) + self._KERNEL_TEMPLATE.format(render_params)
	275	+ return result
	276	+
	277	+ def gen_input_template(self) -> str:
	278	+ """生成输入变量定义代码.
	279	+ 生成包含 M, K, N 等输入变量的定义代码块.
	280	+
	281	+ :return: 输入变量定义代码，如果不存在 _INPUT_TEMPLATE 则返回空字符串.
	282	+ :rtype: str
	283	+ """
	284	+ if not hasattr(self.__class__, "_INPUT_TEMPLATE") or not self.__class__._INPUT_TEMPLATE:
	285	+ return ""
	286	+
	287	+ render_params = self.get_render_params()
	288	+ for key, value in render_params.items():
	289	+ if hasattr(value, "value"):
	290	+ value = value.value
	291	+ render_params[key] = str(value)
	292	+
	293	+ return self.__class__._INPUT_TEMPLATE.format(**render_params)
	294	+
	295	+ def gen_layout_template(self) -> str:
	296	+ """生成 layout 相关信息代码.
	297	+ 生成包含 M, K, N 定义和 layout tag 的代码块.
	298	+
	299	+ :return: layout 相关代码.
	300	+ :rtype: str
	301	+ """
	302	+ render_params = self.get_render_params()
	303	+ for key, value in render_params.items():
	304	+ if hasattr(value, "value"):
	305	+ value = value.value
	306	+ render_params[key] = str(value)
	307	+ return self._LAYOUT_TEMPLATE.format(
	308	+ **render_params,
	309	+ )
	310	+
	311	+ def _insert_params_recursive(self, param: str, params_list: List[str]) -> None:
	312	+ """递归插入参数.
	313	+
	314	+ 如果 param 在 _PARAMS_INSERTIONS 中，插入对应的参数，并递归检查插入的参数。
	315	+
	316	+ :param param: 当前参数名.
	317	+ :type param: str
	318	+ :param params_list: 参数列表（会被修改）.
	319	+ :type params_list: List[str]
	320	+ """
	321	+ if param in self._PARAMS_INSERTIONS:
	322	+ inserted_param = self._PARAMS_INSERTIONS[param]
	323	+ params_list.append(inserted_param)
	324	+ # 递归检查插入的参数是否也需要插入其他参数
	325	+ self._insert_params_recursive(inserted_param, params_list)
	326	+
	327	+ def transform_params_for_construction(self, params_str: str) -> str:
	328	+ """转换参数列表用于构造 Params 对象.
	329	+
	330	+ 子类可以重写此方法来修改参数列表，例如在特定参数后插入额外的参数。
	331	+ 支持链式插入：如果插入的参数本身也在 _PARAMS_INSERTIONS 中，会继续插入。
	332	+ 例如：在 layoutA 后插入 strideA，在 strideA 后插入 strideB。
	333	+
	334	+ :param params_str: 原始参数字符串，格式如 "param1, param2, param3".
	335	+ :type params_str: str
	336	+ :return: 转换后的参数字符串.
	337	+ :rtype: str
	338	+ """
	339	+ if not self._PARAMS_INSERTIONS:
	340	+ return params_str
	341	+ params_list = []
	342	+ for param in params_str.split(", "):
	343	+ params_list.append(param)
	344	+ self._insert_params_recursive(param, params_list)
	345	+ return ", ".join(params_list)
	346	+
	347	+ def _gen_kernel_params_for_def(self) -> str:
	348	+ """生成函数定义时的参数列表，只包含 GM_ADDR 类型的参数和 M, N, K.
	349	+
	350	+ :return: 函数定义参数列表.
	351	+ :rtype: str
	352	+ """
	353	+ # 只包含 GM_ADDR 类型的参数
	354	+ gm_addr_params = [
	355	+ (type_str, var_name)
	356	+ for type_str, var_name in self._PARAMS_DEVICE
	357	+ if type_str == GM_ADDR
	358	+ ]
	359	+ result = ", ".join(
	360	+ [
	361	+ "{} {}".format(get_type_name(type_str), var_name)
	362	+ for type_str, var_name in gm_addr_params
	363	+ ]
	364	+ )
	365	+ result += ", uint32_t M, uint32_t N, uint32_t K"
	366	+ return result

Apython/catlass_cppgen/catlass_cppgen/kernel/visitor_kernel_base.py+40-0

@@ -0,0 +1,40 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from typing import Any, Dict
	11	+
	12	+from catlass_cppgen.kernel.kernel_base import KernelBase
	13	+
	14	+
	15	+class VisitorKernelBase(KernelBase):
	16	+ """Visitor Kernel Base Class
	17	+
	18	+ 继承自 KernelBase，提供 EVG 模板生成功能。
	19	+ """
	20	+ _EVG_TEMPLATE: str = "" # EVG 模板
	21	+
	22	+ def gen_evg_template(self) -> str:
	23	+ """生成 EVG 相关代码.
	24	+
	25	+ 子类可以通过定义 _EVG_TEMPLATE 类变量来生成 EVG 相关代码，
	26	+ 或者重写此方法。
	27	+ 默认检查 _EVG_TEMPLATE，如果存在则使用模板渲染，否则返回空字符串。
	28	+
	29	+ :return: EVG 相关代码.
	30	+ :rtype: str
	31	+ """
	32	+ if not self._EVG_TEMPLATE:
	33	+ return ""
	34	+
	35	+ render_params = self.get_render_params()
	36	+ for key, value in render_params.items():
	37	+ if hasattr(value, "value"):
	38	+ value = value.value
	39	+ render_params[key] = str(value)
	40	+ return self._EVG_TEMPLATE.format(**render_params)

Apython/catlass_cppgen/catlass_cppgen/op/__init__.py+12-0

@@ -0,0 +1,12 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from catlass_cppgen.op.op import OperationBase
	11	+from catlass_cppgen.op.gemm import Gemm
	12	+from catlass_cppgen.op.group_gemm import GroupGemm

Apython/catlass_cppgen/catlass_cppgen/op/gemm.py+268-0

@@ -0,0 +1,268 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+import warnings
	11	+from typing import List, Type, Optional, Dict, Any
	12	+import torch
	13	+import numpy as np
	14	+import math
	15	+
	16	+from catlass_cppgen.op.op import OperationBase
	17	+from catlass_cppgen.common.typing import SupportedDataType, SupportedTensor
	18	+from catlass_cppgen.common.op_tensor import OpTensor
	19	+from catlass_cppgen.catlass.layout.layout import Layout, RowMajor
	20	+from catlass_cppgen.common.data_type import DataType, get_default_accumulator
	21	+from catlass_cppgen.common.utils import extract_info
	22	+from catlass_cppgen.catlass.arch.arch import Arch
	23	+from catlass_cppgen.kernel.gemm.gemm_base import GemmKernelBase
	24	+from catlass_cppgen.kernel.gemm.basic_matmul import BasicMatmulKernel
	25	+from catlass_cppgen.kernel.gemm.batched_matmul import BatchedMatmulKernel
	26	+from catlass_cppgen.kernel.gemm.multi_core_splitk_matmul import MultiCoreSplitkMatmulKernel
	27	+from catlass_cppgen.kernel.gemm.streamk_matmul import StreamkMatmulKernel
	28	+
	29	+from catlass_cppgen.kernel.gemm.tail_multi_core_splitk_matmul import TailMultiCoreSplitkMatmulKernel
	30	+from catlass_cppgen.kernel.gemm.basic_matmul_tla_visitor import BasicMatmulTlaVisitorKernel
	31	+
	32	+_CORE_NUM = 8 # Default core num
	33	+
	34	+class Gemm(OperationBase):
	35	+ # 类属性类型注解
	36	+ A: Optional[OpTensor]
	37	+ B: Optional[OpTensor]
	38	+ Bias: Optional[OpTensor]
	39	+ C: Optional[OpTensor]
	40	+ M: Optional[int]
	41	+ K: Optional[int]
	42	+ N: Optional[int]
	43	+ batch_count: Optional[int]
	44	+ is_batched: Optional[bool]
	45	+
	46	+ def __init__(
	47	+ self,
	48	+ alpha: float = 1.0,
	49	+ beta: float = 0.0,
	50	+ element_accumulator: Optional[SupportedDataType] = None,
	51	+ element: Optional[SupportedDataType] = None,
	52	+ layout: Optional[Type[Layout]] = None,
	53	+ element_A: Optional[SupportedDataType] = None,
	54	+ element_B: Optional[SupportedDataType] = None,
	55	+ element_Bias: Optional[SupportedDataType] = None,
	56	+ element_C: Optional[SupportedDataType] = None,
	57	+ layout_A: Optional[Type[Layout]] = None,
	58	+ layout_B: Optional[Type[Layout]] = None,
	59	+ layout_Bias: Optional[Type[Layout]] = None,
	60	+ evg_config: Optional[Dict[str, Any]] = None,
	61	+ atlas_arch: Optional[Arch] = None,
	62	+ core_num: Optional[int] = None,
	63	+ A: Optional[OpTensor] = None,
	64	+ B: Optional[OpTensor] = None,
	65	+ Bias: Optional[OpTensor] = None,
	66	+ C: Optional[OpTensor] = None,
	67	+ ):
	68	+ # 保存 A、B、Bias 和 C OpTensor
	69	+ self.A = A
	70	+ self.B = B
	71	+ self.Bias = Bias
	72	+ self.C = C
	73	+
	74	+ # 如果传入了 A 和 B OpTensor，从它们中提取信息
	75	+ if A is not None and B is not None:
	76	+ # 从 OpTensor 中提取 shape、element、layout 信息
	77	+ A_shape, element_A_from_tensor, layout_A_from_tensor, _ = extract_info(A, element_A or element, layout_A or layout)
	78	+ B_shape, element_B_from_tensor, layout_B_from_tensor, _ = extract_info(B, element_B or element, layout_B or layout)
	79	+
	80	+ # 确定最终使用的 element 和 layout（传入的参数优先，否则使用 OpTensor 中的信息）
	81	+ element_A = element_A_from_tensor or element_A or element
	82	+ element_B = element_B_from_tensor or element_B or element
	83	+ layout_A = layout_A_from_tensor or layout_A or layout
	84	+ layout_B = layout_B_from_tensor or layout_B or layout
	85	+
	86	+ # 判断是否 batched
	87	+ is_batched = len(A_shape) == 3 and len(B_shape) == 3
	88	+
	89	+ # 提取 M, K, N
	90	+ if is_batched:
	91	+ if A_shape[0] != B_shape[0]:
	92	+ raise ValueError(
	93	+ f"A.shape[0] ({A_shape[0]}) must be equal to B.shape[0] ({B_shape[0]}) for batched matmul"
	94	+ )
	95	+ self.batch_count = A_shape[0]
	96	+ self.M = A_shape[1]
	97	+ self.K = A_shape[2]
	98	+ self.N = B_shape[2]
	99	+ else:
	100	+ self.batch_count = None
	101	+ self.M = A_shape[0]
	102	+ self.K = A_shape[1]
	103	+ self.N = B_shape[1]
	104	+
	105	+ self.is_batched = is_batched
	106	+ else:
	107	+ if element is None and not all([element_A, element_B, element_C]):
	108	+ raise ValueError(
	109	+ "must provide 'element', or specify element_A, element_B, element_C separately"
	110	+ )
	111	+ if layout is None and not all([layout_A, layout_B]):
	112	+ raise ValueError(
	113	+ "must provide 'layout', or specify layout_A, layout_B separately"
	114	+ )
	115	+
	116	+ element_A = element_A or element
	117	+ element_B = element_B or element
	118	+ layout_A = layout_A or layout
	119	+ layout_B = layout_B or layout
	120	+
	121	+ self.batch_count = None
	122	+ self.M = None
	123	+ self.K = None
	124	+ self.N = None
	125	+ self.is_batched = None
	126	+
	127	+ # 如果传入了 Bias OpTensor，从它中提取信息
	128	+ if Bias is not None:
	129	+ _, element_Bias_from_tensor, layout_Bias_from_tensor, _ = extract_info(Bias, element_Bias, layout_Bias)
	130	+ element_Bias = element_Bias_from_tensor or element_Bias
	131	+ layout_Bias = layout_Bias_from_tensor or layout_Bias
	132	+
	133	+ # 如果传入了 C OpTensor，从它中提取信息
	134	+ if C is not None:
	135	+ _, element_C_from_tensor, layout_C_from_tensor, _ = extract_info(C, element_C or element, None)
	136	+ element_C = element_C_from_tensor or element_C or element
	137	+
	138	+ self.element_A = element_A
	139	+ self.element_B = element_B
	140	+ self.element_Bias = element_Bias
	141	+ self.element_C = element_C or element
	142	+ if element_accumulator is None and not all([self.element_A, self.element_B]):
	143	+ raise ValueError("element_accumulator must be provided, or element_A, element_B should be given both so that accumulator type can be auto-derived")
	144	+ self.element_accumulator = element_accumulator or get_default_accumulator(
	145	+ self.element_A, self.element_B
	146	+ )
	147	+ self.layout_A = layout_A
	148	+ self.layout_B = layout_B
	149	+ # layout_Bias不提供则为 None
	150	+ self.layout_Bias = layout_Bias
	151	+ # layout_C 固定为 RowMajor，如果 M 和 N 已确定则实例化
	152	+ if self.M is not None and self.N is not None:
	153	+ self.layout_C = RowMajor((self.M, self.N))
	154	+ else:
	155	+ # 如果 M 和 N 未确定，先保存类，在 get_kernels 中实例化
	156	+ self.layout_C = RowMajor
	157	+ self.atlas_arch = atlas_arch
	158	+ self.alpha = alpha
	159	+ self.beta = beta
	160	+ # 如果 core_num 未提供，从 driver 获取设备属性
	161	+ if core_num is None:
	162	+ _override_hint = "override by passing 'core_num' to Gemm()"
	163	+ try:
	164	+ from triton.runtime.driver import driver
	165	+ device = driver.active.get_current_device()
	166	+ prop = driver.active.utils.get_device_properties(device)
	167	+ core_num = prop["num_aicore"]
	168	+ except ModuleNotFoundError:
	169	+ warnings.warn(
	170	+ "'triton' is not installed on your environment, cannot obtain driver info."
	171	+ f"core_num defaults to ({_CORE_NUM}). ({_override_hint})",
	172	+ RuntimeWarning,
	173	+ stacklevel=2
	174	+ )
	175	+ core_num = _CORE_NUM
	176	+ except Exception as e:
	177	+ warnings.warn(
	178	+ "An unexpected error occurred; "
	179	+ f"core_num defaults to ({_CORE_NUM}). ({_override_hint})\nError details: {e!r}",
	180	+ RuntimeWarning,
	181	+ stacklevel=2
	182	+ )
	183	+ core_num = _CORE_NUM
	184	+
	185	+ self.core_num = core_num
	186	+
	187	+ # 如果 evg_config 不为 None，必须包含 fn_src 和 example_inputs
	188	+ if evg_config is not None:
	189	+ if "fn_src" not in evg_config or "example_inputs" not in evg_config:
	190	+ raise ValueError("evg_config must contain 'fn_src' and 'example_inputs'")
	191	+ self.evg = evg_config
	192	+
	193	+ def can_implement(self) -> bool:
	194	+ return math.isclose(self.alpha, 1.0) and (math.isclose(self.beta, 0.0) or math.isclose(self.beta, 1.0))
	C C Chen_HaoWen7月10日 [正确性 / 一致性] `can_implement()` 允许 `beta=1.0`（返回 `True`），但 `get_kernels()` 中非 EVG 分支仅处理 `alpha=1.0 and beta=0.0` 的情况（见 L240）。当 `beta=1.0` 时： `can_implement()` 返回 `True` → 用户认为该配置可用 `get_kernels()` 走到 L262 的 `else` 分支 → 发出 `UserWarning` 并返回空列表这会导致用户困惑：`can_implement()` 说可以，但实际拿不到 kernel。建议统一两处的判断逻辑——要么 `can_implement()` 也限制 `beta` 仅允许 `0.0`，要么 `get_kernels()` 支持 `beta=1.0` 的场景。 I init__zhb__ 7月10日评论： 7月10日评论：
	195	+
	196	+ def get_kernels(self) -> List[GemmKernelBase]:
	197	+ # 必须使用 __init__ 中传入的 A 和 B
	198	+ if self.A is None or self.B is None:
	199	+ raise ValueError("A 和 B 必须在 Gemm.__init__ 中传入")
	200	+
	201	+ # 判断是否 batched 和提取 M, K, N（使用 __init__ 中保存的值）
	202	+ if self.M is None or self.K is None or self.N is None:
	203	+ raise ValueError("无法确定 M, K, N，请确保在 Gemm.__init__ 中传入了 A 和 B OpTensor")
	204	+
	205	+ # 如果 layout_C 还是类（未实例化），则实例化它
	206	+ if isinstance(self.layout_C, type) and issubclass(self.layout_C, Layout):
	207	+ self.layout_C = self.layout_C((self.M, self.N))
	208	+
	209	+ # 只提取 Bias 的 shape（用于后续判断）
	210	+ Bias_shape = self.Bias.shape if self.Bias is not None else None
	211	+
	212	+ params = {
	213	+ # 不传递 tensor 对象，只传递元数据（使用 OpTensor 时避免实例化）
	214	+ "element_accumulator": self.element_accumulator,
	215	+ "element_A": self.element_A,
	216	+ "element_B": self.element_B,
	217	+ "element_Bias": self.element_Bias,
	218	+ "element_C": self.element_C,
	219	+ "layout_A": self.layout_A,
	220	+ "layout_B": self.layout_B,
	221	+ "layout_Bias": self.layout_Bias,
	222	+ "layout_C": self.layout_C,
	223	+ "arch_tag": self.atlas_arch,
	224	+ "M": self.M,
	225	+ "K": self.K,
	226	+ "N": self.N,
	227	+ "batchCount": self.batch_count if self.is_batched else None,
	228	+ "evg": self.evg,
	229	+ }
	230	+
	231	+ if self.evg is not None:
	232	+ return [BasicMatmulTlaVisitorKernel(**params)]
	233	+
	234	+ if self.is_batched:
	235	+ # BatchedMatmul不支持Bias，带Bias的话返回空列表
	236	+ if self.element_Bias is not None:
	237	+ return []
	238	+ return [BatchedMatmulKernel(**params)]
	239	+ else:
	240	+ if math.isclose(self.alpha, 1.0) and math.isclose(self.beta, 0.0):
	241	+ if self.element_Bias is not None and Bias_shape is not None and len(Bias_shape) > 1:
	242	+ return []
	243	+ _threshold1 = 4096
	244	+ _threshold2 = 2048
	245	+ _default_ksplit_tile = (256, 256, 128, 32)
	246	+ res = []
	247	+ if self.K > _threshold2:
	248	+ # prefer k-split template
	249	+ num_task = math.ceil(self.M / _default_ksplit_tile[0]) * math.ceil(
	250	+ self.N / _default_ksplit_tile[1]
	251	+ )
	252	+ if num_task <= (0.5 * self.core_num):
	253	+ res.append(MultiCoreSplitkMatmulKernel(**params))
	254	+ elif num_task <= (0.9 * self.core_num) or (num_task % self.core_num) <= (0.9 * self.core_num):
	C C Chen_HaoWen7月10日 [正确性 / 逻辑缺陷] 条件 `(num_task % self.core_num) <= (0.9 * self.core_num)` 几乎总是为 `True`：取模结果 `num_task % core_num` 恒满足 `0 <= result < core_num` `0.9 * core_num` 接近 `core_num` 仅当余数恰好落在 `(0.9 * core_num, core_num)` 区间时才为 `False`，概率约 10% 这导致 `StreamkMatmulKernel` 在绝大多数 `K > 2048` 的场景下都会被添加，可能并非预期的 kernel 选择策略。请确认此条件的设计意图——如果是为了在 `num_task` 接近 `core_num` 的整数倍时优先使用 StreamK，建议改用更明确的条件，例如 `num_task % core_num == 0` 或 `abs(num_task % core_num - core_num) < threshold`。 I init__zhb__ 7月10日评论： 7月10日评论：
	255	+ res.append(StreamkMatmulKernel(**params))
	256	+ if num_task > self.core_num and num_task < (1.5 * self.core_num):
	257	+ res.append(TailMultiCoreSplitkMatmulKernel(**params))
	258	+ if self.K < _threshold1:
	259	+ res.append(BasicMatmulKernel(**params))
	260	+ return res
	261	+ else:
	262	+ warnings.warn(
	263	+ f"Only alpha=1.0 and beta=0.0 are supported for gemm, "
	264	+ f"got alpha={self.alpha}, beta={self.beta}; returning an empty kernel list.",
	265	+ UserWarning,
	266	+ stacklevel=2,
	267	+ )
	268	+ return []

Apython/catlass_cppgen/catlass_cppgen/op/group_gemm.py+316-0

@@ -0,0 +1,316 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+import warnings
	11	+from typing import List, Type, Optional, Dict, Any
	12	+import math
	13	+import numpy as np
	14	+
	15	+from catlass_cppgen.op.op import OperationBase
	16	+from catlass_cppgen.common.typing import SupportedDataType, SupportedTensor
	17	+from catlass_cppgen.common.op_tensor import OpTensor
	18	+from catlass_cppgen.common.data_type import DataType, get_default_accumulator
	19	+from catlass_cppgen.catlass.layout.layout import Layout, RowMajor
	20	+from catlass_cppgen.common.utils import extract_info
	21	+from catlass_cppgen.catlass.arch.arch import Arch
	22	+from catlass_cppgen.kernel.gemm.gemm_base import GemmKernelBase
	23	+from catlass_cppgen.kernel.group_gemm.grouped_matmul_slice_m import GroupedMatmulSliceMKernel
	24	+
	25	+_CORE_NUM = 8 # Default core num
	26	+
	27	+class GroupGemm(OperationBase):
	28	+ # 类属性类型注解
	29	+ A: Optional[OpTensor]
	30	+ B: Optional[OpTensor]
	31	+ Bias: Optional[OpTensor]
	32	+ C: Optional[OpTensor]
	33	+ M: Optional[int]
	34	+ K: Optional[int]
	35	+ N: Optional[int]
	36	+ problemCount: Optional[int]
	37	+ groupList: Optional[OpTensor]
	38	+
	39	+ def __init__(
	40	+ self,
	41	+ alpha: float = 1.0,
	42	+ beta: float = 0.0,
	43	+ element_accumulator: Optional[SupportedDataType] = None,
	44	+ element: Optional[SupportedDataType] = None,
	45	+ layout: Optional[Type[Layout]] = None,
	46	+ element_A: Optional[SupportedDataType] = None,
	47	+ element_B: Optional[SupportedDataType] = None,
	48	+ element_Bias: Optional[SupportedDataType] = None,
	49	+ element_C: Optional[SupportedDataType] = None,
	50	+ layout_A: Optional[Type[Layout]] = None,
	51	+ layout_B: Optional[Type[Layout]] = None,
	52	+ layout_Bias: Optional[Type[Layout]] = None,
	53	+ epilogue=None,
	54	+ evg_config: Optional[Dict[str, Any]] = None,
	55	+ atlas_arch: Optional[Arch] = None,
	56	+ core_num: Optional[int] = None,
	57	+ A: Optional[OpTensor] = None,
	58	+ B: Optional[OpTensor] = None,
	59	+ Bias: Optional[OpTensor] = None,
	60	+ C: Optional[OpTensor] = None,
	61	+ groupList: Optional[OpTensor] = None,
	62	+ ):
	63	+ # 保存 A、B、Bias、C 和 groupList OpTensor
	64	+ self.A = A
	65	+ self.B = B
	66	+ self.Bias = Bias
	67	+ self.C = C
	68	+ self.groupList = groupList
	69	+
	70	+ # 必须提供 groupList，从它提取 problemCount
	71	+ if groupList is None:
	72	+ raise ValueError("groupList must be provided")
	73	+
	74	+ groupList_shape, groupList_element, _, _ = extract_info(groupList, None, None)
	75	+ if groupList_shape is None:
	76	+ raise ValueError("groupList must have valid shape")
	77	+
	78	+ if len(groupList_shape) != 1:
	79	+ raise ValueError("groupList must be 1D tensor")
	80	+ # 从 groupList 的长度提取 problemCount
	81	+ self.problemCount = groupList_shape[0]
	82	+
	83	+ # 如果传入了 A 和 B OpTensor，从它们中提取信息
	84	+ if A is not None and B is not None:
	85	+ # 从 OpTensor 中提取 shape、element、layout 信息
	86	+ A_shape, element_A_from_tensor, layout_A_from_tensor, _ = extract_info(A, element_A or element, layout_A or layout)
	87	+ B_shape, element_B_from_tensor, layout_B_from_tensor, _ = extract_info(B, element_B or element, layout_B or layout)
	88	+
	89	+ # 确定最终使用的 element 和 layout（传入的参数优先，否则使用 OpTensor 中的信息）
	90	+ element_A = element_A_from_tensor or element_A or element
	91	+ element_B = element_B_from_tensor or element_B or element
	92	+ layout_A = layout_A_from_tensor or layout_A or layout
	93	+ layout_B = layout_B_from_tensor or layout_B or layout
	94	+
	95	+ if len(A_shape) != 2 or len(B_shape) != 3:
	96	+ raise ValueError("A must be 2D tensor (`m, k`) and B must be 3D tensor (`problemCount, k, n`) for group gemm")
	97	+
	98	+ # B 是 3D: [problemCount, k, n]
	99	+ # 验证第一维是否等于 problemCount
	100	+ if B_shape[0] != self.problemCount:
	101	+ raise ValueError(
	102	+ f"B's first dimension ({B_shape[0]}) must equal problemCount ({self.problemCount})"
	103	+ )
	104	+ # 提取 K, N 从最后两维
	105	+ B_k = B_shape[1]
	106	+ B_n = B_shape[2]
	107	+
	108	+ # 验证 A 的 K 维和 B 的 K 维是否匹配
	109	+ if A_shape[1] != B_k:
	110	+ raise ValueError(
	111	+ f"A's K dimension ({A_shape[1]}) must match B's K dimension ({B_k})"
	112	+ )
	113	+ # 提取 M, K, N
	114	+ self.M = A_shape[0]
	115	+ self.K = A_shape[1] # 使用 A 的 K
	116	+ self.N = B_n # 每个 group 的 N 维度
	117	+ else:
	118	+ if element is None and not all([element_A, element_B, element_C]):
	119	+ raise ValueError(
	120	+ "must provide 'element', or specify element_A, element_B, element_C separately"
	121	+ )
	122	+ if layout is None and not all([layout_A, layout_B]):
	123	+ raise ValueError(
	124	+ "must provide 'layout', or specify layout_A, layout_B separately"
	125	+ )
	126	+
	127	+ element_A = element_A or element
	128	+ element_B = element_B or element
	129	+ layout_A = layout_A or layout
	130	+ layout_B = layout_B or layout
	131	+
	132	+ self.M = None
	133	+ self.K = None
	134	+ self.N = None
	135	+
	136	+ # 如果传入了 Bias OpTensor，从它中提取信息
	137	+ if Bias is not None:
	138	+ _, element_Bias_from_tensor, layout_Bias_from_tensor, _ = extract_info(Bias, element_Bias, layout_Bias)
	139	+ element_Bias = element_Bias_from_tensor or element_Bias
	140	+ layout_Bias = layout_Bias_from_tensor or layout_Bias
	141	+
	142	+ # 如果传入了 C OpTensor，从它中提取信息
	143	+ if C is not None:
	144	+ _, element_C_from_tensor, layout_C_from_tensor, _ = extract_info(C, element_C or element, None)
	145	+ element_C = element_C_from_tensor or element_C or element
	146	+
	147	+ self.element_A = element_A
	148	+ self.element_B = element_B
	149	+ self.element_Bias = element_Bias
	150	+ self.element_C = element_C or element
	151	+ if element_accumulator is None and not all([self.element_A, self.element_B]):
	152	+ raise ValueError("element_accumulator must be provided, or element_A, element_B should be given both so that accumulator type can be auto-derived")
	153	+ self.element_accumulator = element_accumulator or get_default_accumulator(
	154	+ self.element_A, self.element_B
	155	+ )
	156	+ self.layout_A = layout_A
	157	+ self.layout_B = layout_B
	158	+ # layout_Bias不提供则为 None
	159	+ self.layout_Bias = layout_Bias
	160	+ # layout_C 固定为 RowMajor，如果 M 和 N 已确定则实例化
	161	+ if self.M is not None and self.N is not None:
	162	+ self.layout_C = RowMajor((self.M, self.N))
	163	+ else:
	164	+ # 如果 M 和 N 未确定，先保存类，在 get_kernels 中实例化
	165	+ self.layout_C = RowMajor
	166	+ self.atlas_arch = atlas_arch
	167	+ self.alpha = alpha
	168	+ self.beta = beta
	169	+ # 如果 core_num 未提供，从 driver 获取设备属性
	170	+ if core_num is None:
	171	+ _override_hint = "override by passing 'core_num' to GroupGemm()"
	172	+ try:
	173	+ from triton.runtime.driver import driver
	174	+ device = driver.active.get_current_device()
	175	+ prop = driver.active.utils.get_device_properties(device)
	176	+ core_num = prop["num_aicore"]
	177	+ except ModuleNotFoundError:
	178	+ warnings.warn(
	179	+ "'triton' is not installed on your environment, cannot obtain driver info."
	180	+ f"core_num defaults to ({_CORE_NUM}). ({_override_hint})",
	181	+ RuntimeWarning,
	182	+ stacklevel=2
	183	+ )
	184	+ core_num = _CORE_NUM
	185	+ except Exception as e:
	186	+ warnings.warn(
	187	+ "An unexpected error occurred; "
	188	+ f"core_num defaults to ({_CORE_NUM}). ({_override_hint})\nError details: {e!r}",
	189	+ RuntimeWarning,
	190	+ stacklevel=2
	191	+ )
	192	+ core_num = _CORE_NUM
	193	+
	194	+ self.core_num = core_num
	195	+
	196	+ # 如果 evg_config 不为 None，必须包含 fn_src 和 example_inputs
	197	+ if evg_config is not None:
	198	+ if "fn_src" not in evg_config or "example_inputs" not in evg_config:
	199	+ raise ValueError("evg_config must contain 'fn_src' and 'example_inputs'")
	200	+ self.evg = evg_config
	201	+
	202	+ def can_implement(self) -> bool:
	203	+ return math.isclose(self.alpha, 1.0) and (math.isclose(self.beta, 0.0) or math.isclose(self.beta, 1.0))
	204	+
	205	+ def get_kernels(
	206	+ self,
	207	+ A: Optional[SupportedTensor] = None,
	208	+ B: Optional[SupportedTensor] = None,
	209	+ Bias: Optional[SupportedTensor] = None,
	210	+ C: Optional[SupportedTensor] = None,
	211	+ groupList: Optional[SupportedTensor] = None,
	212	+ ) -> List[GemmKernelBase]:
	213	+ # 处理 groupList：优先使用传入的参数，否则使用 __init__ 中保存的值
	214	+ final_groupList = groupList if groupList is not None else self.groupList
	215	+
	216	+ # 必须提供 groupList，从它提取 problemCount
	217	+ if final_groupList is None:
	218	+ raise ValueError("groupList must be provided, either in __init__ or get_kernels")
	219	+
	220	+ groupList_shape, groupList_element, _, _ = extract_info(final_groupList, None, None)
	221	+ if groupList_shape is None:
	222	+ raise ValueError("groupList must have valid shape")
	223	+
	224	+ if len(groupList_shape) != 1:
	225	+ raise ValueError("groupList must be 1D tensor")
	226	+ # 从 groupList 的长度提取 problemCount
	227	+ final_problemCount = groupList_shape[0]
	228	+
	229	+ # 优先使用传入的参数（向后兼容），否则使用 __init__ 中保存的值
	230	+ use_init_values = (A is None and B is None) and (self.A is not None and self.B is not None)
	231	+
	232	+ if use_init_values:
	233	+ # 使用 __init__ 中保存的值
	234	+ if self.M is None or self.K is None or self.N is None:
	235	+ raise ValueError("无法确定 M, K, N，请确保在 GroupGemm.__init__ 中传入了 A 和 B OpTensor")
	236	+
	237	+ # 如果 layout_C 还是类（未实例化），则实例化它
	238	+ if isinstance(self.layout_C, type) and issubclass(self.layout_C, Layout):
	239	+ self.layout_C = self.layout_C((self.M, self.N))
	240	+
	241	+ element_A = self.element_A
	242	+ element_B = self.element_B
	243	+ element_C = self.element_C
	244	+ element_Bias = self.element_Bias
	245	+
	246	+ layout_A = self.layout_A
	247	+ layout_B = self.layout_B
	248	+ layout_C = self.layout_C
	249	+ layout_Bias = self.layout_Bias
	250	+
	251	+ M = self.M
	252	+ K = self.K
	253	+ N = self.N
	254	+
	255	+ else:
	256	+ # 处理输入：如果传入 OpTensor，直接使用其信息；如果是 torch.Tensor/np.ndarray，提取信息
	257	+ # 使用 OpTensor 时，不需要实例化实际的 tensor 数据
	258	+ if A is None or B is None:
	259	+ raise ValueError("A 和 B 必须提供，可以通过 GroupGemm.__init__ 或 get_kernels 参数传入")
	260	+
	261	+ A_shape, element_A, layout_A, A_tensor = extract_info(A, self.element_A, self.layout_A)
	262	+ B_shape, element_B, layout_B, B_tensor = extract_info(B, self.element_B, self.layout_B)
	263	+ Bias_shape, element_Bias, layout_Bias, Bias_tensor = extract_info(Bias, self.element_Bias, self.layout_Bias)
	264	+ C_shape, element_C, layout_C, C_tensor = extract_info(C, self.element_C, self.layout_C)
	265	+
	266	+ element_A = element_A or self.element_A
	267	+ element_B = element_B or self.element_B
	268	+ element_C = element_C or self.element_C
	269	+ element_Bias = element_Bias or self.element_Bias
	270	+
	271	+ layout_A = layout_A or self.layout_A
	272	+ layout_B = layout_B or self.layout_B
	273	+ layout_C = layout_C or self.layout_C
	274	+ layout_Bias = layout_Bias or self.layout_Bias
	275	+
	276	+ if len(A_shape) != 2 or len(B_shape) != 3:
	277	+ raise ValueError("A must be 2D tensor (`m, k`) and B must be 3D tensor (`problemCount, k, n`) for group gemm")
	278	+
	279	+ # 提取 M, K, N
	280	+ M = A_shape[0]
	281	+ K = A_shape[1]
	282	+ N = B_shape[2]
	283	+
	284	+ # 如果 layout_C 还是类（未实例化），则实例化它
	285	+ if isinstance(layout_C, type) and issubclass(layout_C, Layout):
	286	+ layout_C = layout_C((M, N))
	287	+
	288	+ params = {
	289	+ # 不传递 tensor 对象，只传递元数据（使用 OpTensor 时避免实例化）
	290	+ "element_accumulator": self.element_accumulator,
	291	+ "element_A": element_A,
	292	+ "element_B": element_B,
	293	+ "element_Bias": element_Bias,
	294	+ "element_C": element_C,
	295	+ "layout_A": layout_A,
	296	+ "layout_B": layout_B,
	297	+ "layout_Bias": layout_Bias,
	298	+ "layout_C": layout_C,
	299	+ "arch_tag": self.atlas_arch,
	300	+ "M": M,
	301	+ "K": K,
	302	+ "N": N,
	303	+ "problemCount": final_problemCount,
	304	+ "groupList_element": groupList_element,
	305	+ }
	306	+
	307	+ if math.isclose(self.alpha, 1.0) and math.isclose(self.beta, 0.0):
	308	+ return [GroupedMatmulSliceMKernel(**params)]
	309	+ else:
	310	+ warnings.warn(
	311	+ f"Only alpha=1.0 and beta=0.0 are supported for grouped gemm, "
	312	+ f"got alpha={self.alpha}, beta={self.beta}; returning an empty kernel list.",
	313	+ UserWarning,
	314	+ stacklevel=2,
	315	+ )
	316	+ return []

Apython/catlass_cppgen/catlass_cppgen/op/op.py+26-0

@@ -0,0 +1,26 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+from abc import abstractmethod
	11	+
	12	+from typing import List
	13	+from catlass_cppgen.kernel.kernel_base import KernelBase
	14	+
	15	+
	16	+class OperationBase:
	17	+ @abstractmethod
	18	+ def get_kernels(self, args, *kwargs) -> List[KernelBase]:
	19	+ pass
	20	+
	21	+ def get_best_kernel(self) -> KernelBase:
	22	+ return self.get_kernels()[0]
	23	+
	24	+ @abstractmethod
	25	+ def can_implement(self)->bool:
	26	+ pass

Apython/catlass_cppgen/docs/evg_api.md+103-0

@@ -0,0 +1,103 @@
	1	+# EVG API (Python)
	2	+
	3	+本节记录了 CATLASS 中 [EVG（Epilogue Visitor Graph）模块](../../../docs/zh/2_Design/03_evg/01_evg_design.md)的 Python API。EVG 将 Python 描述的 epilogue 后处理函数解析为 DAG，并生成对应的 C++ Visitor 代码。
	4	+
	5	+核心入口函数定义见 [`evg_extension.py`](../catlass_cppgen/catlass/evg_extension.py)，节点与定义见 [`evg/`](../catlass_cppgen/catlass/evg/)。
	6	+
	7	+---
	8	+
	9	+## `evg()` — 入口函数
	10	+
	11	+解析 epilogue Python 函数，生成 EVG 定义字符串与参数。
	12	+
	13	+```python
	14	+def evg(
	15	+ fn_src: str,
	16	+ example_inputs: Dict[str, OpTensor],
	17	+) -> Tuple[str, str, str, EVGArgRenames]:
	18	+```
	19	+
	20	+\| 参数 \| 说明 \|
	21	+\|------\|------\|
	22	+\| `fn_src` \| epilogue 函数的 Python 源码字符串，函数固定命名为 `epilogue`，最后一个 `return` 语句的返回值作为输出 \|
	23	+\| `example_inputs` \| 输入/输出张量的元数据字典，key 为变量名，value 为 `OpTensor` \|
	24	+
	25	+返回值 `(callback_name, evg_args, evg_str, arg_renames)`：
	26	+
	27	+\| 字段 \| 类型 \| 说明 \|
	28	+\|------\|------\|------\|
	29	+\| `callback_name` \| `str` \| 生成的 epilogue callback 名称（固定为 `"EVGResult"`） \|
	30	+\| `evg_args` \| `str` \| EVG 参数结构体声明 C++ 代码，包含 `Arguments` 和 `computeLength` \|
	31	+\| `evg_str` \| `str` \| EVG Visitor 类型定义 C++ 代码，包含 `VisitorAccLoad`、`VisitorAuxLoad`、`VisitorCompute`、`TreeVisitor`、`TopologicalVisitor` 等 \|
	32	+\| `arg_renames` \| `EVGArgRenames` \| 参数重命名信息（当前版本暂返回 `None`） \|
	33	+
	34	+### 使用示例
	35	+
	36	+```python
	37	+from catlass_cppgen.catlass.evg_extension import evg
	38	+from catlass_cppgen.common.data_type import DataType
	39	+from catlass_cppgen.common.op_tensor import OpTensor
	40	+
	41	+fn_src = """
	42	+def epilogue(accum, bias):
	43	+ result = accum + bias
	44	+ return result
	45	+"""
	46	+
	47	+inputs = {
	48	+ "accum": OpTensor.from_shape_stride((128, 256), (256, 1), DataType.FLOAT),
	49	+ "bias": OpTensor.from_shape_stride((1, 256), (256, 1), DataType.FLOAT),
	50	+ "result": OpTensor.from_shape_stride((128, 256), (256, 1), DataType.FLOAT),
	51	+}
	52	+
	53	+callback_name, evg_args, evg_str, _ = evg(fn_src=fn_src, example_inputs=inputs)
	54	+print(callback_name) # "EVGResult"
	55	+print(evg_args) # typename EVGResult::Arguments evg_args{...};
	56	+print(evg_str) # using Result = Catlass::Epilogue::Fusion::VisitorAuStore<...>; ...
	57	+```
	58	+
	59	+### 支持的 Epilogue 算子
	60	+
	61	+在 `fn_src` 的 `epilogue` 函数中，可使用以下算子：
	62	+
	63	+\| 类别 \| 写法 \| 说明 \|
	64	+\|------\|------\|------\|
	65	+\| add \| `accum + bias` \| 二元加法 \|
	66	+\| sub \| `accum - bias` \| 二元减法 \|
	67	+\| mul \| `accum * scale` \| 二元乘法 \|
	68	+\| div \| `accum / scale` \| 二元除法 \|
	69	+\| relu \| `relu(accum)` \| ReLU 激活 \|
	70	+\| leakyRelu \| `leakyRelu(accum, alpha)` \| LeakyReLU 激活 \|
	71	+\| Prelu \| `Prelu(accum, weight)` \| PReLU 激活 \|
	72	+\| sigmoid \| `sigmoid(accum)` \| Sigmoid 激活 \|
	73	+\| silu \| `silu(accum)` \| SiLU 激活 \|
	74	+\| maximum \| `maximum(a, b)` \| 逐元素取最大值 \|
	75	+\| minimum \| `minimum(a, b)` \| 逐元素取最小值 \|
	76	+\| cast \| `cast(accum, "float16", "float")` \| 类型转换（参数：目标类型, 源类型 [, RoundMode]） \|
	77	+\| constant \| `constant(1.0, "float")` \| 创建常量值 \|
	78	+
	79	+算子可串联组合，例如：`relu(accum) + bias` 会生成包含 `VisitorCompute(Relu)` 和 `VisitorCompute(Add)` 的 TreeVisitor 链。
	80	+
	81	+广播：当前支持 `RowBroadcast`（行广播）。输入张量形状为 `(1, N)` 时，自动匹配 `(M, N)` 累加器形状并标记为 `RowBroadcast`。`ColumnBroadcast` 暂不支持。
	82	+
	83	+---
	84	+
	85	+## EVG 使用参考
	86	+
	87	+在`catlass_cppgen`中，要使用 EVG 特性，有下述三种办法：
	88	+ - 在创建kernel时直接通过 `Gemm(evg_config=...)` 传入，参考示例：
	89	+```python
	90	+evg_config = {
	91	+ "fn_src": "def epilogue(accum, bias):\n return relu(accum + bias)",
	92	+ "example_inputs": {"accum": ..., "bias": ..., "result": ...},
	93	+}
	94	+gemm = Gemm(..., evg_config=evg_config, A=a, B=b)
	95	+```
	96	+ - 对已有 Kernel 直接调用 `to_evg(evg_config)`，参考示例：
	97	+```python
	98	+gemm = Gemm(...)
	99	+
	100	+gemm = gemm.to_evg(evg_config=evg_config)
	101	+```
	102	+ - 直接调用`evg(fn_src, example_inputs)` 生成 EVG 定义，如上述[使用示例](###使用示例)所示。
	103	+

Apython/catlass_cppgen/docs/kernel_api.md+480-0

@@ -0,0 +1,480 @@
	1	+# Kernel 基类 API (Python)
	2	+
	3	+本节记录了 CATLASS 中 kernel 基类及其代码生成接口的 Python API，接口定义见 [`kernel_base.py`](../catlass_cppgen/kernel/kernel_base.py)。
	4	+
	5	+---
	6	+
	7	+## `KernelBase` 类
	8	+
	9	+该类是所有 kernel 定义以及代码生成的基类。
	10	+
	11	+### 主要方法
	12	+
	13	+- `get_default_tile_shape(self) -> Tuple[GemmShape, GemmShape]`
	14	+ 抽象方法。返回 kernel 的默认 L1 和 L0 tile 形状。
	15	+
	16	+- `get_workspace_size(self) -> int`
	17	+ 返回 kernel 所需的工作区大小（默认 0）。
	18	+
	19	+- `need_workspace(self) -> bool`
	20	+ 是否需要外部工作区（默认为 `False`）。
	21	+
	22	+- `get_core_num(self) -> int`
	23	+ 返回 kernel 使用的核心数（默认 0）。
	24	+
	25	+- `get_default_dispatch_policy_list(self) -> List[DispatchPolicy]`
	26	+ 获取默认的 dispatch_policy 列表。子类可以重写此方法以定义自己的默认 dispatch_policy 列表。如果子类不重写，默认返回空列表。
	27	+
	28	+#### 调优接口
	29	+
	30	+- `set_l1_tile_shape(self, l1_tile_shape: GemmShape)`
	31	+ 设置 L1 tile 形状。
	32	+
	33	+- `set_l0_tile_shape(self, l0_tile_shape: GemmShape)`
	34	+ 设置 L0 tile 形状。
	35	+
	36	+- `set_dispatch_policy(self, dispatch_policy: Union[DispatchPolicy, List[DispatchPolicy]])`
	37	+ 设置 dispatch policy。可以传入单个 policy 或 policy 列表。如果传入单个 policy，会自动转换为包含该 policy 的列表。
	38	+
	39	+- `get_dispatch_policy(self) -> List[DispatchPolicy]`
	40	+ 获取 dispatch policy 列表。每个算子类（如 `BasicMatmulKernel`）可以定义自己的默认 dispatch_policy 列表，通过重写 `get_default_dispatch_policy_list()` 方法。
	41	+
	42	+ 如果未显式设置，会自动使用算子类定义的默认列表。列表的第一个元素 `[0]` 是默认策略。
	43	+ 目前默认返回只包含一个 policy 的列表，但接口设计支持未来扩展为多个 policy。
	44	+
	45	+- `set_block_scheduler(self, block_scheduler)`
	46	+ 设置 block scheduler（保留参数）。
	47	+
	48	+- `tune(self, l1_tile_shape: Optional[GemmShape] = None, l0_tile_shape: Optional[GemmShape] = None, dispatch_policy: Optional[Union[DispatchPolicy, List[DispatchPolicy]]] = None, block_scheduler: Optional[BlockScheduler] = None)`
	49	+ 一次性调优 kernel 的 tile 形状和调度策略。所有参数均为可选，如果为 `None` 则使用当前值或默认值。
	50	+
	51	+#### 特性支持查询
	52	+
	53	+- `is_support(self, feature: str) -> bool`
	54	+ 查询该 kernel 是否支持某特性。
	55	+
	56	+- `is_support_evg(self) -> bool`
	57	+ 查询该 kernel 是否支持 `evg` 特性。
	58	+
	59	+- `is_support_hf32(self) -> bool`
	60	+ 查询该 kernel 是否支持 `hf32` 特性。
	61	+
	62	+#### 代码生成与渲染
	63	+
	64	+- `get_render_params(self) -> Dict[str, Any]`
	65	+ 抽象方法。返回用于模板渲染的参数字典。
	66	+
	67	+- `gen_includes(self) -> str`
	68	+ 生成 C++ 的 `#include` 文件头，根据 kernel 需求自动拉齐。
	69	+
	70	+- `gen_kernel_name(self) -> str`
	71	+ 根据参数生成 kernel 的函数名。
	72	+
	73	+- `gen_params_device(self, def_mode: bool = False) -> str`
	74	+ 生成 device 端函数参数列表，`def_mode=True` 用于函数定义，否则用于调用。
	75	+
	76	+- `gen_kernel_template(self) -> str`
	77	+ 使用渲染参数填充 kernel 模板，生成核函数"核心代码块"。
	78	+
	79	+- `gen_layout_template(self) -> str`
	80	+ 生成 layout 相关信息代码，包括 M, K, N 的定义和 layout tag 的创建。
	81	+
	82	+- `codegen(self) -> str`
	83	+ 生成完整 C++ kernel，包括头文件、函数签名、参数、代码块与 kernel 启动。
	84	+
	85	+---
	86	+
	87	+## 示例：生成 Kernel 代码
	88	+
	89	+### 简单用法（GEMM kernel 示例）
	90	+
	91	+```python
	92	+from catlass_cppgen.kernel.gemm.basic_matmul import BasicMatmulKernel
	93	+from catlass_cppgen.kernel.gemm.gemm_base import GemmKernelBase
	94	+from catlass_cppgen.catlass.gemm_coord import GemmShape
	95	+from catlass_cppgen.common.data_type import DataType
	96	+from catlass_cppgen.catlass.layout.layout import RowMajor
	97	+from catlass_cppgen.catlass.arch.arch import Arch
	98	+
	99	+# 假设 BasicMatmulKernel 是 GemmKernelBase 的子类，已实现必备方法
	100	+gemm_kernel = BasicMatmulKernel(
	101	+ element_accumulator=DataType.FLOAT,
	102	+ element_A=DataType.FLOAT,
	103	+ element_B=DataType.FLOAT,
	104	+ element_C=DataType.FLOAT,
	105	+ element_Bias=DataType.FLOAT,
	106	+ layout_A=RowMajor((128, 256)),
	107	+ layout_B=RowMajor((256, 384)),
	108	+ layout_Bias=RowMajor((128, 384)),
	109	+ arch_tag=Arch.Ascend950
	110	+)
	111	+
	112	+# 优化 tile 形状
	113	+gemm_kernel.tune(
	114	+ l1_tile_shape=GemmShape(128, 256, 64),
	115	+ l0_tile_shape=GemmShape(128, 256, 64)
	116	+)
	117	+
	118	+# 生成头文件
	119	+print("Includes:")
	120	+print(gemm_kernel.gen_includes())
	121	+
	122	+# 生成核函数参数（定义模式）
	123	+print("Params (def_mode=True):")
	124	+print(gemm_kernel.gen_params_device(def_mode=True))
	125	+
	126	+# 生成核函数参数（调用模式）
	127	+print("Params (def_mode=False):")
	128	+print(gemm_kernel.gen_params_device(def_mode=False))
	129	+
	130	+# 生成核函数模板
	131	+print("Kernel Template:")
	132	+print(gemm_kernel.gen_kernel_template())
	133	+
	134	+# 生成 layout 模板
	135	+print("Layout Template:")
	136	+print(gemm_kernel.gen_layout_template())
	137	+```
	138	+
	139	+### 所有 Kernel 类型的使用示例
	140	+
	141	+#### 1. BasicMatmulKernel（基础矩阵乘法）
	142	+
	143	+适用于标准的 2D 矩阵乘法运算，支持可选的 Bias 参数。
	144	+
	145	+```python
	146	+from catlass_cppgen.op.gemm import Gemm
	147	+from catlass_cppgen.common.op_tensor import OpTensor
	148	+from catlass_cppgen.common.data_type import DataType
	149	+from catlass_cppgen.catlass.layout.layout import RowMajor
	150	+from catlass_cppgen.catlass.gemm_coord import GemmShape
	151	+from catlass_cppgen.catlass.arch.arch import Arch
	152	+from catlass_cppgen.catlass.gemm.dispatch_policy import MmadPingpong
	153	+
	154	+def test_basic_matmul_kernel():
	155	+ # 使用 OpTensor.from_shape_stride 创建输入（避免实例化实际 tensor 数据）
	156	+ a = OpTensor.from_shape_stride(
	157	+ shape=(128, 256),
	158	+ stride=(256, 1), # RowMajor stride: (n, 1)
	159	+ dtype=DataType.FLOAT
	160	+ )
	161	+ b = OpTensor.from_shape_stride(
	162	+ shape=(256, 384),
	163	+ stride=(384, 1), # RowMajor stride: (n, 1)
	164	+ dtype=DataType.FLOAT
	165	+ )
	166	+ gemm_plan = Gemm(atlas_arch=Arch.Ascend950, element=DataType.FLOAT, layout=RowMajor)
	167	+ kernels = gemm_plan.get_kernels(A=a, B=b)
	168	+ basic_kernel = kernels[0] # BasicMatmulKernel
	169	+
	170	+ basic_kernel.tune(
	171	+ GemmShape(128, 256, 64),
	172	+ GemmShape(128, 256, 64),
	173	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950)
	174	+ )
	175	+
	176	+ # 生成头文件
	177	+ print("Includes:")
	178	+ print(basic_kernel.gen_includes())
	179	+
	180	+ # 生成核函数参数（定义模式）
	181	+ print("Params (def_mode=True):")
	182	+ print(basic_kernel.gen_params_device(def_mode=True))
	183	+
	184	+ # 生成核函数模板
	185	+ print("Kernel Template:")
	186	+ print(basic_kernel.gen_kernel_template())
	187	+
	188	+ # 生成 layout 模板
	189	+ print("Layout Template:")
	190	+ print(basic_kernel.gen_layout_template())
	191	+```
	192	+
	193	+#### 2. BatchedMatmulKernel（批处理矩阵乘法）
	194	+
	195	+适用于批处理场景，输入张量 A 和 B 为 3 维（batchCount, M, K）和（batchCount, K, N）。
	196	+
	197	+```python
	198	+from catlass_cppgen.op.gemm import Gemm
	199	+from catlass_cppgen.common.op_tensor import OpTensor
	200	+from catlass_cppgen.common.data_type import DataType
	201	+from catlass_cppgen.catlass.layout.layout import RowMajor
	202	+from catlass_cppgen.catlass.gemm_coord import GemmShape
	203	+from catlass_cppgen.catlass.arch.arch import Arch
	204	+from catlass_cppgen.catlass.gemm.dispatch_policy import MmadPingpong
	205	+
	206	+def test_batched_matmul_kernel():
	207	+ # 批处理：batchCount=8, M=128, K=256, N=384
	208	+ # 使用 OpTensor.from_shape_stride 创建输入（避免实例化实际 tensor 数据）
	209	+ a = OpTensor.from_shape_stride(
	210	+ shape=(8, 128, 256), # (batchCount, M, K)
	211	+ stride=(32768, 256, 1), # batched RowMajor stride: (m*n, n, 1)
	212	+ dtype=DataType.FLOAT
	213	+ )
	214	+ b = OpTensor.from_shape_stride(
	215	+ shape=(8, 256, 384), # (batchCount, K, N)
	216	+ stride=(98304, 384, 1), # batched RowMajor stride: (k*n, n, 1)
	217	+ dtype=DataType.FLOAT
	218	+ )
	219	+ gemm_plan = Gemm(atlas_arch=Arch.Ascend950, element=DataType.FLOAT, layout=RowMajor)
	220	+ kernels = gemm_plan.get_kernels(A=a, B=b)
	221	+ batched_kernel = kernels[0] # BatchedMatmulKernel
	222	+
	223	+ batched_kernel.tune(
	224	+ GemmShape(128, 256, 64),
	225	+ GemmShape(128, 256, 64),
	226	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950, enable_unit_flag=True)
	227	+ )
	228	+
	229	+ # 生成头文件
	230	+ print("Includes:")
	231	+ print(batched_kernel.gen_includes())
	232	+
	233	+ # 生成核函数参数（定义模式）
	234	+ print("Params (def_mode=True):")
	235	+ print(batched_kernel.gen_params_device(def_mode=True))
	236	+
	237	+ # 生成核函数模板
	238	+ print("Kernel Template:")
	239	+ print(batched_kernel.gen_kernel_template())
	240	+
	241	+ # 生成 layout 模板（包含 stride 信息）
	242	+ print("Layout Template:")
	243	+ print(batched_kernel.gen_layout_template())
	244	+```
	245	+
	246	+#### 3. StreamkMatmulKernel（StreamK 矩阵乘法）
	247	+
	248	+适用于大规模矩阵乘法，使用 StreamK 调度策略优化性能。
	249	+
	250	+```python
	251	+from catlass_cppgen.op.gemm import Gemm
	252	+from catlass_cppgen.common.op_tensor import OpTensor
	253	+from catlass_cppgen.common.data_type import DataType
	254	+from catlass_cppgen.catlass.layout.layout import RowMajor
	255	+from catlass_cppgen.catlass.gemm_coord import GemmShape
	256	+from catlass_cppgen.catlass.arch.arch import Arch
	257	+from catlass_cppgen.catlass.gemm.dispatch_policy import MmadPingpong
	258	+
	259	+def test_streamk_matmul_kernel():
	260	+ # 使用 OpTensor.from_shape_stride 创建输入（避免实例化实际 tensor 数据）
	261	+ a = OpTensor.from_shape_stride(
	262	+ shape=(128, 256),
	263	+ stride=(256, 1), # RowMajor stride: (n, 1)
	264	+ dtype=DataType.FLOAT
	265	+ )
	266	+ b = OpTensor.from_shape_stride(
	267	+ shape=(256, 384),
	268	+ stride=(384, 1), # RowMajor stride: (n, 1)
	269	+ dtype=DataType.FLOAT
	270	+ )
	271	+ gemm_plan = Gemm(atlas_arch=Arch.Ascend950, element=DataType.FLOAT, layout=RowMajor)
	272	+ kernels = gemm_plan.get_kernels(A=a, B=b)
	273	+ streamk_kernel = kernels[2] # StreamkMatmulKernel（索引可能因实现而异）
	274	+
	275	+ # StreamK kernel 默认使用较大的 tile shape，并指定 dispatch_policy
	276	+ streamk_kernel.tune(
	277	+ GemmShape(256, 256, 128),
	278	+ GemmShape(256, 256, 32),
	279	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950, enable_unit_flag=True)
	280	+ )
	281	+
	282	+ # 生成头文件
	283	+ print("Includes:")
	284	+ print(streamk_kernel.gen_includes())
	285	+
	286	+ # 生成核函数参数（定义模式，包含 aicCoreNum 参数）
	287	+ print("Params (def_mode=True):")
	288	+ print(streamk_kernel.gen_params_device(def_mode=True))
	289	+
	290	+ # 生成核函数模板
	291	+ print("Kernel Template:")
	292	+ print(streamk_kernel.gen_kernel_template())
	293	+
	294	+ # 生成 layout 模板
	295	+ print("Layout Template:")
	296	+ print(streamk_kernel.gen_layout_template())
	297	+```
	298	+
	299	+#### 4. MultiCoreSplitkMatmulKernel（多核 SplitK 矩阵乘法）
	300	+
	301	+适用于需要多核并行计算的大规模矩阵乘法，通过 SplitK 策略提高并行度。
	302	+
	303	+```python
	304	+from catlass_cppgen.op.gemm import Gemm
	305	+from catlass_cppgen.common.op_tensor import OpTensor
	306	+from catlass_cppgen.common.data_type import DataType
	307	+from catlass_cppgen.catlass.layout.layout import RowMajor
	308	+from catlass_cppgen.catlass.gemm_coord import GemmShape
	309	+from catlass_cppgen.catlass.arch.arch import Arch
	310	+from catlass_cppgen.catlass.gemm.dispatch_policy import MmadPingpong
	311	+
	312	+def test_multi_core_splitk_matmul_kernel():
	313	+ # 使用 OpTensor.from_shape_stride 创建输入（避免实例化实际 tensor 数据）
	314	+ a = OpTensor.from_shape_stride(
	315	+ shape=(128, 256),
	316	+ stride=(256, 1), # RowMajor stride: (n, 1)
	317	+ dtype=DataType.FLOAT
	318	+ )
	319	+ b = OpTensor.from_shape_stride(
	320	+ shape=(256, 384),
	321	+ stride=(384, 1), # RowMajor stride: (n, 1)
	322	+ dtype=DataType.FLOAT
	323	+ )
	324	+ gemm_plan = Gemm(atlas_arch=Arch.Ascend950, element=DataType.FLOAT, layout=RowMajor)
	325	+ kernels = gemm_plan.get_kernels(A=a, B=b)
	326	+ splitk_kernel = kernels[1] # MultiCoreSplitkMatmulKernel（索引可能因实现而异）
	327	+
	328	+ # SplitK kernel 默认使用较大的 tile shape，并指定 dispatch_policy
	329	+ splitk_kernel.tune(
	330	+ GemmShape(256, 256, 128),
	331	+ GemmShape(256, 256, 32),
	332	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950, enable_unit_flag=True)
	333	+ )
	334	+
	335	+ # 生成头文件
	336	+ print("Includes:")
	337	+ print(splitk_kernel.gen_includes())
	338	+
	339	+ # 生成核函数参数（定义模式，包含 aicCoreNum 参数）
	340	+ print("Params (def_mode=True):")
	341	+ print(splitk_kernel.gen_params_device(def_mode=True))
	342	+
	343	+ # 生成核函数模板
	344	+ print("Kernel Template:")
	345	+ print(splitk_kernel.gen_kernel_template())
	346	+
	347	+ # 生成 layout 模板
	348	+ print("Layout Template:")
	349	+ print(splitk_kernel.gen_layout_template())
	350	+```
	351	+
	352	+#### 5. TailMultiCoreSplitkMatmulKernel（尾部多核 SplitK 矩阵乘法）
	353	+
	354	+适用于处理 SplitK 策略中的尾部计算，与 MultiCoreSplitkMatmulKernel 配合使用。
	355	+
	356	+```python
	357	+from catlass_cppgen.op.gemm import Gemm
	358	+from catlass_cppgen.common.op_tensor import OpTensor
	359	+from catlass_cppgen.common.data_type import DataType
	360	+from catlass_cppgen.catlass.layout.layout import RowMajor
	361	+from catlass_cppgen.catlass.gemm_coord import GemmShape
	362	+from catlass_cppgen.catlass.arch.arch import Arch
	363	+from catlass_cppgen.catlass.gemm.dispatch_policy import MmadPingpong
	364	+
	365	+def test_tail_multi_core_splitk_matmul_kernel():
	366	+ # 使用 OpTensor.from_shape_stride 创建输入（避免实例化实际 tensor 数据）
	367	+ a = OpTensor.from_shape_stride(
	368	+ shape=(128, 256),
	369	+ stride=(256, 1), # RowMajor stride: (n, 1)
	370	+ dtype=DataType.FLOAT
	371	+ )
	372	+ b = OpTensor.from_shape_stride(
	373	+ shape=(256, 384),
	374	+ stride=(384, 1), # RowMajor stride: (n, 1)
	375	+ dtype=DataType.FLOAT
	376	+ )
	377	+ gemm_plan = Gemm(atlas_arch=Arch.Ascend950, element=DataType.FLOAT, layout=RowMajor)
	378	+ kernels = gemm_plan.get_kernels(A=a, B=b)
	379	+ tail_splitk_kernel = kernels[3] # TailMultiCoreSplitkMatmulKernel（索引可能因实现而异）
	380	+
	381	+ # Tail SplitK kernel 默认使用较大的 tile shape，并指定 dispatch_policy
	382	+ tail_splitk_kernel.tune(
	383	+ GemmShape(256, 256, 128),
	384	+ GemmShape(256, 256, 32),
	385	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950, enable_unit_flag=True)
	386	+ )
	387	+
	388	+ # 生成头文件
	389	+ print("Includes:")
	390	+ print(tail_splitk_kernel.gen_includes())
	391	+
	392	+ # 生成核函数参数（定义模式，包含 aicCoreNum 参数）
	393	+ print("Params (def_mode=True):")
	394	+ print(tail_splitk_kernel.gen_params_device(def_mode=True))
	395	+
	396	+ # 生成核函数模板
	397	+ print("Kernel Template:")
	398	+ print(tail_splitk_kernel.gen_kernel_template())
	399	+
	400	+ # 生成 layout 模板
	401	+ print("Layout Template:")
	402	+ print(tail_splitk_kernel.gen_layout_template())
	403	+```
	404	+
	405	+#### 6. GroupedMatmulSliceMKernel（分组矩阵乘法 - Slice M）
	406	+
	407	+适用于分组 GEMM 场景，多个不同大小的矩阵乘法问题，通过 Slice M 策略优化。
	408	+
	409	+```python
	410	+from catlass_cppgen.op.group_gemm import GroupGemm
	411	+from catlass_cppgen.common.op_tensor import OpTensor
	412	+from catlass_cppgen.common.data_type import DataType
	413	+from catlass_cppgen.catlass.layout.layout import RowMajor
	414	+from catlass_cppgen.catlass.gemm_coord import GemmShape
	415	+from catlass_cppgen.catlass.arch.arch import Arch
	416	+from catlass_cppgen.catlass.gemm.dispatch_policy import MmadPingpong
	417	+
	418	+def test_grouped_matmul_slice_m_kernel():
	419	+ # 分组 GEMM：每个问题的 A 和 B 矩阵维度相同
	420	+ # 使用 OpTensor.from_shape_stride 创建输入（避免实例化实际 tensor 数据）
	421	+ from catlass_cppgen.catlass.layout.layout import VectorLayout
	422	+
	423	+ a = OpTensor.from_shape_stride(
	424	+ shape=(128, 256), # 单个问题的 A 矩阵
	425	+ stride=(256, 1), # RowMajor stride: (n, 1)
	426	+ dtype=DataType.FLOAT
	427	+ )
	428	+ b = OpTensor.from_shape_stride(
	429	+ shape=(256, 384), # 单个问题的 B 矩阵
	430	+ stride=(384, 1), # RowMajor stride: (n, 1)
	431	+ dtype=DataType.FLOAT
	432	+ )
	433	+ # 创建 groupList OpTensor（一维，int64_t 类型）
	434	+ # groupList 的长度即为 problemCount（这里是 4）
	435	+ groupList = OpTensor(
	436	+ dtype=DataType.INT64,
	437	+ layout=VectorLayout(4), # 4 个 group
	438	+ shape=(4,)
	439	+ )
	440	+
	441	+ group_gemm_plan = GroupGemm(atlas_arch=Arch.Ascend950, element=DataType.FLOAT, layout=RowMajor)
	442	+ kernels = group_gemm_plan.get_kernels(A=a, B=b, groupList=groupList)
	443	+ grouped_kernel = kernels[0] # GroupedMatmulSliceMKernel
	444	+
	445	+ # Grouped kernel 默认使用较大的 tile shape，并指定 dispatch_policy
	446	+ grouped_kernel.tune(
	447	+ GemmShape(256, 256, 256),
	448	+ GemmShape(256, 256, 64),
	449	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950, enable_unit_flag=True)
	450	+ )
	451	+
	452	+ # 生成头文件
	453	+ print("Includes:")
	454	+ print(grouped_kernel.gen_includes())
	455	+
	456	+ # 生成核函数参数（定义模式，包含 problemCount 和 deviceGroupList 参数）
	457	+ print("Params (def_mode=True):")
	458	+ print(grouped_kernel.gen_params_device(def_mode=True))
	459	+
	460	+ # 生成核函数模板
	461	+ print("Kernel Template:")
	462	+ print(grouped_kernel.gen_kernel_template())
	463	+
	464	+ # 生成 layout 模板
	465	+ print("Layout Template:")
	466	+ print(grouped_kernel.gen_layout_template())
	467	+```
	468	+
	469	+---
	470	+
	471	+## 说明
	472	+
	473	+- 如需拓展 kernel 代码生成，继承 `KernelBase` 并实现 `get_render_params` 和 `get_default_tile_shape`。
	474	+- 子类需设置模板元字段（如 `_INCLUDES`, `_PARAMS_DEVICE`, `_KERNEL_TEMPLATE`）以适配专用 kernel。
	475	+- `codegen()` 方法整合所有渲染环节，一键输出最终 C++ kernel 源码。
	476	+
	477	+更多细节与二次开发见 [`kernel_base.py`](../catlass_cppgen/kernel/kernel_base.py)。
	478	+
	479	+
	480	+

Apython/catlass_cppgen/docs/optensor_api.md+517-0

@@ -0,0 +1,517 @@
	1	+# OpTensor 基础 Api (Python)
	2	+
	3	+本文档详细描述了 `catlass_cppgen` 中所有支持的`OpTensor`创建方式，及关联的基础特性，如数据类型，布局分布等。
	4	+
	5	+---
	6	+
	7	+## 目录
	8	+
	9	+- [OpTensor 构造方法](#optensor-构造方法)
	10	+ - [方式1: 直接构造函数](#方式1-直接构造函数)
	11	+ - [方式2: from_shape_stride 类方法](#方式2-from_shape_stride-类方法)
	12	+ - [方式3: from_tensor 类方法](#方式3-from_tensor-类方法)
	13	+- [Operation 支持的输入类型](#operation-支持的输入类型)
	14	+- [数据类型支持](#数据类型支持)
	15	+- [布局类型支持](#布局类型支持)
	16	+- [使用示例](#使用示例)
	17	+
	18	+---
	19	+
	20	+## OpTensor 构造方法
	21	+
	22	+`OpTensor` 是操作（Operation）的输入输出 tensor 的抽象，提供了统一的 tensor 表示。它可以从多种方式创建，支持避免实例化实际 tensor 数据，从而提高代码生成效率。
	23	+
	24	+### 方式1: 直接构造函数
	25	+
	26	+方法签名：
	27	+```python
	28	+OpTensor(
	29	+ dtype: DataType,
	30	+ layout: Layout,
	31	+ shape: Optional[tuple[int, ...]] = None,
	32	+ data_ptr: Optional[ctypes.c_void_p] = None
	33	+)
	34	+```
	35	+
	36	+参数说明：
	37	+- `dtype`: tensor 的数据类型（`DataType` 枚举）
	38	+- `layout`: tensor 的布局（`Layout` 对象，如 `RowMajor`、`ColumnMajor`）
	39	+- `shape`: 可选的完整形状。如果不提供，则使用 `layout.shape`。对于 batched tensor，需要显式指定包含 batch 维度的完整 shape
	40	+- `data_ptr`: 可选的数据指针（`ctypes.c_void_p`），用于实际数据访问
	41	+
	42	+示例：
	43	+
	44	+```python
	45	+from catlass_cppgen.common.op_tensor import OpTensor
	46	+from catlass_cppgen.common.data_type import DataType
	47	+from catlass_cppgen.catlass.layout.layout import RowMajor, ColumnMajor
	48	+import ctypes
	49	+
	50	+# 1a. 基本创建（只指定 dtype 和 layout）
	51	+a_op = OpTensor(dtype=DataType.FLOAT, layout=RowMajor((128, 256)))
	52	+b_op = OpTensor(dtype=DataType.FLOAT, layout=RowMajor((256, 384)))
	53	+print(f"A: shape={a_op.shape}, dtype={a_op.dtype}")
	54	+# 输出: A: shape=(128, 256), dtype=DataType.FLOAT
	55	+
	56	+# 1b. 指定 shape（用于 batched tensor）
	57	+a_op_batched = OpTensor(
	58	+ dtype=DataType.FLOAT,
	59	+ layout=RowMajor((128, 256)), # 内层矩阵的 layout
	60	+ shape=(8, 128, 256) # 完整的 shape，包含 batch 维度
	61	+)
	62	+print(f"A (batched): shape={a_op_batched.shape}")
	63	+# 输出: A (batched): shape=(8, 128, 256)
	64	+
	65	+# 1c. 指定 data_ptr（用于实际数据指针）
	66	+data_ptr = ctypes.c_void_p(0x12345678)
	67	+a_op_with_ptr = OpTensor(
	68	+ dtype=DataType.FLOAT,
	69	+ layout=RowMajor((128, 256)),
	70	+ data_ptr=data_ptr
	71	+)
	72	+print(f"A (with data_ptr): data_ptr={a_op_with_ptr.data_ptr}")
	73	+
	74	+# 1d. 使用 ColumnMajor layout
	75	+a_op_col = OpTensor(dtype=DataType.FLOAT, layout=ColumnMajor((128, 256)))
	76	+print(f"A (ColumnMajor): stride={a_op_col.stride}")
	77	+# 输出: A (ColumnMajor): stride=(1, 128)
	78	+
	79	+# 1e. 使用不同的数据类型
	80	+a_op_fp16 = OpTensor(dtype=DataType.FLOAT16, layout=RowMajor((128, 256)))
	81	+print(f"A (FLOAT16): dtype={a_op_fp16.dtype}")
	82	+```
	83	+
	84	+---
	85	+
	86	+### 方式2: from_shape_stride 类方法
	87	+
	88	+方法签名：
	89	+```python
	90	+@classmethod
	91	+OpTensor.from_shape_stride(
	92	+ cls,
	93	+ shape: tuple[int, ...],
	94	+ stride: tuple[int, ...],
	95	+ dtype: DataType
	96	+) -> "OpTensor"
	97	+```
	98	+
	99	+参数说明：
	100	+- `shape`: tensor 的形状（可以是 2D 或 3D，支持 batched）
	101	+- `stride`: tensor 的步长
	102	+- `dtype`: tensor 的数据类型
	103	+
	104	+说明：
	105	+- 该方法会自动从 `shape` 和 `stride` 推断 `Layout` 类型
	106	+- 对于 batched tensor（3D），只推断内层矩阵的布局
	107	+- 适合在已知 shape 和 stride 但不想实例化实际 tensor 的场景使用
	108	+
	109	+示例：
	110	+
	111	+```python
	112	+from catlass_cppgen.common.op_tensor import OpTensor
	113	+from catlass_cppgen.common.data_type import DataType
	114	+
	115	+# 2a. 2D tensor，RowMajor stride
	116	+a_op = OpTensor.from_shape_stride(
	117	+ shape=(128, 256),
	118	+ stride=(256, 1), # RowMajor stride: (n, 1)
	119	+ dtype=DataType.FLOAT
	120	+)
	121	+print(f"A: shape={a_op.shape}, stride={a_op.stride}, layout={a_op.layout}")
	122	+# 输出: A: shape=(128, 256), stride=(256, 1), layout=RowMajor((128, 256))
	123	+
	124	+# 2b. 2D tensor，ColumnMajor stride
	125	+a_op = OpTensor.from_shape_stride(
	126	+ shape=(128, 256),
	127	+ stride=(1, 128), # ColumnMajor stride: (1, m)
	128	+ dtype=DataType.FLOAT
	129	+)
	130	+print(f"A: shape={a_op.shape}, stride={a_op.stride}, layout={a_op.layout}")
	131	+# 输出: A: shape=(128, 256), stride=(1, 128), layout=ColumnMajor((128, 256))
	132	+
	133	+# 2c. 3D batched tensor
	134	+a_op = OpTensor.from_shape_stride(
	135	+ shape=(8, 128, 256), # (batch, m, n)
	136	+ stride=(32768, 256, 1), # batched RowMajor stride
	137	+ dtype=DataType.FLOAT
	138	+)
	139	+print(f"A (batched): shape={a_op.shape}, layout={a_op.layout}")
	140	+# 输出: A (batched): shape=(8, 128, 256), layout=RowMajor((128, 256))
	141	+```
	142	+
	143	+---
	144	+
	145	+### 方式3: from_tensor 类方法
	146	+
	147	+方法签名：
	148	+```python
	149	+@classmethod
	150	+OpTensor.from_tensor(
	151	+ cls,
	152	+ tensor: SupportedTensor, # torch.Tensor 或 np.ndarray
	153	+ layout: Optional[Layout] = None,
	154	+ dtype: Optional[DataType] = None
	155	+) -> "OpTensor"
	156	+```
	157	+
	158	+参数说明：
	159	+- `tensor`: `torch.Tensor` 或 `np.ndarray` 对象
	160	+- `layout`: 可选的 Layout，如果不提供则从 tensor 的 stride 自动推断
	161	+- `dtype`: 可选的 DataType，如果不提供则从 tensor 的 dtype 自动推断
	162	+
	163	+说明：
	164	+- 该方法会从实际的 tensor 对象中提取 shape、stride、dtype 等信息
	165	+- 如果提供了 `layout` 或 `dtype`，会覆盖从 tensor 推断的值
	166	+- 适合从已有的 tensor 对象创建 `OpTensor`
	167	+
	168	+示例：
	169	+
	170	+```python
	171	+import torch
	172	+import numpy as np
	173	+from catlass_cppgen.common.op_tensor import OpTensor
	174	+from catlass_cppgen.common.data_type import DataType
	175	+from catlass_cppgen.catlass.layout.layout import RowMajor, ColumnMajor
	176	+
	177	+# 3a. 从 torch.Tensor 创建（自动推断所有信息）
	178	+torch_a = torch.ones(128, 256, dtype=torch.float32)
	179	+a_op = OpTensor.from_tensor(torch_a)
	180	+print(f"A: shape={a_op.shape}, dtype={a_op.dtype}, stride={a_op.stride}")
	181	+# 输出: A: shape=(128, 256), dtype=DataType.FLOAT, stride=(256, 1)
	182	+
	183	+# 3b. 从 torch.Tensor 创建，指定 layout
	184	+torch_a = torch.ones(128, 256, dtype=torch.float32)
	185	+a_op = OpTensor.from_tensor(torch_a, layout=RowMajor((128, 256)))
	186	+print(f"A: shape={a_op.shape}, layout={a_op.layout}")
	187	+
	188	+# 3c. 从 torch.Tensor 创建，指定 dtype
	189	+torch_a = torch.ones(128, 256, dtype=torch.float32)
	190	+a_op = OpTensor.from_tensor(torch_a, dtype=DataType.FLOAT16)
	191	+print(f"A: shape={a_op.shape}, dtype={a_op.dtype}")
	192	+
	193	+# 3d. 从 torch.Tensor 创建，同时指定 layout 和 dtype
	194	+torch_a = torch.ones(128, 256, dtype=torch.float32)
	195	+a_op = OpTensor.from_tensor(
	196	+ torch_a,
	197	+ layout=ColumnMajor((128, 256)),
	198	+ dtype=DataType.FLOAT16
	199	+)
	200	+print(f"A: shape={a_op.shape}, dtype={a_op.dtype}, layout={a_op.layout}")
	201	+
	202	+# 3e. 从 np.ndarray 创建
	203	+np_a = np.ones((128, 256), dtype=np.float32)
	204	+a_op = OpTensor.from_tensor(np_a)
	205	+print(f"A: shape={a_op.shape}, dtype={a_op.dtype}, stride={a_op.stride}")
	206	+
	207	+# 3f. 从 batched torch.Tensor 创建
	208	+torch_a_batched = torch.ones(8, 128, 256, dtype=torch.float32)
	209	+a_op = OpTensor.from_tensor(torch_a_batched)
	210	+print(f"A (batched): shape={a_op.shape}, layout={a_op.layout}")
	211	+# 输出: A (batched): shape=(8, 128, 256), layout=RowMajor((128, 256))
	212	+```
	213	+
	214	+---
	215	+
	216	+## Operation 支持的输入类型
	217	+
	218	+Operation（如 `Gemm`、`GroupGemm`）的 `get_kernels()` 方法支持以下输入类型：
	219	+
	220	+### 支持的输入类型
	221	+
	222	+1. OpTensor（推荐）
	223	+ - 使用 `OpTensor` 可以避免实例化实际的 tensor 数据
	224	+ - 适合在代码生成阶段使用，只需要 tensor 的元数据（shape、dtype、layout）
	225	+
	226	+2. torch.Tensor（向后兼容）
	227	+ - 直接传入 `torch.Tensor` 对象
	228	+ - 系统会自动提取信息，但需要实际的数据对象
	229	+
	230	+3. np.ndarray（向后兼容）
	231	+ - 直接传入 `numpy.ndarray` 对象
	232	+ - 系统会自动提取信息，但需要实际的数据对象
	233	+
	234	+4. None
	235	+ - 传入 `None` 表示使用 Operation 初始化时指定的默认值
	236	+
	237	+### 示例
	238	+
	239	+```python
	240	+from catlass_cppgen.op.gemm import Gemm
	241	+from catlass_cppgen.common.op_tensor import OpTensor
	242	+from catlass_cppgen.common.data_type import DataType
	243	+from catlass_cppgen.catlass.layout.layout import RowMajor
	244	+from catlass_cppgen.catlass.arch.arch import Arch
	245	+import torch
	246	+import numpy as np
	247	+
	248	+gemm_plan = Gemm(atlas_arch=Arch.Ascend950, element=DataType.FLOAT, layout=RowMajor)
	249	+
	250	+# 方式1: 使用 OpTensor（推荐）
	251	+a_op = OpTensor(dtype=DataType.FLOAT, layout=RowMajor((128, 256)))
	252	+b_op = OpTensor(dtype=DataType.FLOAT, layout=RowMajor((256, 384)))
	253	+kernels = gemm_plan.get_kernels(A=a_op, B=b_op)
	254	+print(f"获取到 {len(kernels)} 个 kernels")
	255	+
	256	+# 方式2: 使用 torch.Tensor（向后兼容）
	257	+torch_a = torch.ones(128, 256, dtype=torch.float32)
	258	+torch_b = torch.ones(256, 384, dtype=torch.float32)
	259	+kernels = gemm_plan.get_kernels(A=torch_a, B=torch_b)
	260	+print(f"获取到 {len(kernels)} 个 kernels")
	261	+
	262	+# 方式3: 使用 np.ndarray（向后兼容）
	263	+np_a = np.ones((128, 256), dtype=np.float32)
	264	+np_b = np.ones((256, 384), dtype=np.float32)
	265	+kernels = gemm_plan.get_kernels(A=np_a, B=np_b)
	266	+print(f"获取到 {len(kernels)} 个 kernels")
	267	+
	268	+# 方式4: 混合使用（部分使用 OpTensor，部分使用 None）
	269	+a_op = OpTensor(dtype=DataType.FLOAT, layout=RowMajor((128, 256)))
	270	+kernels = gemm_plan.get_kernels(A=a_op, B=None) # B 使用默认值
	271	+```
	272	+
	273	+---
	274	+
	275	+## 数据类型支持
	276	+
	277	+### 数据类型转换
	278	+
	279	+从 `torch.Tensor` 或 `np.ndarray` 创建 `OpTensor` 时，系统会自动将 tensor 的 dtype 转换为对应的 `DataType`：
	280	+
	281	+```python
	282	+import torch
	283	+from catlass_cppgen.common.op_tensor import OpTensor
	284	+
	285	+# torch.float32 -> DataType.FLOAT
	286	+torch_a = torch.ones(128, 256, dtype=torch.float32)
	287	+a_op = OpTensor.from_tensor(torch_a)
	288	+print(a_op.dtype) # DataType.FLOAT
	289	+
	290	+# torch.float16 -> DataType.FLOAT16
	291	+torch_b = torch.ones(128, 256, dtype=torch.float16)
	292	+b_op = OpTensor.from_tensor(torch_b)
	293	+print(b_op.dtype) # DataType.FLOAT16
	294	+```
	295	+
	296	+---
	297	+
	298	+## 布局类型支持
	299	+
	300	+`OpTensor` 支持以下布局类型（通过 `Layout` 类及其子类）：
	301	+
	302	+### 基本布局
	303	+
	304	+- RowMajor - 行主序布局（C 风格）
	305	+ - 2D: `RowMajor((m, n))`，stride = `(n, 1)`
	306	+ - 示例: `RowMajor((128, 256))`
	307	+
	308	+- ColumnMajor - 列主序布局（Fortran 风格）
	309	+ - 2D: `ColumnMajor((m, n))`，stride = `(1, m)`
	310	+ - 示例: `ColumnMajor((128, 256))`
	311	+
	312	+- VectorLayout - 向量布局
	313	+ - 1D: `VectorLayout(n)`
	314	+ - 示例: `VectorLayout(256)`
	315	+
	316	+### 布局推断
	317	+
	318	+当使用 `from_shape_stride` 或 `from_tensor` 时，系统会自动从 stride 推断布局：
	319	+
	320	+```python
	321	+from catlass_cppgen.common.op_tensor import OpTensor
	322	+from catlass_cppgen.common.data_type import DataType
	323	+
	324	+# 自动推断为 RowMajor
	325	+a_op = OpTensor.from_shape_stride(
	326	+ shape=(128, 256),
	327	+ stride=(256, 1), # RowMajor stride
	328	+ dtype=DataType.FLOAT
	329	+)
	330	+print(a_op.layout) # RowMajor((128, 256))
	331	+
	332	+# 自动推断为 ColumnMajor
	333	+b_op = OpTensor.from_shape_stride(
	334	+ shape=(128, 256),
	335	+ stride=(1, 128), # ColumnMajor stride
	336	+ dtype=DataType.FLOAT
	337	+)
	338	+print(b_op.layout) # ColumnMajor((128, 256))
	339	+```
	340	+
	341	+---
	342	+
	343	+## 使用示例
	344	+
	345	+### 示例1: 基本 GEMM 操作
	346	+
	347	+```python
	348	+from catlass_cppgen.op.gemm import Gemm
	349	+from catlass_cppgen.common.op_tensor import OpTensor
	350	+from catlass_cppgen.common.data_type import DataType
	351	+from catlass_cppgen.catlass.layout.layout import RowMajor
	352	+from catlass_cppgen.catlass.arch.arch import Arch
	353	+
	354	+# 创建 GEMM 计划
	355	+gemm_plan = Gemm(
	356	+ atlas_arch=Arch.Ascend950,
	357	+ element=DataType.FLOAT,
	358	+ layout=RowMajor
	359	+)
	360	+
	361	+# 使用 OpTensor 创建输入
	362	+a_op = OpTensor(dtype=DataType.FLOAT, layout=RowMajor((128, 256)))
	363	+b_op = OpTensor(dtype=DataType.FLOAT, layout=RowMajor((256, 384)))
	364	+
	365	+# 获取 kernels
	366	+kernels = gemm_plan.get_kernels(A=a_op, B=b_op)
	367	+print(f"获取到 {len(kernels)} 个 kernels")
	368	+```
	369	+
	370	+### 示例2: Batched GEMM 操作
	371	+
	372	+```python
	373	+from catlass_cppgen.op.gemm import Gemm
	374	+from catlass_cppgen.common.op_tensor import OpTensor
	375	+from catlass_cppgen.common.data_type import DataType
	376	+from catlass_cppgen.catlass.layout.layout import RowMajor
	377	+from catlass_cppgen.catlass.arch.arch import Arch
	378	+
	379	+gemm_plan = Gemm(atlas_arch=Arch.Ascend950, element=DataType.FLOAT, layout=RowMajor)
	380	+
	381	+# 方式1: 直接构造函数，指定 batched shape
	382	+batch_count = 8
	383	+a_op = OpTensor(
	384	+ dtype=DataType.FLOAT,
	385	+ layout=RowMajor((128, 256)), # 内层矩阵的 layout
	386	+ shape=(batch_count, 128, 256) # 完整的 shape，包含 batch 维度
	387	+)
	388	+b_op = OpTensor(
	389	+ dtype=DataType.FLOAT,
	390	+ layout=RowMajor((256, 384)),
	391	+ shape=(batch_count, 256, 384)
	392	+)
	393	+
	394	+kernels = gemm_plan.get_kernels(A=a_op, B=b_op)
	395	+print(f"获取到 {len(kernels)} 个 kernels")
	396	+if len(kernels) > 0:
	397	+ batched_kernel = kernels[0] # BatchedMatmulKernel
	398	+ print(f"Kernel 类型: {type(batched_kernel).__name__}")
	399	+ print(f"BatchCount: {batched_kernel.batchCount}")
	400	+
	401	+# 方式2: 使用 from_shape_stride
	402	+a_op = OpTensor.from_shape_stride(
	403	+ shape=(batch_count, 128, 256),
	404	+ stride=(32768, 256, 1), # batched RowMajor stride
	405	+ dtype=DataType.FLOAT
	406	+)
	407	+b_op = OpTensor.from_shape_stride(
	408	+ shape=(batch_count, 256, 384),
	409	+ stride=(98304, 384, 1),
	410	+ dtype=DataType.FLOAT
	411	+)
	412	+kernels = gemm_plan.get_kernels(A=a_op, B=b_op)
	413	+```
	414	+
	415	+### 示例3: GroupGemm 操作
	416	+
	417	+```python
	418	+from catlass_cppgen.op.group_gemm import GroupGemm
	419	+from catlass_cppgen.common.op_tensor import OpTensor
	420	+from catlass_cppgen.common.data_type import DataType
	421	+from catlass_cppgen.catlass.layout.layout import RowMajor
	422	+from catlass_cppgen.catlass.arch.arch import Arch
	423	+
	424	+group_gemm_plan = GroupGemm(
	425	+ atlas_arch=Arch.Ascend950,
	426	+ element=DataType.FLOAT,
	427	+ layout=RowMajor
	428	+)
	429	+problem_count = 4 # 4 个不同的 GEMM 问题
	430	+
	431	+# 使用 OpTensor 创建（2D tensor，每个问题的 A 和 B 矩阵维度相同）
	432	+a_op = OpTensor(dtype=DataType.FLOAT, layout=RowMajor((128, 256)))
	433	+b_op = OpTensor(dtype=DataType.FLOAT, layout=RowMajor((256, 384)))
	434	+
	435	+kernels = group_gemm_plan.get_kernels(
	436	+ A=a_op,
	437	+ B=b_op,
	438	+ problemCount=problem_count
	439	+)
	440	+print(f"获取到 {len(kernels)} 个 kernels")
	441	+if len(kernels) > 0:
	442	+ grouped_kernel = kernels[0] # GroupedMatmulSliceMKernel
	443	+ print(f"Kernel 类型: {type(grouped_kernel).__name__}")
	444	+ print(f"ProblemCount: {grouped_kernel.problemCount}")
	445	+```
	446	+
	447	+### 示例4: 从实际 tensor 创建
	448	+
	449	+```python
	450	+import torch
	451	+from catlass_cppgen.op.gemm import Gemm
	452	+from catlass_cppgen.common.op_tensor import OpTensor
	453	+from catlass_cppgen.common.data_type import DataType
	454	+from catlass_cppgen.catlass.layout.layout import RowMajor
	455	+from catlass_cppgen.catlass.arch.arch import Arch
	456	+
	457	+gemm_plan = Gemm(atlas_arch=Arch.Ascend950, element=DataType.FLOAT, layout=RowMajor)
	458	+
	459	+# 从 torch.Tensor 创建 OpTensor
	460	+torch_a = torch.ones(128, 256, dtype=torch.float32)
	461	+torch_b = torch.ones(256, 384, dtype=torch.float32)
	462	+
	463	+a_op = OpTensor.from_tensor(torch_a)
	464	+b_op = OpTensor.from_tensor(torch_b)
	465	+
	466	+kernels = gemm_plan.get_kernels(A=a_op, B=b_op)
	467	+print(f"获取到 {len(kernels)} 个 kernels")
	468	+```
	469	+
	470	+---
	471	+
	472	+### 4. 数据类型和布局的一致性
	473	+
	474	+确保 Operation 初始化时指定的数据类型和布局与输入 tensor 一致，或者让系统自动推断：
	475	+
	476	+```python
	477	+# 方式1: Operation 指定默认值，输入使用 None
	478	+gemm_plan = Gemm(
	479	+ atlas_arch=Arch.Ascend950,
	480	+ element=DataType.FLOAT,
	481	+ layout=RowMajor
	482	+)
	483	+kernels = gemm_plan.get_kernels(A=None, B=None) # 使用默认值
	484	+
	485	+# 方式2: 输入显式指定，覆盖默认值
	486	+a_op = OpTensor(dtype=DataType.FLOAT16, layout=ColumnMajor((128, 256)))
	487	+kernels = gemm_plan.get_kernels(A=a_op, B=b_op) # 使用 a_op 指定的值
	488	+```
	489	+
	490	+---
	491	+
	492	+## 总结
	493	+
	494	+本文档介绍了 catlass_cppgen 中所有支持的输入方式：
	495	+
	496	+1. OpTensor 的三种构造方法：
	497	+ - 直接构造函数：`OpTensor(dtype, layout, shape=None, data_ptr=None)`
	498	+ - `from_shape_stride`：从 shape 和 stride 创建，自动推断 layout
	499	+ - `from_tensor`：从 `torch.Tensor` 或 `np.ndarray` 创建
	500	+
	501	+2. Operation 支持的输入类型：
	502	+ - `OpTensor`（推荐）
	503	+ - `torch.Tensor`（向后兼容）
	504	+ - `np.ndarray`（向后兼容）
	505	+ - `None`（使用默认值）
	506	+
	507	+3. 数据类型和布局：
	508	+ - 支持多种数据类型（FLOAT、FLOAT16、INT8 等）
	509	+ - 支持多种布局（RowMajor、ColumnMajor、VectorLayout 等）
	510	+
	511	+4. 最佳实践：
	512	+ - 优先使用 `OpTensor` 避免实例化数据
	513	+ - 根据场景选择合适的构造方法
	514	+ - 正确处理 batched tensor 的维度
	515	+
	516	+通过合理使用这些输入方式，可以高效地进行代码生成和 kernel 调优。
	517	+

Apython/catlass_cppgen/pyproject.toml+54-0

@@ -0,0 +1,54 @@
	1	+[build-system]
	2	+requires = ["setuptools>=61.0", "wheel"]
	3	+build-backend = "setuptools.build_meta"
	4	+
	5	+[project]
	6	+name = "catlass-cppgen"
	7	+dynamic = ["version"]
	8	+description = "通过pylib框架构建CATLASS算子"
	9	+readme = "README.md"
	10	+requires-python = ">=3.8"
	11	+dependencies = [
	12	+ "ml-dtypes",
	13	+ "numpy<2",
	14	+ "torch",
	15	+]
	16	+
	17	+[tool.setuptools]
	18	+packages = [
	19	+ "catlass_cppgen",
	20	+ "catlass_cppgen.catlass",
	21	+ "catlass_cppgen.catlass.arch",
	22	+ "catlass_cppgen.catlass.evg",
	23	+ "catlass_cppgen.catlass.gemm",
	24	+ "catlass_cppgen.catlass.layout",
	25	+ "catlass_cppgen.common",
	26	+ "catlass_cppgen.kernel",
	27	+ "catlass_cppgen.kernel.gemm",
	28	+ "catlass_cppgen.kernel.group_gemm",
	29	+ "catlass_cppgen.op",
	30	+]
	31	+
	32	+[tool.setuptools.dynamic]
	33	+version = {attr = "catlass_cppgen._version.__version__"}
	34	+
	35	+[tool.setuptools.package-data]
	36	+catlass_cppgen = ["py.typed"]
	37	+
	38	+[tool.uv.sources]
	39	+torch = [
	40	+ { index = "pytorch-cpu" },
	41	+]
	42	+torchvision = [
	43	+ { index = "pytorch-cpu" },
	44	+]
	45	+
	46	+[[tool.uv.index]]
	47	+url="https://pypi.tuna.tsinghua.edu.cn/simple"
	48	+default=true
	49	+
	50	+[[tool.uv.index]]
	51	+name = "pytorch-cpu"
	52	+url = "https://download.pytorch.org/whl/cpu"
	53	+explicit = true
	54	+

Apython/catlass_cppgen/tests/catlass/test_arch.py+34-0

@@ -0,0 +1,34 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+import unittest
	11	+from catlass_cppgen.catlass.arch.arch import Arch
	12	+
	13	+
	14	+class TestArch(unittest.TestCase):
	15	+ def test_arch_atlas_a2(self):
	16	+ arch = Arch.AtlasA2
	17	+ self.assertEqual(arch.value, "Arch::AtlasA2")
	18	+
	19	+ def test_arch_ascend950(self):
	20	+ arch = Arch.Ascend950
	21	+ self.assertEqual(arch.value, "Arch::Ascend950")
	22	+
	23	+ def test_arch_enum_values(self):
	24	+ """测试所有 Arch 枚举值"""
	25	+ expected_values = {
	26	+ Arch.AtlasA2: "Arch::AtlasA2",
	27	+ Arch.Ascend950: "Arch::Ascend950",
	28	+ }
	29	+ for arch, expected_value in expected_values.items():
	30	+ self.assertEqual(arch.value, expected_value)
	31	+
	32	+
	33	+if __name__ == "__main__":
	34	+ unittest.main()

Apython/catlass_cppgen/tests/catlass/test_dispatch_policy.py+97-0

@@ -0,0 +1,97 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+import unittest
	CheaterAbec7月6日这里面的测试建议全量删除 I init__zhb__ 7月8日评论： 7月8日评论：
	11	+from catlass_cppgen.catlass.gemm.dispatch_policy import (
	12	+ MmadAtlasA2Pingpong,
	13	+ MmadPingpong,
	14	+ MmadPreloadAsyncWithCallback,
	15	+ MmadMultiBatch,
	16	+)
	17	+from catlass_cppgen.catlass.arch.arch import Arch
	18	+
	19	+
	20	+class TestMmadAtlasA2Pingpong(unittest.TestCase):
	21	+ def test_mmad_atlas_a2_pingpong_default(self):
	22	+ mmad = MmadAtlasA2Pingpong()
	23	+ self.assertEqual(mmad.stages, 2)
	24	+ self.assertFalse(mmad.enable_unit_flag)
	25	+
	26	+ def test_mmad_atlas_a2_pingpong_with_flag(self):
	27	+ mmad = MmadAtlasA2Pingpong(enable_unit_flag=True)
	28	+ self.assertEqual(mmad.stages, 2)
	29	+ self.assertTrue(mmad.enable_unit_flag)
	30	+
	31	+ def test_mmad_atlas_a2_pingpong_to_cpp(self):
	32	+ mmad = MmadAtlasA2Pingpong(enable_unit_flag=True)
	33	+ cpp_str = mmad.to_cpp()
	34	+ self.assertIn("Gemm::MmadAtlasA2Pingpong", cpp_str)
	35	+ self.assertIn("true", cpp_str)
	36	+
	37	+
	38	+class TestMmadPingpong(unittest.TestCase):
	39	+ def test_mmad_pingpong(self):
	40	+ mmad = MmadPingpong(
	41	+ arch_tag=Arch.Ascend950,
	42	+ enable_unit_flag=True,
	43	+ use_hf32_mode=True,
	44	+ l0c_stages=2,
	45	+ enable_l1_resident=True
	46	+ )
	47	+ self.assertEqual(mmad.arch_tag, Arch.Ascend950)
	48	+ self.assertFalse(mmad.async_)
	49	+ self.assertEqual(mmad.stages, 2)
	50	+ self.assertTrue(mmad.enable_unit_flag)
	51	+ self.assertTrue(mmad.use_hf32_mode)
	52	+ self.assertEqual(mmad.l0c_stages, 2)
	53	+ self.assertTrue(mmad.enable_l1_resident)
	54	+
	55	+ def test_mmad_pingpong_to_cpp(self):
	56	+ mmad = MmadPingpong(arch_tag=Arch.Ascend950, enable_unit_flag=True)
	57	+ cpp_str = mmad.to_cpp()
	58	+ self.assertIn("Gemm::MmadPingpong", cpp_str)
	59	+ self.assertIn("Arch::Ascend950", cpp_str)
	60	+
	61	+
	62	+class TestMmadPreloadAsyncWithCallback(unittest.TestCase):
	63	+ def test_mmad_preload_async_with_callback(self):
	64	+ mmad = MmadPreloadAsyncWithCallback(
	65	+ arch_tag=Arch.Ascend950,
	66	+ preload_stages=3,
	67	+ l1a_stages=2,
	68	+ l1b_stages=2,
	69	+ l0a_stages=1,
	70	+ l0b_stages=1,
	71	+ l0c_stages=1,
	72	+ enable_unit_flag=True,
	73	+ enable_shuffle_k=True,
	74	+ use_hf32_mode=True
	75	+ )
	76	+ self.assertEqual(mmad.arch_tag, Arch.Ascend950)
	77	+ self.assertTrue(mmad.async_)
	78	+ self.assertEqual(mmad.preload_stages, 3)
	79	+ self.assertTrue(mmad.use_hf32_mode)
	80	+
	81	+
	82	+class TestMmadMultiBatch(unittest.TestCase):
	83	+ def test_mmad_multi_batch(self):
	84	+ mmad = MmadMultiBatch(
	85	+ arch_tag=Arch.Ascend950,
	86	+ use_hf32_mode=True,
	87	+ l0c_stages=3
	88	+ )
	89	+ self.assertEqual(mmad.arch_tag, Arch.Ascend950)
	90	+ self.assertFalse(mmad.async_)
	91	+ self.assertEqual(mmad.stages, 2)
	92	+ self.assertTrue(mmad.use_hf32_mode)
	93	+ self.assertEqual(mmad.l0c_stages, 3)
	94	+
	95	+
	96	+if __name__ == "__main__":
	97	+ unittest.main()

Apython/catlass_cppgen/tests/catlass/test_gemm_coord.py+45-0

@@ -0,0 +1,45 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+import unittest
	11	+from catlass_cppgen.catlass.gemm_coord import GemmShape, GemmCoord, Shape
	12	+import random
	13	+
	14	+
	15	+class TestGemmCoord(unittest.TestCase):
	16	+ def test_gemm_shape_str(self):
	17	+ m, n, k = [random.randint(0, 512)] * 3
	18	+ gemm_shape = GemmShape(m, n, k)
	19	+ self.assertEqual(str(gemm_shape), f"GemmShape<{m}, {n}, {k}>")
	20	+
	21	+ def test_gemm_shape_tla(self):
	22	+ m, n, k = [random.randint(0, 512)] * 3
	23	+ gemm_shape = GemmShape(m, n, k)
	24	+ shape = gemm_shape.tla()
	25	+ self.assertIsInstance(shape, Shape)
	26	+ self.assertEqual(shape.m, m)
	27	+ self.assertEqual(shape.n, n)
	28	+ self.assertEqual(shape.k, k)
	29	+
	30	+ def test_shape_str(self):
	31	+ m, n, k = [random.randint(0, 512)] * 3
	32	+ shape = Shape(m, n, k)
	33	+ expected = f"Shape<Int<{m}>, Int<{n}>, Int<{k}>>"
	34	+ self.assertEqual(str(shape), expected)
	35	+
	36	+ def test_gemm_coord_creation(self):
	37	+ m, n, k = [random.randint(0, 512)] * 3
	38	+ coord = GemmCoord(m, n, k)
	39	+ self.assertEqual(coord.m, m)
	40	+ self.assertEqual(coord.n, n)
	41	+ self.assertEqual(coord.k, k)
	42	+
	43	+
	44	+if __name__ == "__main__":
	45	+ unittest.main()

Apython/catlass_cppgen/tests/catlass/test_layout.py+157-0

@@ -0,0 +1,157 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+import unittest
	11	+from catlass_cppgen.catlass.layout.layout import (
	12	+ Coord,
	13	+ Layout,
	14	+ RowMajor,
	15	+ ColumnMajor,
	16	+ PaddingRowMajor,
	17	+ PaddingColumnMajor,
	18	+ VectorLayout,
	19	+ PrivateLayout,
	20	+ nZ,
	21	+ zN,
	22	+ zZ,
	23	+ nN,
	24	+)
	25	+import random
	26	+
	27	+
	28	+class TestCoord(unittest.TestCase):
	29	+ def test_coord_creation(self):
	30	+ values = [random.randint(0, 100) for _ in range(3)]
	31	+ coord = Coord(values)
	32	+ self.assertEqual(coord.idx, tuple(values))
	33	+
	34	+ def test_coord_single_value(self):
	35	+ coord = Coord([42])
	36	+ self.assertEqual(coord.idx, (42,))
	37	+
	38	+
	39	+class TestRowMajor(unittest.TestCase):
	40	+ def test_row_major_creation(self):
	41	+ m, n = random.randint(1, 100), random.randint(1, 100)
	42	+ layout = RowMajor((m, n))
	43	+ self.assertEqual(layout.shape, (m, n))
	44	+ self.assertEqual(layout.stride, (n, 1))
	45	+
	46	+ def test_row_major_capacity(self):
	47	+ m, n = random.randint(1, 100), random.randint(1, 100)
	48	+ layout = RowMajor((m, n))
	49	+ self.assertEqual(layout.capacity, m * n)
	50	+
	51	+ def test_row_major_is_need_padding_false(self):
	52	+ """测试不需要 padding 的情况"""
	53	+ # stride[0] = 64, align = 16, 64 % 16 == 0
	54	+ layout = RowMajor((4, 64))
	55	+ self.assertFalse(layout.is_need_padding(16))
	56	+
	57	+ def test_row_major_is_need_padding_true(self):
	58	+ """测试需要 padding 的情况"""
	59	+ # stride[0] = 65, align = 16, 65 % 16 != 0
	60	+ layout = RowMajor((4, 65))
	61	+ self.assertTrue(layout.is_need_padding(16))
	62	+
	63	+ def test_row_major_is_need_padding_large_stride(self):
	64	+ """测试 stride >= 65536 的情况"""
	65	+ layout = RowMajor((1000, 65536))
	66	+ self.assertTrue(layout.is_need_padding(16))
	67	+
	68	+ def test_row_major_get_padding_layout_no_padding(self):
	69	+ """测试不需要 padding 时返回自身"""
	70	+ layout = RowMajor((4, 64))
	71	+ result = layout.get_padding_layout(16)
	72	+ self.assertIs(result, layout)
	73	+
	74	+ def test_row_major_get_padding_layout_with_padding(self):
	75	+ """测试需要 padding 时返回 PaddingRowMajor"""
	76	+ pass
	77	+
	78	+
	79	+class TestColumnMajor(unittest.TestCase):
	80	+ def test_column_major_creation(self):
	81	+ m, n = random.randint(1, 100), random.randint(1, 100)
	82	+ layout = ColumnMajor((m, n))
	83	+ self.assertEqual(layout.shape, (m, n))
	84	+ self.assertEqual(layout.stride, (1, m))
	85	+
	86	+ def test_column_major_capacity(self):
	87	+ m, n = random.randint(1, 100), random.randint(1, 100)
	88	+ layout = ColumnMajor((m, n))
	89	+ self.assertEqual(layout.capacity, m * n)
	90	+
	91	+ def test_column_major_is_need_padding_false(self):
	92	+ """测试不需要 padding 的情况"""
	93	+ # stride[0] = 1, align = 16, 1 % 16 != 0，但 ColumnMajor 的 stride[0] 是 1
	94	+ # 实际上 ColumnMajor 的 stride 是 (1, m)，所以 stride[0] = 1
	95	+ layout = ColumnMajor((64, 4))
	96	+ # stride[0] = 1, 1 % 16 != 0，所以需要 padding
	97	+ self.assertTrue(layout.is_need_padding(16))
	98	+
	99	+ def test_column_major_is_need_padding_true(self):
	100	+ """测试需要 padding 的情况"""
	101	+ layout = ColumnMajor((65, 4))
	102	+ self.assertTrue(layout.is_need_padding(16))
	103	+
	104	+ def test_column_major_get_padding_layout_no_padding(self):
	105	+ """测试不需要 padding 时返回自身（这种情况在 ColumnMajor 中很少见）"""
	106	+ # 由于 ColumnMajor 的 stride[0] = 1，通常需要 padding
	107	+ layout = ColumnMajor((16, 4))
	108	+ result = layout.get_padding_layout(1) # align = 1 时不需要 padding
	109	+ self.assertIs(result, layout)
	110	+
	111	+ def test_column_major_get_padding_layout_with_padding(self):
	112	+ """测试需要 padding 时返回 PaddingColumnMajor"""
	113	+ pass
	114	+
	115	+
	116	+class TestPaddingLayouts(unittest.TestCase):
	117	+ def test_padding_row_major_value(self):
	118	+ pass
	119	+
	120	+ def test_padding_column_major_value(self):
	121	+ pass
	122	+
	123	+
	124	+class TestVectorLayout(unittest.TestCase):
	125	+ def test_vector_layout_creation(self):
	126	+ shape = random.randint(1, 100)
	127	+ layout = VectorLayout(shape)
	128	+ self.assertTrue(len(layout.shape) == 1 and layout.shape[0] == shape)
	129	+ self.assertTrue(len(layout.stride) == 1 and layout.stride[0] == 1)
	130	+
	131	+ def test_vector_layout_is_need_padding(self):
	132	+ """VectorLayout 不需要 padding"""
	133	+ layout = VectorLayout(100)
	134	+ self.assertFalse(layout.is_need_padding(16))
	135	+
	136	+
	137	+class TestPrivateLayouts(unittest.TestCase):
	138	+ def test_private_layout_is_need_padding(self):
	139	+ pass
	140	+
	141	+ def test_nz_value(self):
	142	+ pass
	143	+
	144	+ def test_zn_value(self):
	145	+ pass
	146	+
	147	+ def test_zz_value(self):
	148	+ pass
	149	+
	150	+ def test_nn_value(self):
	151	+ pass
	152	+ def test_private_layout_subclasses(self):
	153	+ pass
	154	+
	155	+
	156	+if __name__ == "__main__":
	157	+ unittest.main()

Apython/catlass_cppgen/tests/common/test_data_type.py+153-0

@@ -0,0 +1,153 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+import unittest
	11	+import torch
	12	+from catlass_cppgen.common.data_type import DataType, get_default_accumulator
	13	+
	14	+
	15	+class TestDataType(unittest.TestCase):
	16	+ def test_data_type_enum_values(self):
	17	+ """测试 DataType 枚举值"""
	18	+ self.assertEqual(DataType.FLOAT.value, "float")
	19	+ self.assertEqual(DataType.FLOAT16.value, "half")
	20	+ self.assertEqual(DataType.INT8.value, "int8_t")
	21	+ self.assertEqual(DataType.INT32.value, "int32_t")
	22	+ self.assertEqual(DataType.BF16.value, "bfloat16_t")
	23	+
	24	+ def test_from_dtype_float32(self):
	25	+ """测试从 torch.float32 创建 DataType"""
	26	+ dtype = DataType.from_dtype(torch.float32)
	27	+ self.assertEqual(dtype, DataType.FLOAT)
	28	+
	29	+ def test_from_dtype_float16(self):
	30	+ """测试从 torch.float16 创建 DataType"""
	31	+ dtype = DataType.from_dtype(torch.float16)
	32	+ self.assertEqual(dtype, DataType.FLOAT16)
	33	+
	34	+ def test_from_dtype_int8(self):
	35	+ """测试从 torch.int8 创建 DataType"""
	36	+ dtype = DataType.from_dtype(torch.int8)
	37	+ self.assertEqual(dtype, DataType.INT8)
	38	+
	39	+ def test_from_dtype_int32(self):
	40	+ """测试从 torch.int32 创建 DataType"""
	41	+ dtype = DataType.from_dtype(torch.int32)
	42	+ self.assertEqual(dtype, DataType.INT32)
	43	+
	44	+ def test_from_dtype_int64(self):
	45	+ """测试从 torch.int64 创建 DataType"""
	46	+ dtype = DataType.from_dtype(torch.int64)
	47	+ self.assertEqual(dtype, DataType.INT64)
	48	+
	49	+ def test_from_dtype_float64(self):
	50	+ """测试从 torch.float64 创建 DataType"""
	51	+ dtype = DataType.from_dtype(torch.float64)
	52	+ self.assertEqual(dtype, DataType.DOUBLE)
	53	+
	54	+ def test_from_dtype_bool(self):
	55	+ """测试从 torch.bool 创建 DataType"""
	56	+ dtype = DataType.from_dtype(torch.bool)
	57	+ self.assertEqual(dtype, DataType.BOOL)
	58	+
	59	+ def test_from_dtype_mxfp8(self):
	60	+ """测试 mxfp8 类型 (float8)"""
	61	+ if hasattr(torch, "float8_e5m2"):
	62	+ dtype = DataType.from_dtype(torch.float8_e5m2)
	63	+ self.assertEqual(dtype, DataType.FLOAT8_E5M2)
	64	+ if hasattr(torch, "float8_e4m3fn"):
	65	+ dtype = DataType.from_dtype(torch.float8_e4m3fn)
	66	+ self.assertEqual(dtype, DataType.FLOAT8_E4M3FN)
	67	+ if hasattr(torch, "float8_e8m0fnu"):
	68	+ dtype = DataType.from_dtype(torch.float8_e8m0fnu)
	69	+ self.assertEqual(dtype, DataType.FLOAT8_E8M0)
	70	+
	71	+ def test_from_dtype_mxfp4(self):
	72	+ """测试 mxfp4 类型 (float4)"""
	73	+ if hasattr(torch, "float4_e2m1fn_x2"):
	74	+ dtype = DataType.from_dtype(torch.float4_e2m1fn_x2)
	75	+ self.assertEqual(dtype, DataType.FLOAT4_E2M1)
	76	+ if hasattr(torch, "float4_e1m2fn_x2"):
	77	+ dtype = DataType.from_dtype(torch.float4_e1m2fn_x2)
	78	+ self.assertEqual(dtype, DataType.FLOAT4_E1M2)
	79	+
	80	+ def test_from_dtype_torch_aliases(self):
	81	+ """测试 torch 别名"""
	82	+ self.assertEqual(DataType.from_dtype(torch.float), DataType.FLOAT)
	83	+ self.assertEqual(DataType.from_dtype(torch.half), DataType.FLOAT16)
	84	+ self.assertEqual(DataType.from_dtype(torch.double), DataType.DOUBLE)
	85	+ self.assertEqual(DataType.from_dtype(torch.short), DataType.INT16)
	86	+ self.assertEqual(DataType.from_dtype(torch.int), DataType.INT32)
	87	+ self.assertEqual(DataType.from_dtype(torch.long), DataType.INT64)
	88	+
	89	+ def test_from_dtype_bfloat16(self):
	90	+ """测试 bfloat16"""
	91	+ if hasattr(torch, "bfloat16"):
	92	+ dtype = DataType.from_dtype(torch.bfloat16)
	93	+ self.assertEqual(dtype, DataType.BF16)
	94	+
	95	+ def test_from_dtype_undefined(self):
	96	+ """测试未定义的类型返回 UNDEFINED"""
	97	+ # 传入一个不存在的类型
	98	+ dtype = DataType.from_dtype("unknown_type")
	99	+ self.assertEqual(dtype, DataType.UNDEFINED)
	100	+
	101	+ def test_data_size_float(self):
	102	+ """测试 FLOAT 的数据大小"""
	103	+ self.assertEqual(DataType.FLOAT.data_size(), 4)
	104	+
	105	+ def test_data_size_float16(self):
	106	+ """测试 FLOAT16 的数据大小"""
	107	+ self.assertEqual(DataType.FLOAT16.data_size(), 2)
	108	+
	109	+ def test_data_size_bf16(self):
	110	+ """测试 BF16 的数据大小"""
	111	+ self.assertEqual(DataType.BF16.data_size(), 2)
	112	+
	113	+ def test_data_size_int8(self):
	114	+ """测试 INT8 的数据大小"""
	115	+ self.assertEqual(DataType.INT8.data_size(), 1)
	116	+
	117	+ def test_data_size_int32(self):
	118	+ """测试 INT32 的数据大小"""
	119	+ self.assertEqual(DataType.INT32.data_size(), 4)
	120	+
	121	+ def test_data_size_int64(self):
	122	+ """测试 INT64 的数据大小"""
	123	+ self.assertEqual(DataType.INT64.data_size(), 8)
	124	+
	125	+
	126	+class TestGetDefaultAccumulator(unittest.TestCase):
	127	+ def test_get_default_accumulator_float16(self):
	128	+ """测试 FLOAT16 + FLOAT16 -> FLOAT"""
	129	+ acc = get_default_accumulator(DataType.FLOAT16, DataType.FLOAT16)
	130	+ self.assertEqual(acc, DataType.FLOAT)
	131	+
	132	+ def test_get_default_accumulator_float_float16(self):
	133	+ """测试 FLOAT + FLOAT16 -> FLOAT"""
	134	+ pass
	135	+
	136	+ def test_get_default_accumulator_bf16(self):
	137	+ """测试 BF16 + BF16 -> FLOAT"""
	138	+ acc = get_default_accumulator(DataType.BF16, DataType.BF16)
	139	+ self.assertEqual(acc, DataType.FLOAT)
	140	+
	141	+ def test_get_default_accumulator_int8(self):
	142	+ """测试 INT8 + INT8 -> INT32"""
	143	+ acc = get_default_accumulator(DataType.INT8, DataType.INT8)
	144	+ self.assertEqual(acc, DataType.INT32)
	145	+
	146	+ def test_get_default_accumulator_float(self):
	147	+ """测试 FLOAT + FLOAT -> FLOAT (默认返回自身)"""
	148	+ acc = get_default_accumulator(DataType.FLOAT, DataType.FLOAT)
	149	+ self.assertEqual(acc, DataType.FLOAT)
	150	+
	151	+
	152	+if __name__ == "__main__":
	153	+ unittest.main()

Apython/catlass_cppgen/tests/common/test_op_tensor.py+75-0

@@ -0,0 +1,75 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+import unittest
	11	+import ctypes
	12	+from catlass_cppgen.common.op_tensor import OpTensor
	13	+from catlass_cppgen.common.data_type import DataType
	14	+from catlass_cppgen.catlass.layout.layout import RowMajor, ColumnMajor
	15	+import random
	16	+
	17	+
	18	+class TestOpTensor(unittest.TestCase):
	19	+ def test_op_tensor_creation(self):
	20	+ """测试 OpTensor 的创建"""
	21	+ dtype = DataType.FLOAT
	22	+ layout = RowMajor((128, 256))
	23	+ tensor = OpTensor(dtype, layout)
	24	+ self.assertEqual(tensor.dtype, dtype)
	25	+ self.assertEqual(tensor.layout, layout)
	26	+
	27	+ def test_op_tensor_with_data_ptr(self):
	28	+ """测试带 data_ptr 的 OpTensor 创建"""
	29	+ dtype = DataType.FLOAT
	30	+ layout = RowMajor((128, 256))
	31	+ data_ptr = ctypes.c_void_p(0x12345678)
	32	+ tensor = OpTensor(dtype, layout, data_ptr)
	33	+ self.assertEqual(tensor.dtype, dtype)
	34	+ self.assertEqual(tensor.layout, layout)
	35	+
	36	+ def test_op_tensor_shape_property(self):
	37	+ """测试 shape 属性"""
	38	+ m, n = random.randint(1, 100), random.randint(1, 100)
	39	+ layout = RowMajor((m, n))
	40	+ tensor = OpTensor(DataType.FLOAT, layout)
	41	+ self.assertEqual(tensor.shape, (m, n))
	42	+
	43	+ def test_op_tensor_stride_property(self):
	44	+ """测试 stride 属性"""
	45	+ m, n = random.randint(1, 100), random.randint(1, 100)
	46	+ layout = RowMajor((m, n))
	47	+ tensor = OpTensor(DataType.FLOAT, layout)
	48	+ self.assertEqual(tensor.stride, (n, 1))
	49	+
	50	+ def test_op_tensor_capacity_property(self):
	51	+ """测试 capacity 属性"""
	52	+ m, n = random.randint(1, 100), random.randint(1, 100)
	53	+ layout = RowMajor((m, n))
	54	+ tensor = OpTensor(DataType.FLOAT, layout)
	55	+ self.assertEqual(tensor.capacity, m * n)
	56	+
	57	+ def test_op_tensor_column_major(self):
	58	+ """测试使用 ColumnMajor layout"""
	59	+ m, n = random.randint(1, 100), random.randint(1, 100)
	60	+ layout = ColumnMajor((m, n))
	61	+ tensor = OpTensor(DataType.FLOAT16, layout)
	62	+ self.assertEqual(tensor.shape, (m, n))
	63	+ self.assertEqual(tensor.stride, (1, m))
	64	+ self.assertEqual(tensor.capacity, m * n)
	65	+
	66	+ def test_op_tensor_different_dtypes(self):
	67	+ """测试不同的数据类型"""
	68	+ layout = RowMajor((64, 128))
	69	+ for dtype in [DataType.FLOAT, DataType.FLOAT16, DataType.INT8, DataType.INT32]:
	70	+ tensor = OpTensor(dtype, layout)
	71	+ self.assertEqual(tensor.dtype, dtype)
	72	+
	73	+
	74	+if __name__ == "__main__":
	75	+ unittest.main()

Apython/catlass_cppgen/tests/op/assertion_helper.py+97-0

@@ -0,0 +1,97 @@
	1	+#!/usr/bin/env python3
	2	+# -- coding: utf-8 --
	3	+# ----------------------------------------------------------------------------
	4	+# This program is free software, you can redistribute it and/or modify.
	5	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	6	+# This file is a part of the CANN Open Software.
	7	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	8	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	9	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	10	+# See LICENSE in the root of the software repository for the full text of the License.
	11	+# ----------------------------------------------------------------------------
	12	+
	13	+from collections.abc import Iterable
	14	+
	15	+import re
	16	+import unittest
	17	+from typing import Union, List, Tuple
	18	+
	19	+from catlass_cppgen.catlass.gemm_coord import GemmShape
	20	+from catlass_cppgen.common.data_type import DataType
	21	+from catlass_cppgen.catlass.arch.arch import Arch
	22	+from catlass_cppgen.catlass.library import BroadcastTag, BroadcastType, EpilogueOp, EpilogueOpTag
	23	+
	24	+class TestAssertions:
	CheaterAbec7月6日建议文件重命名为helper或者类似的 I init__zhb__ 7月8日评论： 7月8日评论：
	25	+ def __init__(self, test_case: unittest.TestCase):
	26	+ self.t = test_case
	27	+
	28	+ def test_params(self, params_str: str, params_list: Union[list, tuple]):
	29	+ actual_params_list = [p.strip() for p in params_str.split(",")]
	30	+ self.t.assertEqual(len(actual_params_list), len(params_list))
	31	+ for params, actual_params in zip(params_list, actual_params_list):
	32	+ self.t.assertEqual(params, actual_params)
	33	+
	34	+ def test_tileshape(self, kernel_str: str, tiling: GemmShape, pos: str = "L1"):
	35	+ self.t.assertIn(pos, ("L1", "L0"), "argument 'pos' can only be 'L1' or 'L0'")
	36	+
	37	+ pattern = r"using\s+L1TileShape\s=\sShape<Int<(\d+)>,\sInt<(\d+)>,\sInt<(\d+)>>;" if pos == "L1" else \
	38	+ r"using\s+L0TileShape\s=\sShape<Int<(\d+)>,\sInt<(\d+)>,\sInt<(\d+)>>;"
	39	+ match_tiling = re.search(pattern, kernel_str)
	40	+ self.t.assertIsNotNone(match_tiling)
	41	+ self.t.assertEqual(match_tiling.group(1), str(tiling.m))
	42	+ self.t.assertEqual(match_tiling.group(2), str(tiling.n))
	43	+ self.t.assertEqual(match_tiling.group(3), str(tiling.k))
	44	+
	45	+ def test_element_dtype(self, template_str: str, dtype: DataType, pos: str = "A"):
	46	+ match = re.search(rf"using\s+Element{pos}\s=\s(\S+);", template_str)
	47	+ self.t.assertIsNotNone(match, f"Element{pos} not found for dtype {dtype.value}")
	48	+ self.t.assertEqual(match.group(1), dtype.value)
	49	+
	50	+ def test_accu_dtype(self, template_str: str, dtype: DataType):
	51	+ match = re.search(r"Catlass::Epilogue::Fusion::VisitorAccLoad<(\S+)>;", template_str)
	52	+ self.t.assertIsNotNone(match, f"VistorAccLoad not found for dtype {dtype.value}")
	53	+ self.t.assertEqual(match.group(1), dtype.value)
	54	+
	55	+ def test_layout(self, template_str: str, layout_tag: str, pos: str = "TagA"):
	56	+ self.t.assertIn(pos, ("TagA", "TagB", "TagC", "TagBias"), "invalid layout pos.")
	57	+ match = re.search(rf"using\s+Layout{pos}\s=\s(?:\w+::)*layout::(\S+);", template_str)
	58	+ self.t.assertIsNotNone(match, f"Layout{pos} not found for layout {layout_tag}")
	59	+ self.t.assertEqual(layout_tag, match.group(1))
	60	+
	61	+ def test_kernel(self, template_str: str, kernel_name: str):
	62	+ match = re.search(r"using\s+GemmKernel\s=\sGemm::Kernel::(\S+)<", template_str)
	63	+ self.t.assertIsNotNone(match, f"Kernel {kernel_name} not found in template")
	64	+ self.t.assertEqual(match.group(1), kernel_name)
	65	+
	66	+ def test_dispatch_policy(self, template_str: str, dispatch_name: str):
	67	+ match = re.search(r"using\s+DispatchPolicy\s=\sGemm::(\S+)<", template_str)
	68	+ self.t.assertIsNotNone(match, f"DispatchPolicy {dispatch_name} not found in template")
	69	+ self.t.assertEqual(match.group(1), dispatch_name)
	70	+
	71	+ def test_arch_tag(self, template_str: str, arch: Arch):
	72	+ arch_str = arch.value
	73	+ match = re.search(r"using\s+ArchTag\s=\s(\S+);", template_str)
	74	+ self.t.assertIsNotNone(match, f"ArchTag not found for {arch_str}")
	75	+ self.t.assertEqual(match.group(1), arch_str)
	76	+
	77	+class TestEvgAssertions(TestAssertions):
	78	+ def __init__(self, test_case: unittest.TestCase):
	79	+ super().__init__(test_case)
	80	+
	81	+ def test_boardcast(self, template_str: str, node: str, boardcast: BroadcastType):
	82	+ match = re.search(rf"using\s+{node}\s=\sCatlass::Epilogue::Fusion::Visitor(\S+)<", template_str)
	83	+ self.t.assertIsNotNone(match, f"Boardcast {boardcast} not found for node {node}")
	84	+ self.t.assertEqual(match.group(1), BroadcastTag[boardcast])
	85	+
	86	+ def test_visitor_compute(self, template_str: str, op: Union[EpilogueOp, List[EpilogueOp]]):
	87	+ vistor_pattern = r"Catlass::Epilogue::Fusion::VisitorCompute<(\S+),\s*\S+>;"
	88	+ op_list = (EpilogueOpTag[x] for x in op) if isinstance(op, Iterable) else EpilogueOpTag[op]
	89	+ for match in re.finditer(vistor_pattern, template_str):
	90	+ self.t.assertIsNotNone(match, f"VisitorCompute not found")
	91	+ self.t.assertIn(match.group(1), op_list)
	92	+
	93	+ def test_tree_visitor(self, template_str: str, vistor_params: List[str]):
	94	+ vistor_pattern = r"Catlass::Epilogue::Fusion::TreeVisitor<(\S+)>;"
	95	+ for match, vistor_param in zip(re.finditer(vistor_pattern, template_str), vistor_params):
	96	+ self.t.assertIsNotNone(match, f"TreeVisitor not found")
	97	+ self.t.assertEqual(match.group(1), vistor_param)

Apython/catlass_cppgen/tests/op/test_evg.py+418-0

@@ -0,0 +1,418 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+import unittest
	11	+from catlass_cppgen.op.gemm import Gemm
	12	+from catlass_cppgen.catlass.library import BroadcastType, EpilogueOp
	13	+from catlass_cppgen.common.op_tensor import OpTensor
	14	+from catlass_cppgen.common.data_type import DataType
	15	+from catlass_cppgen.catlass.layout.layout import RowMajor
	16	+from catlass_cppgen.catlass.gemm_coord import GemmShape
	17	+from catlass_cppgen.catlass.arch.arch import Arch
	18	+from catlass_cppgen.catlass.evg_extension import evg
	19	+
	20	+from assertion_helper import TestEvgAssertions
	21	+
	22	+class TestBaseOp(unittest.TestCase):
	23	+ def setUp(self):
	24	+ self.check = TestEvgAssertions(self)
	25	+
	26	+ def test_codegen_with_evg_direct(self):
	27	+ fn_src = """
	28	+def epilogue(accum, bias):
	29	+ result = accum + bias
	30	+ return result
	31	+"""
	32	+ example_inputs = {
	33	+ "accum": OpTensor.from_shape_stride(
	34	+ shape=(128, 256),
	35	+ stride=(256, 1),
	36	+ dtype=DataType.FLOAT
	37	+ ),
	38	+ "bias": OpTensor.from_shape_stride(
	39	+ shape=(1, 256),
	40	+ stride=(256, 1),
	41	+ dtype=DataType.FLOAT
	42	+ ),
	43	+ "result": OpTensor.from_shape_stride(
	44	+ shape=(128, 256),
	45	+ stride=(256, 1),
	46	+ dtype=DataType.FLOAT
	47	+ ),
	48	+ }
	49	+
	50	+ from catlass_cppgen.kernel.gemm.basic_matmul_tla_visitor import BasicMatmulTlaVisitorKernel
	51	+
	52	+ kernel = BasicMatmulTlaVisitorKernel(
	53	+ element_accumulator=DataType.FLOAT,
	54	+ element_A=DataType.FLOAT,
	55	+ element_B=DataType.FLOAT,
	56	+ element_C=DataType.FLOAT,
	57	+ element_Bias=DataType.FLOAT,
	58	+ layout_A=RowMajor((128, 256)),
	59	+ layout_B=RowMajor((256, 384)),
	60	+ arch_tag=Arch.Ascend950,
	61	+ M=128,
	62	+ K=256,
	63	+ N=384,
	64	+ evg={
	65	+ "fn_src": fn_src,
	66	+ "example_inputs": example_inputs,
	67	+ }
	68	+ )
	69	+
	70	+ evg_template = kernel.gen_evg_template()
	71	+ kernel_template = kernel.gen_kernel_template()
	72	+
	73	+ self.check.test_element_dtype(kernel_template, DataType.FLOAT, "A")
	74	+ self.check.test_element_dtype(kernel_template, DataType.FLOAT, "B")
	75	+ self.check.test_element_dtype(kernel_template, DataType.FLOAT, "C")
	76	+ self.check.test_layout(kernel_template, "RowMajor", "TagA")
	77	+ self.check.test_layout(kernel_template, "RowMajor", "TagB")
	78	+ self.check.test_layout(kernel_template, "RowMajor", "TagC")
	79	+
	80	+ self.check.test_arch_tag(kernel_template, Arch.Ascend950)
	81	+
	82	+ self.check.test_accu_dtype(evg_template, DataType.FLOAT)
	83	+ self.check.test_layout(evg_template, "RowMajor", "TagBias")
	84	+ self.check.test_boardcast(evg_template, "Bias", BroadcastType.RowBroadcast)
	85	+ self.check.test_visitor_compute(evg_template, EpilogueOp.Add)
	86	+ self.check.test_tree_visitor(evg_template, (
	87	+ "Compute0, Accum, Bias",
	88	+ "Result, EVGCompute0"
	89	+ )
	90	+ )
	91	+
	92	+ def test_codegen_without_evg(self):
	93	+ from catlass_cppgen.kernel.gemm.basic_matmul_tla_visitor import BasicMatmulTlaVisitorKernel
	94	+
	95	+ kernel = BasicMatmulTlaVisitorKernel(
	96	+ element_accumulator=DataType.FLOAT,
	97	+ element_A=DataType.FLOAT,
	98	+ element_B=DataType.FLOAT,
	99	+ element_C=DataType.FLOAT,
	100	+ element_Bias=DataType.FLOAT,
	101	+ layout_A=RowMajor((128, 256)),
	102	+ layout_B=RowMajor((256, 384)),
	103	+ arch_tag=Arch.Ascend950,
	104	+ M=128,
	105	+ K=256,
	106	+ N=384,
	107	+ evg=None,
	108	+ )
	109	+
	110	+ kernel_template = kernel.gen_kernel_template()
	111	+
	112	+ self.check.test_element_dtype(kernel_template, DataType.FLOAT, "A")
	113	+ self.check.test_element_dtype(kernel_template, DataType.FLOAT, "B")
	114	+ self.check.test_element_dtype(kernel_template, DataType.FLOAT, "C")
	115	+ self.check.test_layout(kernel_template, "RowMajor", "TagA")
	116	+ self.check.test_layout(kernel_template, "RowMajor", "TagB")
	117	+ self.check.test_layout(kernel_template, "RowMajor", "TagC")
	118	+ self.check.test_arch_tag(kernel_template, Arch.Ascend950)
	119	+
	120	+ def test_codegen_with_evg(self):
	121	+ a = OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT)
	122	+ b = OpTensor.from_shape_stride(shape=(256, 384), stride=(384, 1), dtype=DataType.FLOAT)
	123	+
	124	+ fn_src = """
	125	+def epilogue(accum, bias):
	126	+ result = accum + bias
	127	+ return result
	128	+"""
	129	+ example_inputs = {
	130	+ "accum": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	131	+ "bias": OpTensor.from_shape_stride(shape=(1, 256), stride=(256, 1), dtype=DataType.FLOAT),
	132	+ "result": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	133	+ }
	134	+
	135	+ gemm_plan = Gemm(
	136	+ atlas_arch=Arch.Ascend950,
	137	+ element=DataType.FLOAT,
	138	+ layout=RowMajor,
	139	+ evg_config={"fn_src": fn_src, "example_inputs": example_inputs},
	140	+ A=a, B=b
	141	+ )
	142	+ kernels = gemm_plan.get_kernels()
	143	+ kernel = kernels[0]
	144	+
	145	+ from catlass_cppgen.kernel.gemm.basic_matmul_tla_visitor import BasicMatmulTlaVisitorKernel
	146	+ self.assertIsInstance(kernel, BasicMatmulTlaVisitorKernel)
	147	+
	148	+ evg_template = kernel.gen_evg_template()
	149	+ kernel_template = kernel.gen_kernel_template()
	150	+
	151	+ self.check.test_element_dtype(kernel_template, DataType.FLOAT, "A")
	152	+ self.check.test_element_dtype(kernel_template, DataType.FLOAT, "B")
	153	+ self.check.test_element_dtype(kernel_template, DataType.FLOAT, "C")
	154	+ self.check.test_layout(kernel_template, "RowMajor", "TagA")
	155	+ self.check.test_layout(kernel_template, "RowMajor", "TagB")
	156	+ self.check.test_layout(kernel_template, "RowMajor", "TagC")
	157	+ self.check.test_arch_tag(kernel_template, Arch.Ascend950)
	158	+
	159	+ self.check.test_accu_dtype(evg_template, DataType.FLOAT)
	160	+ self.check.test_layout(evg_template, "RowMajor", "TagBias")
	161	+ self.check.test_boardcast(evg_template, "Bias", BroadcastType.RowBroadcast)
	162	+ self.check.test_visitor_compute(evg_template, EpilogueOp.Add)
	163	+ self.check.test_tree_visitor(evg_template, (
	164	+ "Compute0, Accum, Bias",
	165	+ "Result, EVGCompute0"
	166	+ )
	167	+ )
	168	+
	169	+ def test_tune_functionality(self):
	170	+ a = OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT)
	171	+ b = OpTensor.from_shape_stride(shape=(256, 384), stride=(384, 1), dtype=DataType.FLOAT)
	172	+
	173	+ fn_src = """
	174	+def epilogue(accum, bias):
	175	+ result = accum + bias
	176	+ return result
	177	+"""
	178	+ example_inputs = {
	179	+ "accum": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	180	+ "bias": OpTensor.from_shape_stride(shape=(1, 256), stride=(256, 1), dtype=DataType.FLOAT),
	181	+ "result": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	182	+ }
	183	+
	184	+ gemm_plan = Gemm(
	185	+ atlas_arch=Arch.Ascend950,
	186	+ element=DataType.FLOAT,
	187	+ layout=RowMajor,
	188	+ evg_config={"fn_src": fn_src, "example_inputs": example_inputs},
	189	+ A=a, B=b
	190	+ )
	191	+ kernels = gemm_plan.get_kernels()
	192	+ kernel = kernels[0]
	193	+
	194	+ default_l1, default_l0 = kernel.get_default_tile_shape()
	195	+ self.assertEqual(default_l1, GemmShape(256, 256, 128))
	196	+ self.assertEqual(default_l0, GemmShape(256, 256, 32))
	197	+
	198	+ new_l1 = GemmShape(128, 256, 64)
	199	+ new_l0 = GemmShape(128, 256, 64)
	200	+ kernel.tune(l1_tile_shape=new_l1, l0_tile_shape=new_l0)
	201	+ self.assertEqual(kernel.l1_tile_shape, new_l1)
	202	+ self.assertEqual(kernel.l0_tile_shape, new_l0)
	203	+
	204	+ evg_template = kernel.gen_evg_template()
	205	+ kernel_template = kernel.gen_kernel_template()
	206	+
	207	+ self.check.test_tileshape(kernel_template, GemmShape(128, 256, 64))
	208	+ self.check.test_tileshape(kernel_template, GemmShape(128, 256, 64), pos="L0")
	209	+
	210	+ self.check.test_element_dtype(kernel_template, DataType.FLOAT, "A")
	211	+ self.check.test_element_dtype(kernel_template, DataType.FLOAT, "B")
	212	+ self.check.test_element_dtype(kernel_template, DataType.FLOAT, "C")
	213	+ self.check.test_layout(kernel_template, "RowMajor", "TagA")
	214	+ self.check.test_layout(kernel_template, "RowMajor", "TagB")
	215	+ self.check.test_layout(kernel_template, "RowMajor", "TagC")
	216	+ self.check.test_arch_tag(kernel_template, Arch.Ascend950)
	217	+
	218	+ self.check.test_accu_dtype(evg_template, DataType.FLOAT)
	219	+ self.check.test_layout(evg_template, "RowMajor", "TagBias")
	220	+ self.check.test_boardcast(evg_template, "Bias", BroadcastType.RowBroadcast)
	221	+ self.check.test_visitor_compute(evg_template, EpilogueOp.Add)
	222	+ self.check.test_tree_visitor(evg_template, (
	223	+ "Compute0, Accum, Bias",
	224	+ "Result, EVGCompute0"
	225	+ )
	226	+ )
	227	+
	228	+
	229	+######################
	230	+
	231	+
	232	+class TestEvgOp(unittest.TestCase):
	233	+ def setUp(self):
	234	+ self.check = TestEvgAssertions(self)
	235	+
	236	+ def _build_evg(self,
	237	+ fn_src: str,
	238	+ inputs: str,
	239	+ output_dtype: DataType = DataType.FLOAT):
	240	+ callback_name, evg_args, evg_str, _ = evg(
	241	+ fn_src=fn_src,
	242	+ example_inputs=inputs
	243	+ )
	244	+ self.assertEqual(callback_name, "EVGResult") # fixed callback
	245	+
	246	+ return evg_str, evg_args
	247	+
	248	+ def test_evg_add(self):
	249	+ fn_src = """
	250	+def epilogue(accum, bias):
	251	+ result = accum + bias
	252	+ return result
	253	+"""
	254	+ inputs = {
	255	+ "accum": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	256	+ "bias": OpTensor.from_shape_stride(shape=(1, 256), stride=(256, 1), dtype=DataType.FLOAT),
	257	+ "result": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	258	+ }
	259	+ evg_str, _ = self._build_evg(fn_src, inputs)
	260	+ self.check.test_accu_dtype(evg_str, DataType.FLOAT)
	261	+ self.check.test_layout(evg_str, "RowMajor", "TagBias")
	262	+ self.check.test_boardcast(evg_str, "Bias", BroadcastType.RowBroadcast)
	263	+ self.check.test_visitor_compute(evg_str, EpilogueOp.Add)
	264	+ self.check.test_tree_visitor(evg_str, (
	265	+ "Compute0, Accum, Bias",
	266	+ "Result, EVGCompute0"
	267	+ )
	268	+ )
	269	+
	270	+ def test_evg_mul(self):
	271	+ fn_src = """
	272	+def epilogue(accum, scale):
	273	+ result = accum * scale
	274	+ return result
	275	+"""
	276	+ inputs = {
	277	+ "accum": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	278	+ "scale": OpTensor.from_shape_stride(shape=(1, 256), stride=(256, 1), dtype=DataType.FLOAT),
	279	+ "result": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	280	+ }
	281	+ evg_str, _ = self._build_evg(fn_src, inputs)
	282	+ self.check.test_accu_dtype(evg_str, DataType.FLOAT)
	283	+ self.check.test_boardcast(evg_str, "Scale", BroadcastType.RowBroadcast)
	284	+ self.check.test_visitor_compute(evg_str, EpilogueOp.Mul)
	285	+ self.check.test_tree_visitor(evg_str, (
	286	+ "Compute0, Accum, Scale",
	287	+ "Result, EVGCompute0"
	288	+ )
	289	+ )
	290	+
	291	+ def test_evg_relu(self):
	292	+ fn_src = """
	293	+def epilogue(accum):
	294	+ result = relu(accum)
	295	+ return result
	296	+"""
	297	+ inputs = {
	298	+ "accum": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	299	+ "result": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	300	+ }
	301	+ evg_str, _ = self._build_evg(fn_src, inputs)
	302	+ self.check.test_accu_dtype(evg_str, DataType.FLOAT)
	303	+ self.check.test_visitor_compute(evg_str, EpilogueOp.Relu)
	304	+ self.check.test_tree_visitor(evg_str, (
	305	+ "Compute0, Accum",
	306	+ "Result, EVGCompute0"
	307	+ )
	308	+ )
	309	+
	310	+ def test_evg_leaky_relu(self):
	311	+ fn_src = """
	312	+def epilogue(accum, alpha):
	313	+ result = leakyRelu(accum, alpha)
	314	+ return result
	315	+"""
	316	+ inputs = {
	317	+ "accum": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	318	+ "alpha": OpTensor.from_shape_stride(shape=(1, 256), stride=(256, 1), dtype=DataType.FLOAT),
	319	+ "result": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	320	+ }
	321	+ evg_str, _ = self._build_evg(fn_src, inputs)
	322	+ self.check.test_accu_dtype(evg_str, DataType.FLOAT)
	323	+ self.check.test_boardcast(evg_str, "Alpha", BroadcastType.RowBroadcast)
	324	+ self.check.test_visitor_compute(evg_str, EpilogueOp.LeakyRelu)
	325	+ self.check.test_tree_visitor(evg_str, (
	326	+ "Compute0, Accum, Alpha",
	327	+ "Result, EVGCompute0"
	328	+ )
	329	+ )
	330	+
	331	+ def test_evg_prelu(self):
	332	+ fn_src = """
	333	+def epilogue(accum, weight):
	334	+ result = Prelu(accum, weight)
	335	+ return result
	336	+"""
	337	+ inputs = {
	338	+ "accum": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	339	+ "weight": OpTensor.from_shape_stride(shape=(1, 256), stride=(256, 1), dtype=DataType.FLOAT),
	340	+ "result": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	341	+ }
	342	+ evg_str, _ = self._build_evg(fn_src, inputs)
	343	+ self.check.test_accu_dtype(evg_str, DataType.FLOAT)
	344	+ self.check.test_boardcast(evg_str, "Weight", BroadcastType.RowBroadcast)
	345	+ self.check.test_visitor_compute(evg_str, EpilogueOp.Prelu)
	346	+ self.check.test_tree_visitor(evg_str, (
	347	+ "Compute0, Accum, Weight",
	348	+ "Result, EVGCompute0"
	349	+ )
	350	+ )
	351	+
	352	+ def test_evg_sigmoid(self):
	353	+ fn_src = """
	354	+def epilogue(accum):
	355	+ result = sigmoid(accum)
	356	+ return result
	357	+"""
	358	+ inputs = {
	359	+ "accum": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	360	+ "result": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	361	+ }
	362	+ evg_str, _ = self._build_evg(fn_src, inputs)
	363	+ self.check.test_accu_dtype(evg_str, DataType.FLOAT)
	364	+ self.check.test_visitor_compute(evg_str, EpilogueOp.Sigmoid)
	365	+ self.check.test_tree_visitor(evg_str, (
	366	+ "Compute0, Accum",
	367	+ "Result, EVGCompute0"
	368	+ )
	369	+ )
	370	+
	371	+ def test_evg_cast(self):
	372	+ fn_src = """
	373	+def epilogue(accum):
	374	+ result = cast(accum, "half", "float")
	375	+ return result
	376	+"""
	377	+ inputs = {
	378	+ "accum": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	379	+ "result": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT16),
	380	+ }
	381	+ evg_str, _ = self._build_evg(fn_src, inputs)
	382	+ self.check.test_accu_dtype(evg_str, DataType.FLOAT)
	383	+ self.check.test_visitor_compute(evg_str, EpilogueOp.Cast)
	384	+ self.check.test_tree_visitor(evg_str, (
	385	+ "Compute0, Accum",
	386	+ "Result, EVGCompute0"
	387	+ )
	388	+ )
	389	+
	390	+ def test_evg_relu_add(self):
	391	+ fn_src = """
	392	+def epilogue(accum, bias):
	393	+ relu_result = relu(accum)
	394	+ result = relu_result + bias
	395	+ return result
	396	+"""
	397	+ inputs = {
	398	+ "accum": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	399	+ "bias": OpTensor.from_shape_stride(shape=(1, 256), stride=(256, 1), dtype=DataType.FLOAT),
	400	+ "result": OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT),
	401	+ }
	402	+ evg_str, _ = self._build_evg(fn_src, inputs)
	403	+ self.check.test_accu_dtype(evg_str, DataType.FLOAT)
	404	+ self.check.test_boardcast(evg_str, "Bias", BroadcastType.RowBroadcast)
	405	+ self.check.test_visitor_compute(evg_str, [EpilogueOp.Relu, EpilogueOp.Add])
	406	+ self.check.test_tree_visitor(evg_str, (
	407	+ "Compute0, Accum",
	408	+ "Compute1, EVGCompute0, Bias",
	409	+ "Result, EVGCompute1"
	410	+ )
	411	+ )
	412	+
	413	+
	414	+######################################
	415	+
	416	+
	417	+if __name__ == "__main__":
	418	+ unittest.main()

Apython/catlass_cppgen/tests/op/test_gemm.py+736-0

@@ -0,0 +1,736 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+import unittest
	11	+from catlass_cppgen.op.gemm import Gemm
	12	+from catlass_cppgen.common.op_tensor import OpTensor
	13	+from catlass_cppgen.common.data_type import DataType
	14	+from catlass_cppgen.catlass.layout.layout import RowMajor
	15	+from catlass_cppgen.catlass.gemm_coord import GemmShape
	16	+from catlass_cppgen.catlass.arch.arch import Arch
	17	+from catlass_cppgen.catlass.gemm.dispatch_policy import (
	18	+ MmadPingpong,
	19	+ MmadPreloadAsyncWithCallback,
	20	+ MmadMultiBatch,
	21	+ MmadAtlasA2Pingpong,
	22	+)
	23	+from catlass_cppgen.kernel.gemm import (
	24	+ BasicMatmulKernel,
	25	+ BatchedMatmulKernel,
	26	+ BasicMatmulTlaVisitorKernel,
	27	+ MultiCoreSplitkMatmulKernel,
	28	+ StreamkMatmulKernel,
	29	+ TailMultiCoreSplitkMatmulKernel,
	30	+)
	31	+
	32	+from assertion_helper import TestAssertions
	33	+
	34	+
	35	+def find_kernel_by_type(kernels, kernel_type):
	36	+ """根据类型查找 kernel"""
	37	+ for kernel in kernels:
	38	+ if isinstance(kernel, kernel_type):
	39	+ return kernel
	40	+ return None
	41	+
	42	+class TestGemm(unittest.TestCase):
	43	+ def setUp(self):
	44	+ self.check = TestAssertions(self)
	45	+
	46	+ def test_basic_matmul_kernel(self):
	47	+ a = OpTensor.from_shape_stride(
	48	+ shape=(128, 256),
	49	+ stride=(256, 1),
	50	+ dtype=DataType.FLOAT
	51	+ )
	52	+ b = OpTensor.from_shape_stride(
	53	+ shape=(256, 384),
	54	+ stride=(384, 1),
	55	+ dtype=DataType.FLOAT
	56	+ )
	57	+
	58	+ gemm_plan = Gemm(atlas_arch=Arch.Ascend950, element=DataType.FLOAT, layout=RowMajor, A=a, B=b)
	59	+ kernels = gemm_plan.get_kernels()
	60	+
	61	+ # 根据类型查找 kernel，而不是使用固定索引
	62	+ basic_kernel = find_kernel_by_type(kernels, BasicMatmulKernel)
	63	+ if basic_kernel is None:
	64	+ raise ValueError(f"No kernel named 'BasicMatmulKernel' found, available kernel list: {[type(k).__name__ for k in kernels]}")
	65	+
	66	+ # 基础特征检查
	67	+ self.assertEqual(type(basic_kernel), BasicMatmulKernel)
	68	+ self.assertFalse(basic_kernel.relu_enable) # default to False
	69	+
	70	+ basic_kernel.tune(
	71	+ GemmShape(128, 256, 64),
	72	+ GemmShape(128, 256, 64),
	73	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950)
	74	+ )
	75	+
	76	+ params = basic_kernel.gen_params_device(def_mode=False)
	77	+ self.check.test_params(params, (
	78	+ "problemShape",
	79	+ "deviceA", "layoutA",
	80	+ "deviceB", "layoutB",
	81	+ "deviceC", "layoutC",
	82	+ "deviceBias"
	83	+ ))
	84	+
	85	+ kernel_str = basic_kernel.gen_kernel_template()
	86	+
	87	+ self.check.test_tileshape(kernel_str, GemmShape(128, 256, 64))
	88	+ self.check.test_tileshape(kernel_str, GemmShape(128, 256, 64), pos="L0")
	89	+ self.check.test_element_dtype(kernel_str, DataType.FLOAT)
	90	+ self.check.test_layout(kernel_str, "RowMajor")
	91	+ self.check.test_layout(kernel_str, "RowMajor", "TagB")
	92	+ self.check.test_layout(kernel_str, "RowMajor", "TagC")
	93	+ self.check.test_kernel(kernel_str, "BasicMatmulTla")
	94	+ self.check.test_dispatch_policy(kernel_str, "MmadPingpong")
	95	+ self.check.test_arch_tag(kernel_str, Arch.Ascend950)
	96	+
	97	+ def test_batched_matmul(self):
	98	+ a = OpTensor.from_shape_stride(
	99	+ shape=(8, 128, 256),
	100	+ stride=(32768, 256, 1),
	101	+ dtype=DataType.FLOAT
	102	+ )
	103	+ b = OpTensor.from_shape_stride(
	104	+ shape=(8, 256, 384),
	105	+ stride=(98304, 384, 1),
	106	+ dtype=DataType.FLOAT
	107	+ )
	108	+ gemm_plan = Gemm(atlas_arch=Arch.Ascend950, element_C=DataType.FLOAT, core_num=8, A=a, B=b)
	109	+ kernels = gemm_plan.get_kernels()
	110	+
	111	+ batched_kernel = find_kernel_by_type(kernels, BatchedMatmulKernel)
	112	+ if batched_kernel is None:
	113	+ raise ValueError(f"No kernel named 'BatchedMatmulKernel' found, available kernel list: {[type(k).__name__ for k in kernels]}")
	114	+
	115	+ # 基础特征检查
	116	+ self.assertEqual(type(batched_kernel), BatchedMatmulKernel)
	117	+ self.assertFalse(batched_kernel.relu_enable) # default to False
	118	+
	119	+ batched_kernel.tune(
	120	+ GemmShape(128, 256, 64),
	121	+ GemmShape(128, 256, 64),
	122	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950, enable_unit_flag=True)
	123	+ )
	124	+
	125	+ kernel_str = batched_kernel.gen_kernel_template()
	126	+
	127	+ self.check.test_tileshape(kernel_str, GemmShape(128, 256, 64))
	128	+ self.check.test_tileshape(kernel_str, GemmShape(128, 256, 64), pos="L0")
	129	+ self.check.test_element_dtype(kernel_str, DataType.FLOAT)
	130	+ self.check.test_layout(kernel_str, "RowMajor")
	131	+ self.check.test_layout(kernel_str, "RowMajor", "TagB")
	132	+ self.check.test_layout(kernel_str, "RowMajor", "TagC")
	133	+ self.check.test_kernel(kernel_str, "BatchedMatmulTla")
	134	+ self.check.test_dispatch_policy(kernel_str, "MmadPingpong")
	135	+ self.check.test_arch_tag(kernel_str, Arch.Ascend950)
	136	+
	137	+ def test_basic_matmul_kernel_with_relu(self):
	138	+ a = OpTensor.from_shape_stride(
	139	+ shape=(128, 256),
	140	+ stride=(256, 1),
	141	+ dtype=DataType.FLOAT
	142	+ )
	143	+ b = OpTensor.from_shape_stride(
	144	+ shape=(256, 384),
	145	+ stride=(384, 1),
	146	+ dtype=DataType.FLOAT
	147	+ )
	148	+
	149	+ gemm_plan = Gemm(atlas_arch=Arch.Ascend950, element=DataType.FLOAT, layout=RowMajor, A=a, B=b)
	150	+ kernels = gemm_plan.get_kernels()
	151	+
	152	+ basic_kernel = find_kernel_by_type(kernels, BasicMatmulKernel)
	153	+ if basic_kernel is None:
	154	+ raise ValueError(f"No kernel named 'BasicMatmulKernel' found, available kernel list: {[type(k).__name__ for k in kernels]}")
	155	+
	156	+ self.assertEqual(type(basic_kernel), BasicMatmulKernel)
	157	+ self.assertFalse(basic_kernel.relu_enable)
	158	+
	159	+ basic_kernel.tune(
	160	+ GemmShape(128, 256, 64),
	161	+ GemmShape(128, 256, 64),
	162	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950),
	163	+ relu_enable=True
	164	+ )
	165	+ self.assertTrue(basic_kernel.relu_enable)
	166	+
	167	+ kernel_str = basic_kernel.gen_kernel_template()
	168	+
	169	+ self.check.test_tileshape(kernel_str, GemmShape(128, 256, 64))
	170	+ self.check.test_tileshape(kernel_str, GemmShape(128, 256, 64), pos="L0")
	171	+ self.check.test_element_dtype(kernel_str, DataType.FLOAT)
	172	+ self.check.test_layout(kernel_str, "RowMajor")
	173	+ self.check.test_layout(kernel_str, "RowMajor", "TagB")
	174	+ self.check.test_layout(kernel_str, "RowMajor", "TagC")
	175	+ self.check.test_kernel(kernel_str, "BasicMatmulTla")
	176	+ self.check.test_dispatch_policy(kernel_str, "MmadPingpong")
	177	+ self.check.test_arch_tag(kernel_str, Arch.Ascend950)
	178	+
	179	+ def test_basic_matmul_kernel_with_is_hf32(self):
	180	+ a = OpTensor.from_shape_stride(
	181	+ shape=(128, 256),
	182	+ stride=(256, 1),
	183	+ dtype=DataType.FLOAT
	184	+ )
	185	+ b = OpTensor.from_shape_stride(
	186	+ shape=(256, 384),
	187	+ stride=(384, 1),
	188	+ dtype=DataType.FLOAT
	189	+ )
	190	+
	191	+ gemm_plan = Gemm(atlas_arch=Arch.Ascend950, element=DataType.FLOAT, layout=RowMajor, A=a, B=b)
	192	+ kernels = gemm_plan.get_kernels()
	193	+
	194	+ basic_kernel = find_kernel_by_type(kernels, BasicMatmulKernel)
	195	+ if basic_kernel is None:
	196	+ raise ValueError(f"No kernel named 'BasicMatmulKernel' found, available kernel list: {[type(k).__name__ for k in kernels]}")
	197	+
	198	+ self.assertEqual(type(basic_kernel), BasicMatmulKernel)
	199	+
	200	+ basic_kernel.tune(
	201	+ GemmShape(128, 256, 64),
	202	+ GemmShape(128, 256, 64),
	203	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950),
	204	+ is_hf32=True
	205	+ )
	206	+ self.assertTrue(basic_kernel.dispatch_policy[0].use_hf32_mode)
	207	+
	208	+ kernel_str = basic_kernel.gen_kernel_template()
	209	+
	210	+ self.check.test_tileshape(kernel_str, GemmShape(128, 256, 64))
	211	+ self.check.test_tileshape(kernel_str, GemmShape(128, 256, 64), pos="L0")
	212	+ self.check.test_element_dtype(kernel_str, DataType.FLOAT)
	213	+ self.check.test_layout(kernel_str, "RowMajor")
	214	+ self.check.test_layout(kernel_str, "RowMajor", "TagB")
	215	+ self.check.test_layout(kernel_str, "RowMajor", "TagC")
	216	+ self.check.test_kernel(kernel_str, "BasicMatmulTla")
	217	+ self.check.test_dispatch_policy(kernel_str, "MmadPingpong")
	218	+ self.check.test_arch_tag(kernel_str, Arch.Ascend950)
	219	+
	220	+ def test_matmul_with_bias(self):
	221	+ a = OpTensor.from_shape_stride(
	222	+ shape=(128, 256),
	223	+ stride=(256, 1),
	224	+ dtype=DataType.FLOAT
	225	+ )
	226	+ b = OpTensor.from_shape_stride(
	227	+ shape=(256, 384),
	228	+ stride=(384, 1),
	229	+ dtype=DataType.FLOAT
	230	+ )
	231	+ bias = OpTensor.from_shape_stride(
	232	+ shape=(384,),
	233	+ stride=(1,),
	234	+ dtype=DataType.FLOAT
	235	+ )
	236	+
	237	+ gemm_plan = Gemm(
	238	+ atlas_arch=Arch.Ascend950,
	239	+ element=DataType.FLOAT,
	240	+ layout=RowMajor,
	241	+ A=a, B=b, Bias=bias
	242	+ )
	243	+ kernels = gemm_plan.get_kernels()
	244	+
	245	+ basic_kernel = find_kernel_by_type(kernels, BasicMatmulKernel)
	246	+ if basic_kernel is None:
	247	+ raise ValueError(f"No kernel named 'BasicMatmulKernel' found, available kernel list: {[type(k).__name__ for k in kernels]}")
	248	+
	249	+ self.assertEqual(type(basic_kernel), BasicMatmulKernel)
	250	+
	251	+ basic_kernel.tune(
	252	+ GemmShape(128, 256, 64),
	253	+ GemmShape(128, 256, 64),
	254	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950)
	255	+ )
	256	+
	257	+ params = basic_kernel.gen_params_device(def_mode=False)
	258	+ self.check.test_params(params, (
	259	+ "problemShape",
	260	+ "deviceA", "layoutA",
	261	+ "deviceB", "layoutB",
	262	+ "deviceC", "layoutC",
	263	+ "deviceBias"
	264	+ ))
	265	+
	266	+ kernel_str = basic_kernel.gen_kernel_template()
	267	+
	268	+ self.check.test_tileshape(kernel_str, GemmShape(128, 256, 64))
	269	+ self.check.test_tileshape(kernel_str, GemmShape(128, 256, 64), pos="L0")
	270	+ self.check.test_element_dtype(kernel_str, DataType.FLOAT)
	271	+ self.check.test_layout(kernel_str, "RowMajor")
	272	+ self.check.test_layout(kernel_str, "RowMajor", "TagB")
	273	+ self.check.test_layout(kernel_str, "RowMajor", "TagC")
	274	+ self.check.test_kernel(kernel_str, "BasicMatmulTla")
	275	+ self.check.test_dispatch_policy(kernel_str, "MmadPingpong")
	276	+ self.check.test_arch_tag(kernel_str, Arch.Ascend950)
	277	+
	278	+ def test_tile_shape_with_bias(self):
	279	+ a = OpTensor.from_shape_stride(
	280	+ shape=(128, 256),
	281	+ stride=(256, 1),
	282	+ dtype=DataType.FLOAT
	283	+ )
	284	+ b = OpTensor.from_shape_stride(
	285	+ shape=(256, 384),
	286	+ stride=(384, 1),
	287	+ dtype=DataType.FLOAT
	288	+ )
	289	+ bias = OpTensor.from_shape_stride(
	290	+ shape=(384,),
	291	+ stride=(1,),
	292	+ dtype=DataType.FLOAT
	293	+ )
	294	+
	295	+ gemm_plan_without_bias = Gemm(
	296	+ atlas_arch=Arch.Ascend950,
	297	+ element=DataType.FLOAT,
	298	+ layout=RowMajor,
	299	+ A=a, B=b
	300	+ )
	301	+ kernels_without_bias = gemm_plan_without_bias.get_kernels()
	302	+ basic_kernel_without_bias = find_kernel_by_type(kernels_without_bias, BasicMatmulKernel)
	303	+ if basic_kernel_without_bias is None:
	304	+ raise ValueError(f"No kernel named 'BasicMatmulKernel' found, available kernel list: {[type(k).__name__ for k in kernels_without_bias]}")
	305	+
	306	+ default_shape_no_bias = basic_kernel_without_bias.get_default_tile_shape()
	307	+ self.assertEqual(default_shape_no_bias, (GemmShape(256, 256, 128), GemmShape(256, 256, 32)))
	308	+
	309	+ gemm_plan_with_bias = Gemm(
	310	+ atlas_arch=Arch.Ascend950,
	311	+ element=DataType.FLOAT,
	312	+ layout=RowMajor,
	313	+ A=a, B=b, Bias=bias
	314	+ )
	315	+ kernels_with_bias = gemm_plan_with_bias.get_kernels()
	316	+ basic_kernel_with_bias = find_kernel_by_type(kernels_with_bias, BasicMatmulKernel)
	317	+ if basic_kernel_with_bias is None:
	318	+ raise ValueError(f"No kernel named 'BasicMatmulKernel' found, available kernel list: {[type(k).__name__ for k in kernels_with_bias]}")
	319	+
	320	+ default_shape_with_bias = basic_kernel_with_bias.get_default_tile_shape()
	321	+ self.assertEqual(default_shape_with_bias, (GemmShape(240, 256, 128), GemmShape(240, 256, 32)))
	322	+
	323	+ def test_streamk_matmul(self):
	324	+ a = OpTensor.from_shape_stride(
	325	+ shape=(1280, 3000),
	326	+ stride=(3000, 1),
	327	+ dtype=DataType.FLOAT
	328	+ )
	329	+ b = OpTensor.from_shape_stride(
	330	+ shape=(3000, 256),
	331	+ stride=(256, 1),
	332	+ dtype=DataType.FLOAT
	333	+ )
	334	+
	335	+ gemm_plan = Gemm(
	336	+ atlas_arch=Arch.Ascend950,
	337	+ element=DataType.FLOAT,
	338	+ layout=RowMajor,
	339	+ core_num=8,
	340	+ A=a, B=b
	341	+ )
	342	+ kernels = gemm_plan.get_kernels()
	343	+
	344	+ streamk_kernel = find_kernel_by_type(kernels, StreamkMatmulKernel)
	345	+ if streamk_kernel is None:
	346	+ raise ValueError(f"No kernel named 'StreamkMatmulKernel' found, available kernel list: {[type(k).__name__ for k in kernels]}")
	347	+
	348	+ self.assertEqual(type(streamk_kernel), StreamkMatmulKernel)
	349	+ self.assertEqual(streamk_kernel.slice_axis, "K")
	350	+ self.assertFalse(streamk_kernel.relu_enable)
	351	+
	352	+ streamk_kernel.tune(
	353	+ GemmShape(256, 256, 128),
	354	+ GemmShape(256, 256, 32),
	355	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950, enable_unit_flag=True)
	356	+ )
	357	+
	358	+ kernel_str = streamk_kernel.gen_kernel_template()
	359	+
	360	+ self.check.test_tileshape(kernel_str, GemmShape(256, 256, 128))
	361	+ self.check.test_tileshape(kernel_str, GemmShape(256, 256, 32), pos="L0")
	362	+ self.check.test_element_dtype(kernel_str, DataType.FLOAT)
	363	+ self.check.test_layout(kernel_str, "RowMajor")
	364	+ self.check.test_layout(kernel_str, "RowMajor", "TagB")
	365	+ self.check.test_layout(kernel_str, "RowMajor", "TagC")
	366	+ self.check.test_kernel(kernel_str, "StreamkMatmulTla")
	367	+ self.check.test_dispatch_policy(kernel_str, "MmadPingpong")
	368	+ self.check.test_arch_tag(kernel_str, Arch.Ascend950)
	369	+
	370	+ def test_multi_core_splitk_matmul_kernel(self):
	371	+ a = OpTensor.from_shape_stride(
	372	+ shape=(256, 3000),
	373	+ stride=(3000, 1),
	374	+ dtype=DataType.FLOAT
	375	+ )
	376	+ b = OpTensor.from_shape_stride(
	377	+ shape=(3000, 256),
	378	+ stride=(256, 1),
	379	+ dtype=DataType.FLOAT
	380	+ )
	381	+
	382	+ gemm_plan = Gemm(
	383	+ atlas_arch=Arch.Ascend950,
	384	+ element=DataType.FLOAT,
	385	+ layout=RowMajor,
	386	+ core_num=8,
	387	+ A=a, B=b
	388	+ )
	389	+ kernels = gemm_plan.get_kernels()
	390	+
	391	+ splitk_kernel = find_kernel_by_type(kernels, MultiCoreSplitkMatmulKernel)
	392	+ if splitk_kernel is None:
	393	+ raise ValueError(f"No kernel named 'MultiCoreSplitkMatmulKernel' found, available kernel list: {[type(k).__name__ for k in kernels]}")
	394	+
	395	+ self.assertEqual(type(splitk_kernel), MultiCoreSplitkMatmulKernel)
	396	+ self.assertEqual(splitk_kernel.slice_axis, "K")
	397	+ self.assertFalse(splitk_kernel.relu_enable)
	398	+
	399	+ splitk_kernel.tune(
	400	+ GemmShape(256, 256, 128),
	401	+ GemmShape(256, 256, 32),
	402	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950, enable_unit_flag=True)
	403	+ )
	404	+
	405	+ kernel_str = splitk_kernel.gen_kernel_template()
	406	+
	407	+ self.check.test_tileshape(kernel_str, GemmShape(256, 256, 128))
	408	+ self.check.test_tileshape(kernel_str, GemmShape(256, 256, 32), pos="L0")
	409	+ self.check.test_element_dtype(kernel_str, DataType.FLOAT)
	410	+ self.check.test_layout(kernel_str, "RowMajor")
	411	+ self.check.test_layout(kernel_str, "RowMajor", "TagB")
	412	+ self.check.test_layout(kernel_str, "RowMajor", "TagC")
	413	+ self.check.test_kernel(kernel_str, "MultiCoreSplitkMatmulTla")
	414	+ self.check.test_dispatch_policy(kernel_str, "MmadPingpong")
	415	+ self.check.test_arch_tag(kernel_str, Arch.Ascend950)
	416	+
	417	+ def test_tail_multi_core_splitk_matmul_kernel(self):
	418	+ a = OpTensor.from_shape_stride(
	419	+ shape=(768, 3000),
	420	+ stride=(3000, 1),
	421	+ dtype=DataType.FLOAT
	422	+ )
	423	+ b = OpTensor.from_shape_stride(
	424	+ shape=(3000, 768),
	425	+ stride=(768, 1),
	426	+ dtype=DataType.FLOAT
	427	+ )
	428	+
	429	+ gemm_plan = Gemm(
	430	+ atlas_arch=Arch.Ascend950,
	431	+ element=DataType.FLOAT,
	432	+ layout=RowMajor,
	433	+ core_num=8,
	434	+ A=a, B=b
	435	+ )
	436	+ kernels = gemm_plan.get_kernels()
	437	+
	438	+ tail_splitk_kernel = find_kernel_by_type(kernels, TailMultiCoreSplitkMatmulKernel)
	439	+ if tail_splitk_kernel is None:
	440	+ raise ValueError(f"No kernel named 'TailMultiCoreSplitkMatmulKernel' found, available kernel list: {[type(k).__name__ for k in kernels]}")
	441	+
	442	+ self.assertEqual(type(tail_splitk_kernel), TailMultiCoreSplitkMatmulKernel)
	443	+ self.assertEqual(tail_splitk_kernel.slice_axis, "K")
	444	+ self.assertFalse(tail_splitk_kernel.relu_enable)
	445	+
	446	+ tail_splitk_kernel.tune(
	447	+ GemmShape(256, 256, 128),
	448	+ GemmShape(256, 256, 32),
	449	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950, enable_unit_flag=True)
	450	+ )
	451	+
	452	+ kernel_str = tail_splitk_kernel.gen_kernel_template()
	453	+
	454	+ self.check.test_tileshape(kernel_str, GemmShape(256, 256, 128))
	455	+ self.check.test_tileshape(kernel_str, GemmShape(256, 256, 32), pos="L0")
	456	+ self.check.test_element_dtype(kernel_str, DataType.FLOAT)
	457	+ self.check.test_layout(kernel_str, "RowMajor")
	458	+ self.check.test_layout(kernel_str, "RowMajor", "TagB")
	459	+ self.check.test_layout(kernel_str, "RowMajor", "TagC")
	460	+ self.check.test_kernel(kernel_str, "TailMultiCoreSplitkMatmulTla")
	461	+ self.check.test_dispatch_policy(kernel_str, "MmadPingpong")
	462	+ self.check.test_arch_tag(kernel_str, Arch.Ascend950)
	463	+
	464	+ def test_basic_matmul_tla_visitor_kernel(self):
	465	+ a = OpTensor.from_shape_stride(
	466	+ shape=(128, 256),
	467	+ stride=(256, 1),
	468	+ dtype=DataType.FLOAT
	469	+ )
	470	+ b = OpTensor.from_shape_stride(
	471	+ shape=(256, 384),
	472	+ stride=(384, 1),
	473	+ dtype=DataType.FLOAT
	474	+ )
	475	+
	476	+ function_source = """
	477	+def epilogue(accum):
	478	+ temp = accum
	479	+ return temp
	480	+"""
	481	+ example_inputs = {
	482	+ "accum": OpTensor.from_shape_stride(
	483	+ shape=(128, 384),
	484	+ stride=(384, 1),
	485	+ dtype=DataType.FLOAT
	486	+ ),
	487	+ "temp": OpTensor.from_shape_stride(
	488	+ shape=(128, 384),
	489	+ stride=(384, 1),
	490	+ dtype=DataType.FLOAT
	491	+ ),
	492	+ }
	493	+
	494	+ evg_config = {
	495	+ "fn_src": function_source,
	496	+ "example_inputs": example_inputs,
	497	+ }
	498	+
	499	+ gemm_plan = Gemm(
	500	+ atlas_arch=Arch.Ascend950,
	501	+ evg_config=evg_config,
	502	+ element=DataType.FLOAT,
	503	+ layout=RowMajor,
	504	+ A=a, B=b
	505	+ )
	506	+ kernels = gemm_plan.get_kernels()
	507	+
	508	+ visitor_kernel = find_kernel_by_type(kernels, BasicMatmulTlaVisitorKernel)
	509	+ if visitor_kernel is None:
	510	+ raise ValueError(f"No kernel named 'BasicMatmulTlaVisitorKernel' found, available kernel list: {[type(k).__name__ for k in kernels]}")
	511	+
	512	+ self.assertEqual(type(visitor_kernel), BasicMatmulTlaVisitorKernel)
	513	+ self.assertEqual(visitor_kernel.slice_axis, None)
	514	+ self.assertFalse(visitor_kernel.relu_enable)
	515	+
	516	+ visitor_kernel.tune(
	517	+ GemmShape(128, 256, 64),
	518	+ GemmShape(128, 256, 64),
	519	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950)
	520	+ )
	521	+
	522	+ kernel_str = visitor_kernel.gen_kernel_template()
	523	+
	524	+ self.check.test_tileshape(kernel_str, GemmShape(128, 256, 64))
	525	+ self.check.test_tileshape(kernel_str, GemmShape(128, 256, 64), pos="L0")
	526	+ self.check.test_element_dtype(kernel_str, DataType.FLOAT)
	527	+ self.check.test_layout(kernel_str, "RowMajor")
	528	+ self.check.test_layout(kernel_str, "RowMajor", "TagB")
	529	+ self.check.test_layout(kernel_str, "RowMajor", "TagC")
	530	+ self.check.test_kernel(kernel_str, "BasicMatmulTlaVisitor")
	531	+ self.check.test_dispatch_policy(kernel_str, "MmadPingpong")
	532	+ self.check.test_arch_tag(kernel_str, Arch.Ascend950)
	533	+
	534	+ def test_basic_matmul_tla_visitor_kernel_with_constant(self):
	535	+ a = OpTensor.from_shape_stride(
	536	+ shape=(128, 256),
	537	+ stride=(256, 1),
	538	+ dtype=DataType.FLOAT
	539	+ )
	540	+ b = OpTensor.from_shape_stride(
	541	+ shape=(256, 384),
	542	+ stride=(384, 1),
	543	+ dtype=DataType.FLOAT
	544	+ )
	545	+
	546	+ function_source = """
	547	+def epilogue(accum):
	548	+ scale = constant(0.1, "float")
	549	+ result = accum * scale
	550	+ return result
	551	+"""
	552	+ example_inputs = {
	553	+ "accum": OpTensor.from_shape_stride(
	554	+ shape=(128, 384),
	555	+ stride=(384, 1),
	556	+ dtype=DataType.FLOAT
	557	+ ),
	558	+ "result": OpTensor.from_shape_stride(
	559	+ shape=(128, 384),
	560	+ stride=(384, 1),
	561	+ dtype=DataType.FLOAT
	562	+ ),
	563	+ }
	564	+
	565	+ evg_config = {
	566	+ "fn_src": function_source,
	567	+ "example_inputs": example_inputs,
	568	+ }
	569	+
	570	+ gemm_plan = Gemm(
	571	+ atlas_arch=Arch.Ascend950,
	572	+ evg_config=evg_config,
	573	+ element=DataType.FLOAT,
	574	+ layout=RowMajor,
	575	+ A=a, B=b
	576	+ )
	577	+ kernels = gemm_plan.get_kernels()
	578	+
	579	+ visitor_kernel = find_kernel_by_type(kernels, BasicMatmulTlaVisitorKernel)
	580	+ if visitor_kernel is None:
	581	+ raise ValueError(f"No kernel named 'BasicMatmulTlaVisitorKernel' found, available kernel list: {[type(k).__name__ for k in kernels]}")
	582	+
	583	+ self.assertEqual(type(visitor_kernel), BasicMatmulTlaVisitorKernel)
	584	+
	585	+ visitor_kernel.tune(
	586	+ GemmShape(128, 256, 64),
	587	+ GemmShape(128, 256, 64),
	588	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950)
	589	+ )
	590	+
	591	+ kernel_str = visitor_kernel.gen_kernel_template()
	592	+
	593	+ self.check.test_tileshape(kernel_str, GemmShape(128, 256, 64))
	594	+ self.check.test_tileshape(kernel_str, GemmShape(128, 256, 64), pos="L0")
	595	+ self.check.test_element_dtype(kernel_str, DataType.FLOAT)
	596	+ self.check.test_layout(kernel_str, "RowMajor")
	597	+ self.check.test_layout(kernel_str, "RowMajor", "TagB")
	598	+ self.check.test_layout(kernel_str, "RowMajor", "TagC")
	599	+ self.check.test_kernel(kernel_str, "BasicMatmulTlaVisitor")
	600	+ self.check.test_dispatch_policy(kernel_str, "MmadPingpong")
	601	+ self.check.test_arch_tag(kernel_str, Arch.Ascend950)
	602	+
	603	+ def test_to_evg_method(self):
	604	+ a = OpTensor.from_shape_stride(
	605	+ shape=(128, 256),
	606	+ stride=(256, 1),
	607	+ dtype=DataType.FLOAT
	608	+ )
	609	+ b = OpTensor.from_shape_stride(
	610	+ shape=(256, 384),
	611	+ stride=(384, 1),
	612	+ dtype=DataType.FLOAT
	613	+ )
	614	+
	615	+ gemm_plan = Gemm(
	616	+ atlas_arch=Arch.Ascend950,
	617	+ element=DataType.FLOAT,
	618	+ layout=RowMajor,
	619	+ A=a, B=b
	620	+ )
	621	+ kernels = gemm_plan.get_kernels()
	622	+
	623	+ basic_kernel = find_kernel_by_type(kernels, BasicMatmulKernel)
	624	+ if basic_kernel is None:
	625	+ raise ValueError(f"No kernel named 'BasicMatmulKernel' found, available kernel list: {[type(k).__name__ for k in kernels]}")
	626	+
	627	+ self.assertEqual(type(basic_kernel), BasicMatmulKernel)
	628	+
	629	+ is_support_evg = bool(getattr(basic_kernel, "is_support_evg", False))
	630	+ self.assertTrue(is_support_evg)
	631	+
	632	+ function_source = """
	633	+def epilogue(accum, bias):
	634	+ result = accum + bias
	635	+ return result
	636	+"""
	637	+ example_inputs = {
	638	+ "accum": OpTensor.from_shape_stride(
	639	+ shape=(128, 384),
	640	+ stride=(384, 1),
	641	+ dtype=DataType.FLOAT
	642	+ ),
	643	+ "bias": OpTensor.from_shape_stride(
	644	+ shape=(1, 384),
	645	+ stride=(384, 1),
	646	+ dtype=DataType.FLOAT
	647	+ ),
	648	+ "result": OpTensor.from_shape_stride(
	649	+ shape=(128, 384),
	650	+ stride=(384, 1),
	651	+ dtype=DataType.FLOAT
	652	+ ),
	653	+ }
	654	+
	655	+ evg_config = {
	656	+ "fn_src": function_source,
	657	+ "example_inputs": example_inputs,
	658	+ }
	659	+
	660	+ basic_kernel.tune(
	661	+ GemmShape(128, 256, 64),
	662	+ GemmShape(128, 256, 64),
	663	+ dispatch_policy=MmadPingpong(arch_tag=Arch.Ascend950)
	664	+ )
	665	+
	666	+ evg_kernel = basic_kernel.to_evg(evg_config)
	667	+ self.assertIsNotNone(evg_kernel)
	668	+ self.assertIsInstance(evg_kernel, BasicMatmulTlaVisitorKernel)
	669	+ self.assertIsNotNone(evg_kernel.evg)
	670	+
	671	+ kernel_str = evg_kernel.gen_kernel_template()
	672	+
	673	+ self.check.test_tileshape(kernel_str, GemmShape(128, 256, 64))
	674	+ self.check.test_tileshape(kernel_str, GemmShape(128, 256, 64), pos="L0")
	675	+ self.check.test_element_dtype(kernel_str, DataType.FLOAT)
	676	+ self.check.test_layout(kernel_str, "RowMajor")
	677	+ self.check.test_layout(kernel_str, "RowMajor", "TagB")
	678	+ self.check.test_layout(kernel_str, "RowMajor", "TagC")
	679	+ self.check.test_kernel(kernel_str, "BasicMatmulTlaVisitor")
	680	+ self.check.test_dispatch_policy(kernel_str, "MmadPingpong")
	681	+ self.check.test_arch_tag(kernel_str, Arch.Ascend950)
	682	+
	683	+ def test_to_evg_unsupported_kernel(self):
	684	+ a = OpTensor.from_shape_stride(
	685	+ shape=(256, 3000),
	686	+ stride=(3000, 1),
	687	+ dtype=DataType.FLOAT
	688	+ )
	689	+ b = OpTensor.from_shape_stride(
	690	+ shape=(3000, 256),
	691	+ stride=(256, 1),
	692	+ dtype=DataType.FLOAT
	693	+ )
	694	+
	695	+ gemm_plan = Gemm(
	696	+ atlas_arch=Arch.Ascend950,
	697	+ element=DataType.FLOAT,
	698	+ layout=RowMajor,
	699	+ core_num=8,
	700	+ A=a, B=b
	701	+ )
	702	+ kernels = gemm_plan.get_kernels()
	703	+
	704	+ splitk_kernel = find_kernel_by_type(kernels, MultiCoreSplitkMatmulKernel)
	705	+ if splitk_kernel is None:
	706	+ raise ValueError(f"No kernel named 'MultiCoreSplitkMatmulKernel' found, available kernel list: {[type(k).__name__ for k in kernels]}")
	707	+
	708	+ self.assertEqual(type(splitk_kernel), MultiCoreSplitkMatmulKernel)
	709	+
	710	+ is_support_evg = bool(getattr(splitk_kernel, "is_support_evg", False))
	711	+ self.assertFalse(is_support_evg)
	712	+
	713	+ evg_config = {
	714	+ "fn_src": "def epilogue(accum): return accum\n",
	715	+ "example_inputs": {
	716	+ "accum": OpTensor.from_shape_stride(
	717	+ shape=(256, 256),
	718	+ stride=(256, 1),
	719	+ dtype=DataType.FLOAT
	720	+ ),
	721	+ "result": OpTensor.from_shape_stride(
	722	+ shape=(256, 256),
	723	+ stride=(256, 1),
	724	+ dtype=DataType.FLOAT
	725	+ ),
	726	+ },
	727	+ }
	728	+
	729	+ result = splitk_kernel.to_evg(evg_config)
	730	+ self.assertIsNone(result)
	731	+
	732	+######################################
	733	+
	734	+
	735	+if __name__ == "__main__":
	736	+ unittest.main()

Apython/catlass_cppgen/tests/op/test_group_gemm.py+184-0

@@ -0,0 +1,184 @@
	1	+# This program is free software, you can redistribute it and/or modify.
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This file is a part of the CANN Open Software.
	4	+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	7	+# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	8	+# the software repository for the full text of the License.
	9	+
	10	+import unittest
	11	+import re
	12	+from catlass_cppgen.op.group_gemm import GroupGemm
	13	+
	14	+from catlass_cppgen.common.data_type import DataType
	15	+from catlass_cppgen.catlass.layout.layout import RowMajor, ColumnMajor, VectorLayout
	16	+from catlass_cppgen.catlass.gemm_coord import GemmShape
	17	+from catlass_cppgen.catlass.arch.arch import Arch
	18	+from catlass_cppgen.kernel.group_gemm.grouped_matmul_slice_m import GroupedMatmulSliceMKernel
	19	+from catlass_cppgen.common.op_tensor import OpTensor
	20	+
	21	+from assertion_helper import TestAssertions
	22	+
	23	+
	24	+class TestGroupGemm(unittest.TestCase):
	25	+ def setUp(self):
	26	+ self.check = TestAssertions(self)
	27	+
	28	+ def test_group_gemm(self):
	29	+ problemCount = 4
	30	+ M, K, N = 128, 256, 384
	31	+ a = OpTensor.from_shape_stride(
	32	+ shape=(M, K),
	33	+ stride=(K, 1),
	34	+ dtype=DataType.FLOAT
	35	+ )
	36	+ # B 是 3D tensor: [problemCount, k, n] = [4, 256, 384]
	37	+ b = OpTensor.from_shape_stride(
	38	+ shape=(problemCount, K, N),
	39	+ stride=(K * N, N, 1),
	40	+ dtype=DataType.FLOAT
	41	+ )
	42	+ # 创建 groupList OpTensor（一维，int64_t 类型）
	43	+ groupList = OpTensor(
	44	+ dtype=DataType.INT64,
	45	+ layout=VectorLayout(problemCount), # 4 个 group
	46	+ shape=(problemCount,)
	47	+ )
	48	+
	49	+ group_gemm_plan = GroupGemm(
	50	+ atlas_arch=Arch.Ascend950, element=DataType.FLOAT, layout=RowMajor,
	51	+ core_num=8, A=a, B=b, groupList=groupList
	52	+ )
	53	+ kernels = group_gemm_plan.get_kernels()
	54	+ self.assertEqual(len(kernels), 1)
	55	+ self.assertIsInstance(kernels[0], GroupedMatmulSliceMKernel)
	56	+
	57	+ group_gemm_kernel = kernels[0]
	58	+ group_gemm_kernel.tune(GemmShape(256, 256, 128), GemmShape(256, 256, 64))
	59	+
	60	+ kernel_str = group_gemm_kernel.gen_kernel_template()
	61	+ input_str = group_gemm_kernel.gen_input_template()
	62	+ self.check.test_tileshape(kernel_str, GemmShape(256, 256, 128))
	63	+ self.check.test_tileshape(kernel_str, GemmShape(256, 256, 64), pos="L0")
	64	+ self.check.test_element_dtype(kernel_str, DataType.FLOAT)
	65	+ self.check.test_kernel(kernel_str, "GroupedMatmulSliceMTla")
	66	+ self._test_problem_count(input_str, problemCount)
	67	+
	68	+ self.check.test_dispatch_policy(kernel_str, "MmadPingpong")
	69	+ self.check.test_arch_tag(kernel_str, Arch.Ascend950)
	70	+
	71	+ def test_group_gemm_codegen_float16(self):
	72	+ groupList = OpTensor.from_shape_stride(shape=(2,), stride=(1,), dtype=DataType.INT64)
	73	+ A = OpTensor.from_shape_stride(shape=(64, 512), stride=(512, 1), dtype=DataType.FLOAT16)
	74	+ B = OpTensor.from_shape_stride(
	75	+ shape=(2, 512, 256), stride=(131072, 256, 1), dtype=DataType.FLOAT16
	76	+ )
	77	+
	78	+ gemm_plan = GroupGemm(
	79	+ atlas_arch=Arch.AtlasA2, element=DataType.FLOAT16, layout=RowMajor,
	80	+ core_num=8, A=A, B=B, groupList=groupList
	81	+ )
	82	+ kernels = gemm_plan.get_kernels()
	83	+ gemm_kernel = kernels[0]
	84	+ gemm_kernel.tune(GemmShape(256, 256, 128), GemmShape(256, 256, 64))
	85	+
	86	+ kernel_str = gemm_kernel.gen_kernel_template()
	87	+ self.check.test_element_dtype(kernel_str, DataType.FLOAT16)
	88	+
	89	+ def test_group_gemm_codegen_column_major(self):
	90	+ groupList = OpTensor.from_shape_stride(shape=(2,), stride=(1,), dtype=DataType.INT64)
	91	+ A = OpTensor.from_shape_stride(shape=(64, 512), stride=(1, 64), dtype=DataType.FLOAT)
	92	+ B = OpTensor.from_shape_stride(
	93	+ shape=(2, 512, 256), stride=(131072, 1, 512), dtype=DataType.FLOAT
	94	+ )
	95	+
	96	+ gemm_plan = GroupGemm(
	97	+ atlas_arch=Arch.AtlasA2, element=DataType.FLOAT, layout=ColumnMajor,
	98	+ core_num=8, A=A, B=B, groupList=groupList
	99	+ )
	100	+ kernels = gemm_plan.get_kernels()
	101	+ self.assertEqual(len(kernels), 1)
	102	+ self.assertIsInstance(kernels[0], GroupedMatmulSliceMKernel)
	103	+
	104	+ def test_group_gemm_missing_grouplist(self):
	105	+ A = OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT)
	106	+ B = OpTensor.from_shape_stride(shape=(256, 512), stride=(512, 1), dtype=DataType.FLOAT)
	107	+
	108	+ with self.assertRaises(ValueError):
	109	+ GroupGemm(
	110	+ atlas_arch=Arch.AtlasA2, element=DataType.FLOAT, layout=RowMajor,
	111	+ core_num=8, A=A, B=B
	112	+ )
	113	+
	114	+ def test_group_gemm_grouplist_shape_mismatch(self):
	115	+ groupList = OpTensor.from_shape_stride(shape=(3,), stride=(1,), dtype=DataType.INT64)
	116	+ A = OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT)
	117	+ B = OpTensor.from_shape_stride(
	118	+ shape=(2, 256, 512), stride=(131072, 512, 1), dtype=DataType.FLOAT
	119	+ )
	120	+
	121	+ with self.assertRaises(ValueError):
	122	+ GroupGemm(
	123	+ atlas_arch=Arch.AtlasA2, element=DataType.FLOAT, layout=RowMajor,
	124	+ core_num=8, A=A, B=B, groupList=groupList
	125	+ )
	126	+
	127	+ def test_group_gemm_a_not_2d(self):
	128	+ groupList = OpTensor.from_shape_stride(shape=(2,), stride=(1,), dtype=DataType.INT64)
	129	+ A = OpTensor.from_shape_stride(
	130	+ shape=(2, 128, 256), stride=(32768, 256, 1), dtype=DataType.FLOAT
	131	+ )
	132	+ B = OpTensor.from_shape_stride(
	133	+ shape=(2, 256, 512), stride=(131072, 512, 1), dtype=DataType.FLOAT
	134	+ )
	135	+
	136	+ with self.assertRaises(ValueError):
	137	+ GroupGemm(
	138	+ atlas_arch=Arch.AtlasA2, element=DataType.FLOAT, layout=RowMajor,
	139	+ core_num=8, A=A, B=B, groupList=groupList
	140	+ )
	141	+
	142	+ def test_group_gemm_can_implement(self):
	143	+ groupList = OpTensor.from_shape_stride(shape=(2,), stride=(1,), dtype=DataType.INT64)
	144	+ A = OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT)
	145	+ B = OpTensor.from_shape_stride(
	146	+ shape=(2, 256, 512), stride=(131072, 512, 1), dtype=DataType.FLOAT
	147	+ )
	148	+
	149	+ gemm_plan = GroupGemm(
	150	+ atlas_arch=Arch.AtlasA2, element=DataType.FLOAT, layout=RowMajor,
	151	+ core_num=8, A=A, B=B, groupList=groupList
	152	+ )
	153	+ self.assertTrue(gemm_plan.can_implement())
	154	+
	155	+ gemm_plan2 = GroupGemm(
	156	+ atlas_arch=Arch.AtlasA2, element=DataType.FLOAT, layout=RowMajor,
	157	+ core_num=8, A=A, B=B, groupList=groupList, alpha=2.0, beta=0.5
	158	+ )
	159	+ self.assertFalse(gemm_plan2.can_implement())
	160	+
	161	+ def test_group_gemm_includes(self):
	162	+ groupList = OpTensor.from_shape_stride(shape=(2,), stride=(1,), dtype=DataType.INT64)
	163	+ A = OpTensor.from_shape_stride(shape=(128, 256), stride=(256, 1), dtype=DataType.FLOAT)
	164	+ B = OpTensor.from_shape_stride(
	165	+ shape=(2, 256, 512), stride=(131072, 512, 1), dtype=DataType.FLOAT
	166	+ )
	167	+
	168	+ gemm_plan = GroupGemm(
	169	+ atlas_arch=Arch.AtlasA2, element=DataType.FLOAT, layout=RowMajor,
	170	+ core_num=8, A=A, B=B, groupList=groupList
	171	+ )
	172	+ kernels = gemm_plan.get_kernels()
	173	+ includes = kernels[0].gen_includes()
	174	+ self.assertIn("catlass/catlass.hpp", includes)
	175	+ self.assertIn("catlass/gemm/kernel/grouped_matmul_slice_m_tla.hpp", includes)
	176	+
	177	+ def _test_problem_count(self, template_str: str, expected_count: int):
	178	+ match = re.search(r"uint32_t\s+problemCount\s=\s(\d+);", template_str)
	179	+ self.assertIsNotNone(match, "problemCount not found in template")
	180	+ self.assertEqual(int(match.group(1)), expected_count)
	181	+
	182	+
	183	+if __name__ == "__main__":
	184	+ unittest.main()

Apython/catlass_cppgen/uv.lock+332-0

@@ -0,0 +1,332 @@
	1	+version = 1
	2	+revision = 3
	3	+requires-python = ">=3.11"
	4	+resolution-markers = [
	5	+ "python_full_version >= '3.13' and sys_platform != 'darwin'",
	6	+ "python_full_version == '3.12.*' and sys_platform != 'darwin'",
	7	+ "python_full_version < '3.12' and sys_platform != 'darwin'",
	8	+ "python_full_version >= '3.13' and sys_platform == 'darwin'",
	9	+ "python_full_version == '3.12.*' and sys_platform == 'darwin'",
	10	+ "python_full_version < '3.12' and sys_platform == 'darwin'",
	11	+]
	12	+
	13	+[[package]]
	14	+name = "catlass-cppgen"
	15	+version = "0.1.0"
	16	+source = { editable = "." }
	17	+dependencies = [
	18	+ { name = "ml-dtypes", version = "0.4.1", source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" }, marker = "python_full_version >= '3.13'" },
	19	+ { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" }, marker = "python_full_version < '3.13'" },
	20	+ { name = "numpy" },
	21	+ { name = "torch", version = "2.6.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform == 'darwin'" },
	22	+ { name = "torch", version = "2.6.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform != 'darwin'" },
	23	+]
	24	+
	25	+[package.metadata]
	26	+requires-dist = [
	27	+ { name = "ml-dtypes", specifier = ">=0.4.1" },
	28	+ { name = "numpy", specifier = "==1.26.4" },
	29	+ { name = "torch", specifier = "==2.6.0", index = "https://download.pytorch.org/whl/cpu" },
	30	+]
	31	+
	32	+[[package]]
	33	+name = "filelock"
	34	+version = "3.20.0"
	35	+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" }
	36	+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/58/46/0028a82567109b5ef6e4d2a1f04a583fb513e6cf9527fcdd09afd817deeb/filelock-3.20.0.tar.gz", hash = "sha256:711e943b4ec6be42e1d4e6690b48dc175c822967466bb31c0c293f34334c13f4", size = 18922, upload-time = "2025-10-08T18:03:50.056Z" }
	37	+wheels = [
	38	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/76/91/7216b27286936c16f5b4d0c530087e4a54eead683e6b0b73dd0c64844af6/filelock-3.20.0-py3-none-any.whl", hash = "sha256:339b4732ffda5cd79b13f4e2711a31b0365ce445d95d243bb996273d072546a2", size = 16054, upload-time = "2025-10-08T18:03:48.35Z" },
	39	+]
	40	+
	41	+[[package]]
	42	+name = "fsspec"
	43	+version = "2025.10.0"
	44	+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" }
	45	+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/24/7f/2747c0d332b9acfa75dc84447a066fdf812b5a6b8d30472b74d309bfe8cb/fsspec-2025.10.0.tar.gz", hash = "sha256:b6789427626f068f9a83ca4e8a3cc050850b6c0f71f99ddb4f542b8266a26a59", size = 309285, upload-time = "2025-10-30T14:58:44.036Z" }
	46	+wheels = [
	47	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/eb/02/a6b21098b1d5d6249b7c5ab69dde30108a71e4e819d4a9778f1de1d5b70d/fsspec-2025.10.0-py3-none-any.whl", hash = "sha256:7c7712353ae7d875407f97715f0e1ffcc21e33d5b24556cb1e090ae9409ec61d", size = 200966, upload-time = "2025-10-30T14:58:42.53Z" },
	48	+]
	49	+
	50	+[[package]]
	51	+name = "jinja2"
	52	+version = "3.1.6"
	53	+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" }
	54	+dependencies = [
	55	+ { name = "markupsafe" },
	56	+]
	57	+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" }
	58	+wheels = [
	59	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" },
	60	+]
	61	+
	62	+[[package]]
	63	+name = "markupsafe"
	64	+version = "3.0.3"
	65	+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" }
	66	+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" }
	67	+wheels = [
	68	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/08/db/fefacb2136439fc8dd20e797950e749aa1f4997ed584c62cfb8ef7c2be0e/markupsafe-3.0.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:1cc7ea17a6824959616c525620e387f6dd30fec8cb44f649e31712db02123dad", size = 11631, upload-time = "2025-09-27T18:36:18.185Z" },
	69	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e1/2e/5898933336b61975ce9dc04decbc0a7f2fee78c30353c5efba7f2d6ff27a/markupsafe-3.0.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:4bd4cd07944443f5a265608cc6aab442e4f74dff8088b0dfc8238647b8f6ae9a", size = 12058, upload-time = "2025-09-27T18:36:19.444Z" },
	70	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1d/09/adf2df3699d87d1d8184038df46a9c80d78c0148492323f4693df54e17bb/markupsafe-3.0.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6b5420a1d9450023228968e7e6a9ce57f65d148ab56d2313fcd589eee96a7a50", size = 24287, upload-time = "2025-09-27T18:36:20.768Z" },
	71	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/30/ac/0273f6fcb5f42e314c6d8cd99effae6a5354604d461b8d392b5ec9530a54/markupsafe-3.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0bf2a864d67e76e5c9a34dc26ec616a66b9888e25e7b9460e1c76d3293bd9dbf", size = 22940, upload-time = "2025-09-27T18:36:22.249Z" },
	72	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/19/ae/31c1be199ef767124c042c6c3e904da327a2f7f0cd63a0337e1eca2967a8/markupsafe-3.0.3-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:bc51efed119bc9cfdf792cdeaa4d67e8f6fcccab66ed4bfdd6bde3e59bfcbb2f", size = 21887, upload-time = "2025-09-27T18:36:23.535Z" },
	73	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b2/76/7edcab99d5349a4532a459e1fe64f0b0467a3365056ae550d3bcf3f79e1e/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:068f375c472b3e7acbe2d5318dea141359e6900156b5b2ba06a30b169086b91a", size = 23692, upload-time = "2025-09-27T18:36:24.823Z" },
	74	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a4/28/6e74cdd26d7514849143d69f0bf2399f929c37dc2b31e6829fd2045b2765/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:7be7b61bb172e1ed687f1754f8e7484f1c8019780f6f6b0786e76bb01c2ae115", size = 21471, upload-time = "2025-09-27T18:36:25.95Z" },
	75	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/62/7e/a145f36a5c2945673e590850a6f8014318d5577ed7e5920a4b3448e0865d/markupsafe-3.0.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:f9e130248f4462aaa8e2552d547f36ddadbeaa573879158d721bbd33dfe4743a", size = 22923, upload-time = "2025-09-27T18:36:27.109Z" },
	76	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/0f/62/d9c46a7f5c9adbeeeda52f5b8d802e1094e9717705a645efc71b0913a0a8/markupsafe-3.0.3-cp311-cp311-win32.whl", hash = "sha256:0db14f5dafddbb6d9208827849fad01f1a2609380add406671a26386cdf15a19", size = 14572, upload-time = "2025-09-27T18:36:28.045Z" },
	77	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/83/8a/4414c03d3f891739326e1783338e48fb49781cc915b2e0ee052aa490d586/markupsafe-3.0.3-cp311-cp311-win_amd64.whl", hash = "sha256:de8a88e63464af587c950061a5e6a67d3632e36df62b986892331d4620a35c01", size = 15077, upload-time = "2025-09-27T18:36:29.025Z" },
	78	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/35/73/893072b42e6862f319b5207adc9ae06070f095b358655f077f69a35601f0/markupsafe-3.0.3-cp311-cp311-win_arm64.whl", hash = "sha256:3b562dd9e9ea93f13d53989d23a7e775fdfd1066c33494ff43f5418bc8c58a5c", size = 13876, upload-time = "2025-09-27T18:36:29.954Z" },
	79	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e", size = 11615, upload-time = "2025-09-27T18:36:30.854Z" },
	80	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/9a/81/7e4e08678a1f98521201c3079f77db69fb552acd56067661f8c2f534a718/markupsafe-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce", size = 12020, upload-time = "2025-09-27T18:36:31.971Z" },
	81	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1e/2c/799f4742efc39633a1b54a92eec4082e4f815314869865d876824c257c1e/markupsafe-3.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d", size = 24332, upload-time = "2025-09-27T18:36:32.813Z" },
	82	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" },
	83	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/2c/54/887f3092a85238093a0b2154bd629c89444f395618842e8b0c41783898ea/markupsafe-3.0.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a", size = 21962, upload-time = "2025-09-27T18:36:35.099Z" },
	84	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c9/2f/336b8c7b6f4a4d95e91119dc8521402461b74a485558d8f238a68312f11c/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b", size = 23760, upload-time = "2025-09-27T18:36:36.001Z" },
	85	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/32/43/67935f2b7e4982ffb50a4d169b724d74b62a3964bc1a9a527f5ac4f1ee2b/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f", size = 21529, upload-time = "2025-09-27T18:36:36.906Z" },
	86	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" },
	87	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d", size = 14540, upload-time = "2025-09-27T18:36:38.761Z" },
	88	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c", size = 15105, upload-time = "2025-09-27T18:36:39.701Z" },
	89	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" },
	90	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/38/2f/907b9c7bbba283e68f20259574b13d005c121a0fa4c175f9bed27c4597ff/markupsafe-3.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e1cf1972137e83c5d4c136c43ced9ac51d0e124706ee1c8aa8532c1287fa8795", size = 11622, upload-time = "2025-09-27T18:36:41.777Z" },
	91	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/9c/d9/5f7756922cdd676869eca1c4e3c0cd0df60ed30199ffd775e319089cb3ed/markupsafe-3.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:116bb52f642a37c115f517494ea5feb03889e04df47eeff5b130b1808ce7c219", size = 12029, upload-time = "2025-09-27T18:36:43.257Z" },
	92	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/00/07/575a68c754943058c78f30db02ee03a64b3c638586fba6a6dd56830b30a3/markupsafe-3.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:133a43e73a802c5562be9bbcd03d090aa5a1fe899db609c29e8c8d815c5f6de6", size = 24374, upload-time = "2025-09-27T18:36:44.508Z" },
	93	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a9/21/9b05698b46f218fc0e118e1f8168395c65c8a2c750ae2bab54fc4bd4e0e8/markupsafe-3.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ccfcd093f13f0f0b7fdd0f198b90053bf7b2f02a3927a30e63f3ccc9df56b676", size = 22980, upload-time = "2025-09-27T18:36:45.385Z" },
	94	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7f/71/544260864f893f18b6827315b988c146b559391e6e7e8f7252839b1b846a/markupsafe-3.0.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:509fa21c6deb7a7a273d629cf5ec029bc209d1a51178615ddf718f5918992ab9", size = 21990, upload-time = "2025-09-27T18:36:46.916Z" },
	95	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c2/28/b50fc2f74d1ad761af2f5dcce7492648b983d00a65b8c0e0cb457c82ebbe/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a4afe79fb3de0b7097d81da19090f4df4f8d3a2b3adaa8764138aac2e44f3af1", size = 23784, upload-time = "2025-09-27T18:36:47.884Z" },
	96	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ed/76/104b2aa106a208da8b17a2fb72e033a5a9d7073c68f7e508b94916ed47a9/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:795e7751525cae078558e679d646ae45574b47ed6e7771863fcc079a6171a0fc", size = 21588, upload-time = "2025-09-27T18:36:48.82Z" },
	97	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b5/99/16a5eb2d140087ebd97180d95249b00a03aa87e29cc224056274f2e45fd6/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8485f406a96febb5140bfeca44a73e3ce5116b2501ac54fe953e488fb1d03b12", size = 23041, upload-time = "2025-09-27T18:36:49.797Z" },
	98	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/19/bc/e7140ed90c5d61d77cea142eed9f9c303f4c4806f60a1044c13e3f1471d0/markupsafe-3.0.3-cp313-cp313-win32.whl", hash = "sha256:bdd37121970bfd8be76c5fb069c7751683bdf373db1ed6c010162b2a130248ed", size = 14543, upload-time = "2025-09-27T18:36:51.584Z" },
	99	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/05/73/c4abe620b841b6b791f2edc248f556900667a5a1cf023a6646967ae98335/markupsafe-3.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:9a1abfdc021a164803f4d485104931fb8f8c1efd55bc6b748d2f5774e78b62c5", size = 15113, upload-time = "2025-09-27T18:36:52.537Z" },
	100	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f0/3a/fa34a0f7cfef23cf9500d68cb7c32dd64ffd58a12b09225fb03dd37d5b80/markupsafe-3.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:7e68f88e5b8799aa49c85cd116c932a1ac15caaa3f5db09087854d218359e485", size = 13911, upload-time = "2025-09-27T18:36:53.513Z" },
	101	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e4/d7/e05cd7efe43a88a17a37b3ae96e79a19e846f3f456fe79c57ca61356ef01/markupsafe-3.0.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:218551f6df4868a8d527e3062d0fb968682fe92054e89978594c28e642c43a73", size = 11658, upload-time = "2025-09-27T18:36:54.819Z" },
	102	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/99/9e/e412117548182ce2148bdeacdda3bb494260c0b0184360fe0d56389b523b/markupsafe-3.0.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:3524b778fe5cfb3452a09d31e7b5adefeea8c5be1d43c4f810ba09f2ceb29d37", size = 12066, upload-time = "2025-09-27T18:36:55.714Z" },
	103	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/bc/e6/fa0ffcda717ef64a5108eaa7b4f5ed28d56122c9a6d70ab8b72f9f715c80/markupsafe-3.0.3-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4e885a3d1efa2eadc93c894a21770e4bc67899e3543680313b09f139e149ab19", size = 25639, upload-time = "2025-09-27T18:36:56.908Z" },
	104	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/96/ec/2102e881fe9d25fc16cb4b25d5f5cde50970967ffa5dddafdb771237062d/markupsafe-3.0.3-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8709b08f4a89aa7586de0aadc8da56180242ee0ada3999749b183aa23df95025", size = 23569, upload-time = "2025-09-27T18:36:57.913Z" },
	105	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4b/30/6f2fce1f1f205fc9323255b216ca8a235b15860c34b6798f810f05828e32/markupsafe-3.0.3-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b8512a91625c9b3da6f127803b166b629725e68af71f8184ae7e7d54686a56d6", size = 23284, upload-time = "2025-09-27T18:36:58.833Z" },
	106	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/58/47/4a0ccea4ab9f5dcb6f79c0236d954acb382202721e704223a8aafa38b5c8/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9b79b7a16f7fedff2495d684f2b59b0457c3b493778c9eed31111be64d58279f", size = 24801, upload-time = "2025-09-27T18:36:59.739Z" },
	107	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/6a/70/3780e9b72180b6fecb83a4814d84c3bf4b4ae4bf0b19c27196104149734c/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:12c63dfb4a98206f045aa9563db46507995f7ef6d83b2f68eda65c307c6829eb", size = 22769, upload-time = "2025-09-27T18:37:00.719Z" },
	108	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/98/c5/c03c7f4125180fc215220c035beac6b9cb684bc7a067c84fc69414d315f5/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8f71bc33915be5186016f675cd83a1e08523649b0e33efdb898db577ef5bb009", size = 23642, upload-time = "2025-09-27T18:37:01.673Z" },
	109	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/80/d6/2d1b89f6ca4bff1036499b1e29a1d02d282259f3681540e16563f27ebc23/markupsafe-3.0.3-cp313-cp313t-win32.whl", hash = "sha256:69c0b73548bc525c8cb9a251cddf1931d1db4d2258e9599c28c07ef3580ef354", size = 14612, upload-time = "2025-09-27T18:37:02.639Z" },
	110	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/2b/98/e48a4bfba0a0ffcf9925fe2d69240bfaa19c6f7507b8cd09c70684a53c1e/markupsafe-3.0.3-cp313-cp313t-win_amd64.whl", hash = "sha256:1b4b79e8ebf6b55351f0d91fe80f893b4743f104bff22e90697db1590e47a218", size = 15200, upload-time = "2025-09-27T18:37:03.582Z" },
	111	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/0e/72/e3cc540f351f316e9ed0f092757459afbc595824ca724cbc5a5d4263713f/markupsafe-3.0.3-cp313-cp313t-win_arm64.whl", hash = "sha256:ad2cf8aa28b8c020ab2fc8287b0f823d0a7d8630784c31e9ee5edea20f406287", size = 13973, upload-time = "2025-09-27T18:37:04.929Z" },
	112	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/33/8a/8e42d4838cd89b7dde187011e97fe6c3af66d8c044997d2183fbd6d31352/markupsafe-3.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:eaa9599de571d72e2daf60164784109f19978b327a3910d3e9de8c97b5b70cfe", size = 11619, upload-time = "2025-09-27T18:37:06.342Z" },
	113	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b5/64/7660f8a4a8e53c924d0fa05dc3a55c9cee10bbd82b11c5afb27d44b096ce/markupsafe-3.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c47a551199eb8eb2121d4f0f15ae0f923d31350ab9280078d1e5f12b249e0026", size = 12029, upload-time = "2025-09-27T18:37:07.213Z" },
	114	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/da/ef/e648bfd021127bef5fa12e1720ffed0c6cbb8310c8d9bea7266337ff06de/markupsafe-3.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f34c41761022dd093b4b6896d4810782ffbabe30f2d443ff5f083e0cbbb8c737", size = 24408, upload-time = "2025-09-27T18:37:09.572Z" },
	115	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/41/3c/a36c2450754618e62008bf7435ccb0f88053e07592e6028a34776213d877/markupsafe-3.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:457a69a9577064c05a97c41f4e65148652db078a3a509039e64d3467b9e7ef97", size = 23005, upload-time = "2025-09-27T18:37:10.58Z" },
	116	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/bc/20/b7fdf89a8456b099837cd1dc21974632a02a999ec9bf7ca3e490aacd98e7/markupsafe-3.0.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e8afc3f2ccfa24215f8cb28dcf43f0113ac3c37c2f0f0806d8c70e4228c5cf4d", size = 22048, upload-time = "2025-09-27T18:37:11.547Z" },
	117	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/9a/a7/591f592afdc734f47db08a75793a55d7fbcc6902a723ae4cfbab61010cc5/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ec15a59cf5af7be74194f7ab02d0f59a62bdcf1a537677ce67a2537c9b87fcda", size = 23821, upload-time = "2025-09-27T18:37:12.48Z" },
	118	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7d/33/45b24e4f44195b26521bc6f1a82197118f74df348556594bd2262bda1038/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:0eb9ff8191e8498cca014656ae6b8d61f39da5f95b488805da4bb029cccbfbaf", size = 21606, upload-time = "2025-09-27T18:37:13.485Z" },
	119	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ff/0e/53dfaca23a69fbfbbf17a4b64072090e70717344c52eaaaa9c5ddff1e5f0/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2713baf880df847f2bece4230d4d094280f4e67b1e813eec43b4c0e144a34ffe", size = 23043, upload-time = "2025-09-27T18:37:14.408Z" },
	120	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/46/11/f333a06fc16236d5238bfe74daccbca41459dcd8d1fa952e8fbd5dccfb70/markupsafe-3.0.3-cp314-cp314-win32.whl", hash = "sha256:729586769a26dbceff69f7a7dbbf59ab6572b99d94576a5592625d5b411576b9", size = 14747, upload-time = "2025-09-27T18:37:15.36Z" },
	121	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/28/52/182836104b33b444e400b14f797212f720cbc9ed6ba34c800639d154e821/markupsafe-3.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:bdc919ead48f234740ad807933cdf545180bfbe9342c2bb451556db2ed958581", size = 15341, upload-time = "2025-09-27T18:37:16.496Z" },
	122	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/6f/18/acf23e91bd94fd7b3031558b1f013adfa21a8e407a3fdb32745538730382/markupsafe-3.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:5a7d5dc5140555cf21a6fefbdbf8723f06fcd2f63ef108f2854de715e4422cb4", size = 14073, upload-time = "2025-09-27T18:37:17.476Z" },
	123	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/3c/f0/57689aa4076e1b43b15fdfa646b04653969d50cf30c32a102762be2485da/markupsafe-3.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:1353ef0c1b138e1907ae78e2f6c63ff67501122006b0f9abad68fda5f4ffc6ab", size = 11661, upload-time = "2025-09-27T18:37:18.453Z" },
	124	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/89/c3/2e67a7ca217c6912985ec766c6393b636fb0c2344443ff9d91404dc4c79f/markupsafe-3.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1085e7fbddd3be5f89cc898938f42c0b3c711fdcb37d75221de2666af647c175", size = 12069, upload-time = "2025-09-27T18:37:19.332Z" },
	125	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f0/00/be561dce4e6ca66b15276e184ce4b8aec61fe83662cce2f7d72bd3249d28/markupsafe-3.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1b52b4fb9df4eb9ae465f8d0c228a00624de2334f216f178a995ccdcf82c4634", size = 25670, upload-time = "2025-09-27T18:37:20.245Z" },
	126	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/50/09/c419f6f5a92e5fadde27efd190eca90f05e1261b10dbd8cbcb39cd8ea1dc/markupsafe-3.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50", size = 23598, upload-time = "2025-09-27T18:37:21.177Z" },
	127	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/22/44/a0681611106e0b2921b3033fc19bc53323e0b50bc70cffdd19f7d679bb66/markupsafe-3.0.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f190daf01f13c72eac4efd5c430a8de82489d9cff23c364c3ea822545032993e", size = 23261, upload-time = "2025-09-27T18:37:22.167Z" },
	128	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5f/57/1b0b3f100259dc9fffe780cfb60d4be71375510e435efec3d116b6436d43/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e56b7d45a839a697b5eb268c82a71bd8c7f6c94d6fd50c3d577fa39a9f1409f5", size = 24835, upload-time = "2025-09-27T18:37:23.296Z" },
	129	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/26/6a/4bf6d0c97c4920f1597cc14dd720705eca0bf7c787aebc6bb4d1bead5388/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:f3e98bb3798ead92273dc0e5fd0f31ade220f59a266ffd8a4f6065e0a3ce0523", size = 22733, upload-time = "2025-09-27T18:37:24.237Z" },
	130	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/14/c7/ca723101509b518797fedc2fdf79ba57f886b4aca8a7d31857ba3ee8281f/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5678211cb9333a6468fb8d8be0305520aa073f50d17f089b5b4b477ea6e67fdc", size = 23672, upload-time = "2025-09-27T18:37:25.271Z" },
	131	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fb/df/5bd7a48c256faecd1d36edc13133e51397e41b73bb77e1a69deab746ebac/markupsafe-3.0.3-cp314-cp314t-win32.whl", hash = "sha256:915c04ba3851909ce68ccc2b8e2cd691618c4dc4c4232fb7982bca3f41fd8c3d", size = 14819, upload-time = "2025-09-27T18:37:26.285Z" },
	132	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1a/8a/0402ba61a2f16038b48b39bccca271134be00c5c9f0f623208399333c448/markupsafe-3.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4faffd047e07c38848ce017e8725090413cd80cbc23d86e55c587bf979e579c9", size = 15426, upload-time = "2025-09-27T18:37:27.316Z" },
	133	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/70/bc/6f1c2f612465f5fa89b95bead1f44dcb607670fd42891d8fdcd5d039f4f4/markupsafe-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:32001d6a8fc98c8cb5c947787c5d08b0a50663d139f1305bac5885d98d9b40fa", size = 14146, upload-time = "2025-09-27T18:37:28.327Z" },
	134	+]
	135	+
	136	+[[package]]
	137	+name = "ml-dtypes"
	138	+version = "0.4.1"
	139	+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" }
	140	+resolution-markers = [
	141	+ "python_full_version >= '3.13' and sys_platform != 'darwin'",
	142	+ "python_full_version >= '3.13' and sys_platform == 'darwin'",
	143	+]
	144	+dependencies = [
	145	+ { name = "numpy", marker = "python_full_version >= '3.13'" },
	146	+]
	147	+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/fd/15/76f86faa0902836cc133939732f7611ace68cf54148487a99c539c272dc8/ml_dtypes-0.4.1.tar.gz", hash = "sha256:fad5f2de464fd09127e49b7fd1252b9006fb43d2edc1ff112d390c324af5ca7a", size = 692594, upload-time = "2024-09-13T19:07:11.624Z" }
	148	+wheels = [
	149	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d1/76/9835c8609c29f2214359e88f29255fc4aad4ea0f613fb48aa8815ceda1b6/ml_dtypes-0.4.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:2d55b588116a7085d6e074cf0cdb1d6fa3875c059dddc4d2c94a4cc81c23e975", size = 397973, upload-time = "2024-09-13T19:06:51.748Z" },
	150	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/7e/99/e68c56fac5de973007a10254b6e17a0362393724f40f66d5e4033f4962c2/ml_dtypes-0.4.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e138a9b7a48079c900ea969341a5754019a1ad17ae27ee330f7ebf43f23877f9", size = 2185134, upload-time = "2024-09-13T19:06:53.197Z" },
	151	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/28/bc/6a2344338ea7b61cd7b46fb24ec459360a5a0903b57c55b156c1e46c644a/ml_dtypes-0.4.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:74c6cfb5cf78535b103fde9ea3ded8e9f16f75bc07789054edc7776abfb3d752", size = 2163661, upload-time = "2024-09-13T19:06:54.519Z" },
	152	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e8/d3/ddfd9878b223b3aa9a930c6100a99afca5cfab7ea703662e00323acb7568/ml_dtypes-0.4.1-cp311-cp311-win_amd64.whl", hash = "sha256:274cc7193dd73b35fb26bef6c5d40ae3eb258359ee71cd82f6e96a8c948bdaa6", size = 126727, upload-time = "2024-09-13T19:06:55.897Z" },
	153	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ba/1a/99e924f12e4b62139fbac87419698c65f956d58de0dbfa7c028fa5b096aa/ml_dtypes-0.4.1-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:827d3ca2097085cf0355f8fdf092b888890bb1b1455f52801a2d7756f056f54b", size = 405077, upload-time = "2024-09-13T19:06:57.538Z" },
	154	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8f/8c/7b610bd500617854c8cc6ed7c8cfb9d48d6a5c21a1437a36a4b9bc8a3598/ml_dtypes-0.4.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:772426b08a6172a891274d581ce58ea2789cc8abc1c002a27223f314aaf894e7", size = 2181554, upload-time = "2024-09-13T19:06:59.196Z" },
	155	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c7/c6/f89620cecc0581dc1839e218c4315171312e46c62a62da6ace204bda91c0/ml_dtypes-0.4.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:126e7d679b8676d1a958f2651949fbfa182832c3cd08020d8facd94e4114f3e9", size = 2160488, upload-time = "2024-09-13T19:07:03.131Z" },
	156	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ae/11/a742d3c31b2cc8557a48efdde53427fd5f9caa2fa3c9c27d826e78a66f51/ml_dtypes-0.4.1-cp312-cp312-win_amd64.whl", hash = "sha256:df0fb650d5c582a9e72bb5bd96cfebb2cdb889d89daff621c8fbc60295eba66c", size = 127462, upload-time = "2024-09-13T19:07:04.916Z" },
	157	+]
	158	+
	159	+[[package]]
	160	+name = "ml-dtypes"
	161	+version = "0.5.4"
	162	+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" }
	163	+resolution-markers = [
	164	+ "python_full_version == '3.12.*' and sys_platform != 'darwin'",
	165	+ "python_full_version < '3.12' and sys_platform != 'darwin'",
	166	+ "python_full_version == '3.12.*' and sys_platform == 'darwin'",
	167	+ "python_full_version < '3.12' and sys_platform == 'darwin'",
	168	+]
	169	+dependencies = [
	170	+ { name = "numpy", marker = "python_full_version < '3.13'" },
	171	+]
	172	+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz", hash = "sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453", size = 692314, upload-time = "2025-11-17T22:32:31.031Z" }
	173	+wheels = [
	174	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c6/5e/712092cfe7e5eb667b8ad9ca7c54442f21ed7ca8979745f1000e24cf8737/ml_dtypes-0.5.4-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:6c7ecb74c4bd71db68a6bea1edf8da8c34f3d9fe218f038814fd1d310ac76c90", size = 679734, upload-time = "2025-11-17T22:31:39.223Z" },
	175	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4f/cf/912146dfd4b5c0eea956836c01dcd2fce6c9c844b2691f5152aca196ce4f/ml_dtypes-0.5.4-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bc11d7e8c44a65115d05e2ab9989d1e045125d7be8e05a071a48bc76eb6d6040", size = 5056165, upload-time = "2025-11-17T22:31:41.071Z" },
	176	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a9/80/19189ea605017473660e43762dc853d2797984b3c7bf30ce656099add30c/ml_dtypes-0.5.4-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:19b9a53598f21e453ea2fbda8aa783c20faff8e1eeb0d7ab899309a0053f1483", size = 5034975, upload-time = "2025-11-17T22:31:42.758Z" },
	177	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b4/24/70bd59276883fdd91600ca20040b41efd4902a923283c4d6edcb1de128d2/ml_dtypes-0.5.4-cp311-cp311-win_amd64.whl", hash = "sha256:7c23c54a00ae43edf48d44066a7ec31e05fdc2eee0be2b8b50dd1903a1db94bb", size = 210742, upload-time = "2025-11-17T22:31:44.068Z" },
	178	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a0/c9/64230ef14e40aa3f1cb254ef623bf812735e6bec7772848d19131111ac0d/ml_dtypes-0.5.4-cp311-cp311-win_arm64.whl", hash = "sha256:557a31a390b7e9439056644cb80ed0735a6e3e3bb09d67fd5687e4b04238d1de", size = 160709, upload-time = "2025-11-17T22:31:46.557Z" },
	179	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a8/b8/3c70881695e056f8a32f8b941126cf78775d9a4d7feba8abcb52cb7b04f2/ml_dtypes-0.5.4-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:a174837a64f5b16cab6f368171a1a03a27936b31699d167684073ff1c4237dac", size = 676927, upload-time = "2025-11-17T22:31:48.182Z" },
	180	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/54/0f/428ef6881782e5ebb7eca459689448c0394fa0a80bea3aa9262cba5445ea/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a7f7c643e8b1320fd958bf098aa7ecf70623a42ec5154e3be3be673f4c34d900", size = 5028464, upload-time = "2025-11-17T22:31:50.135Z" },
	181	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/3a/cb/28ce52eb94390dda42599c98ea0204d74799e4d8047a0eb559b6fd648056/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ad459e99793fa6e13bd5b7e6792c8f9190b4e5a1b45c63aba14a4d0a7f1d5ff", size = 5009002, upload-time = "2025-11-17T22:31:52.001Z" },
	182	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/f5/f0/0cfadd537c5470378b1b32bd859cf2824972174b51b873c9d95cfd7475a5/ml_dtypes-0.5.4-cp312-cp312-win_amd64.whl", hash = "sha256:c1a953995cccb9e25a4ae19e34316671e4e2edaebe4cf538229b1fc7109087b7", size = 212222, upload-time = "2025-11-17T22:31:53.742Z" },
	183	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/16/2e/9acc86985bfad8f2c2d30291b27cd2bb4c74cea08695bd540906ed744249/ml_dtypes-0.5.4-cp312-cp312-win_arm64.whl", hash = "sha256:9bad06436568442575beb2d03389aa7456c690a5b05892c471215bfd8cf39460", size = 160793, upload-time = "2025-11-17T22:31:55.358Z" },
	184	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d9/a1/4008f14bbc616cfb1ac5b39ea485f9c63031c4634ab3f4cf72e7541f816a/ml_dtypes-0.5.4-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8c760d85a2f82e2bed75867079188c9d18dae2ee77c25a54d60e9cc79be1bc48", size = 676888, upload-time = "2025-11-17T22:31:56.907Z" },
	185	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d3/b7/dff378afc2b0d5a7d6cd9d3209b60474d9819d1189d347521e1688a60a53/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ce756d3a10d0c4067172804c9cc276ba9cc0ff47af9078ad439b075d1abdc29b", size = 5036993, upload-time = "2025-11-17T22:31:58.497Z" },
	186	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/eb/33/40cd74219417e78b97c47802037cf2d87b91973e18bb968a7da48a96ea44/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:533ce891ba774eabf607172254f2e7260ba5f57bdd64030c9a4fcfbd99815d0d", size = 5010956, upload-time = "2025-11-17T22:31:59.931Z" },
	187	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e1/8b/200088c6859d8221454825959df35b5244fa9bdf263fd0249ac5fb75e281/ml_dtypes-0.5.4-cp313-cp313-win_amd64.whl", hash = "sha256:f21c9219ef48ca5ee78402d5cc831bd58ea27ce89beda894428bc67a52da5328", size = 212224, upload-time = "2025-11-17T22:32:01.349Z" },
	188	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8f/75/dfc3775cb36367816e678f69a7843f6f03bd4e2bcd79941e01ea960a068e/ml_dtypes-0.5.4-cp313-cp313-win_arm64.whl", hash = "sha256:35f29491a3e478407f7047b8a4834e4640a77d2737e0b294d049746507af5175", size = 160798, upload-time = "2025-11-17T22:32:02.864Z" },
	189	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4f/74/e9ddb35fd1dd43b1106c20ced3f53c2e8e7fc7598c15638e9f80677f81d4/ml_dtypes-0.5.4-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:304ad47faa395415b9ccbcc06a0350800bc50eda70f0e45326796e27c62f18b6", size = 702083, upload-time = "2025-11-17T22:32:04.08Z" },
	190	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/74/f5/667060b0aed1aa63166b22897fdf16dca9eb704e6b4bbf86848d5a181aa7/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6a0df4223b514d799b8a1629c65ddc351b3efa833ccf7f8ea0cf654a61d1e35d", size = 5354111, upload-time = "2025-11-17T22:32:05.546Z" },
	191	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/40/49/0f8c498a28c0efa5f5c95a9e374c83ec1385ca41d0e85e7cf40e5d519a21/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:531eff30e4d368cb6255bc2328d070e35836aa4f282a0fb5f3a0cd7260257298", size = 5366453, upload-time = "2025-11-17T22:32:07.115Z" },
	192	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/8c/27/12607423d0a9c6bbbcc780ad19f1f6baa2b68b18ce4bddcdc122c4c68dc9/ml_dtypes-0.5.4-cp313-cp313t-win_amd64.whl", hash = "sha256:cb73dccfc991691c444acc8c0012bee8f2470da826a92e3a20bb333b1a7894e6", size = 225612, upload-time = "2025-11-17T22:32:08.615Z" },
	193	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e5/80/5a5929e92c72936d5b19872c5fb8fc09327c1da67b3b68c6a13139e77e20/ml_dtypes-0.5.4-cp313-cp313t-win_arm64.whl", hash = "sha256:3bbbe120b915090d9dd1375e4684dd17a20a2491ef25d640a908281da85e73f1", size = 164145, upload-time = "2025-11-17T22:32:09.782Z" },
	194	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/72/4e/1339dc6e2557a344f5ba5590872e80346f76f6cb2ac3dd16e4666e88818c/ml_dtypes-0.5.4-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:2b857d3af6ac0d39db1de7c706e69c7f9791627209c3d6dedbfca8c7e5faec22", size = 673781, upload-time = "2025-11-17T22:32:11.364Z" },
	195	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/04/f9/067b84365c7e83bda15bba2b06c6ca250ce27b20630b1128c435fb7a09aa/ml_dtypes-0.5.4-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:805cef3a38f4eafae3a5bf9ebdcdb741d0bcfd9e1bd90eb54abd24f928cd2465", size = 5036145, upload-time = "2025-11-17T22:32:12.783Z" },
	196	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/c6/bb/82c7dcf38070b46172a517e2334e665c5bf374a262f99a283ea454bece7c/ml_dtypes-0.5.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:14a4fd3228af936461db66faccef6e4f41c1d82fcc30e9f8d58a08916b1d811f", size = 5010230, upload-time = "2025-11-17T22:32:14.38Z" },
	197	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e9/93/2bfed22d2498c468f6bcd0d9f56b033eaa19f33320389314c19ef6766413/ml_dtypes-0.5.4-cp314-cp314-win_amd64.whl", hash = "sha256:8c6a2dcebd6f3903e05d51960a8058d6e131fe69f952a5397e5dbabc841b6d56", size = 221032, upload-time = "2025-11-17T22:32:15.763Z" },
	198	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/76/a3/9c912fe6ea747bb10fe2f8f54d027eb265db05dfb0c6335e3e063e74e6e8/ml_dtypes-0.5.4-cp314-cp314-win_arm64.whl", hash = "sha256:5a0f68ca8fd8d16583dfa7793973feb86f2fbb56ce3966daf9c9f748f52a2049", size = 163353, upload-time = "2025-11-17T22:32:16.932Z" },
	199	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/cd/02/48aa7d84cc30ab4ee37624a2fd98c56c02326785750cd212bc0826c2f15b/ml_dtypes-0.5.4-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:bfc534409c5d4b0bf945af29e5d0ab075eae9eecbb549ff8a29280db822f34f9", size = 702085, upload-time = "2025-11-17T22:32:18.175Z" },
	200	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/5a/e7/85cb99fe80a7a5513253ec7faa88a65306be071163485e9a626fce1b6e84/ml_dtypes-0.5.4-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2314892cdc3fcf05e373d76d72aaa15fda9fb98625effa73c1d646f331fcecb7", size = 5355358, upload-time = "2025-11-17T22:32:19.7Z" },
	201	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/79/2b/a826ba18d2179a56e144aef69e57fb2ab7c464ef0b2111940ee8a3a223a2/ml_dtypes-0.5.4-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0d2ffd05a2575b1519dc928c0b93c06339eb67173ff53acb00724502cda231cf", size = 5366332, upload-time = "2025-11-17T22:32:21.193Z" },
	202	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/84/44/f4d18446eacb20ea11e82f133ea8f86e2bf2891785b67d9da8d0ab0ef525/ml_dtypes-0.5.4-cp314-cp314t-win_amd64.whl", hash = "sha256:4381fe2f2452a2d7589689693d3162e876b3ddb0a832cde7a414f8e1adf7eab1", size = 236612, upload-time = "2025-11-17T22:32:22.579Z" },
	203	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ad/3f/3d42e9a78fe5edf792a83c074b13b9b770092a4fbf3462872f4303135f09/ml_dtypes-0.5.4-cp314-cp314t-win_arm64.whl", hash = "sha256:11942cbf2cf92157db91e5022633c0d9474d4dfd813a909383bd23ce828a4b7d", size = 168825, upload-time = "2025-11-17T22:32:23.766Z" },
	204	+]
	205	+
	206	+[[package]]
	207	+name = "mpmath"
	208	+version = "1.3.0"
	209	+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" }
	210	+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" }
	211	+wheels = [
	212	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" },
	213	+]
	214	+
	215	+[[package]]
	216	+name = "networkx"
	217	+version = "3.6"
	218	+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" }
	219	+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/e8/fc/7b6fd4d22c8c4dc5704430140d8b3f520531d4fe7328b8f8d03f5a7950e8/networkx-3.6.tar.gz", hash = "sha256:285276002ad1f7f7da0f7b42f004bcba70d381e936559166363707fdad3d72ad", size = 2511464, upload-time = "2025-11-24T03:03:47.158Z" }
	220	+wheels = [
	221	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/07/c7/d64168da60332c17d24c0d2f08bdf3987e8d1ae9d84b5bbd0eec2eb26a55/networkx-3.6-py3-none-any.whl", hash = "sha256:cdb395b105806062473d3be36458d8f1459a4e4b98e236a66c3a48996e07684f", size = 2063713, upload-time = "2025-11-24T03:03:45.21Z" },
	222	+]
	223	+
	224	+[[package]]
	225	+name = "numpy"
	226	+version = "1.26.4"
	227	+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" }
	228	+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/65/6e/09db70a523a96d25e115e71cc56a6f9031e7b8cd166c1ac8438307c14058/numpy-1.26.4.tar.gz", hash = "sha256:2a02aba9ed12e4ac4eb3ea9421c420301a0c6460d9830d74a9df87efa4912010", size = 15786129, upload-time = "2024-02-06T00:26:44.495Z" }
	229	+wheels = [
	230	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/11/57/baae43d14fe163fa0e4c47f307b6b2511ab8d7d30177c491960504252053/numpy-1.26.4-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:4c66707fabe114439db9068ee468c26bbdf909cac0fb58686a42a24de1760c71", size = 20630554, upload-time = "2024-02-05T23:51:50.149Z" },
	231	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/1a/2e/151484f49fd03944c4a3ad9c418ed193cfd02724e138ac8a9505d056c582/numpy-1.26.4-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:edd8b5fe47dab091176d21bb6de568acdd906d1887a4584a15a9a96a1dca06ef", size = 13997127, upload-time = "2024-02-05T23:52:15.314Z" },
	232	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/79/ae/7e5b85136806f9dadf4878bf73cf223fe5c2636818ba3ab1c585d0403164/numpy-1.26.4-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7ab55401287bfec946ced39700c053796e7cc0e3acbef09993a9ad2adba6ca6e", size = 14222994, upload-time = "2024-02-05T23:52:47.569Z" },
	233	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/3a/d0/edc009c27b406c4f9cbc79274d6e46d634d139075492ad055e3d68445925/numpy-1.26.4-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:666dbfb6ec68962c033a450943ded891bed2d54e6755e35e5835d63f4f6931d5", size = 18252005, upload-time = "2024-02-05T23:53:15.637Z" },
	234	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/09/bf/2b1aaf8f525f2923ff6cfcf134ae5e750e279ac65ebf386c75a0cf6da06a/numpy-1.26.4-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:96ff0b2ad353d8f990b63294c8986f1ec3cb19d749234014f4e7eb0112ceba5a", size = 13885297, upload-time = "2024-02-05T23:53:42.16Z" },
	235	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/df/a0/4e0f14d847cfc2a633a1c8621d00724f3206cfeddeb66d35698c4e2cf3d2/numpy-1.26.4-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:60dedbb91afcbfdc9bc0b1f3f402804070deed7392c23eb7a7f07fa857868e8a", size = 18093567, upload-time = "2024-02-05T23:54:11.696Z" },
	236	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/d2/b7/a734c733286e10a7f1a8ad1ae8c90f2d33bf604a96548e0a4a3a6739b468/numpy-1.26.4-cp311-cp311-win32.whl", hash = "sha256:1af303d6b2210eb850fcf03064d364652b7120803a0b872f5211f5234b399f20", size = 5968812, upload-time = "2024-02-05T23:54:26.453Z" },
	237	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/3f/6b/5610004206cf7f8e7ad91c5a85a8c71b2f2f8051a0c0c4d5916b76d6cbb2/numpy-1.26.4-cp311-cp311-win_amd64.whl", hash = "sha256:cd25bcecc4974d09257ffcd1f098ee778f7834c3ad767fe5db785be9a4aa9cb2", size = 15811913, upload-time = "2024-02-05T23:54:53.933Z" },
	238	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/95/12/8f2020a8e8b8383ac0177dc9570aad031a3beb12e38847f7129bacd96228/numpy-1.26.4-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:b3ce300f3644fb06443ee2222c2201dd3a89ea6040541412b8fa189341847218", size = 20335901, upload-time = "2024-02-05T23:55:32.801Z" },
	239	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/75/5b/ca6c8bd14007e5ca171c7c03102d17b4f4e0ceb53957e8c44343a9546dcc/numpy-1.26.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:03a8c78d01d9781b28a6989f6fa1bb2c4f2d51201cf99d3dd875df6fbd96b23b", size = 13685868, upload-time = "2024-02-05T23:55:56.28Z" },
	240	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/79/f8/97f10e6755e2a7d027ca783f63044d5b1bc1ae7acb12afe6a9b4286eac17/numpy-1.26.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9fad7dcb1aac3c7f0584a5a8133e3a43eeb2fe127f47e3632d43d677c66c102b", size = 13925109, upload-time = "2024-02-05T23:56:20.368Z" },
	241	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/0f/50/de23fde84e45f5c4fda2488c759b69990fd4512387a8632860f3ac9cd225/numpy-1.26.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:675d61ffbfa78604709862923189bad94014bef562cc35cf61d3a07bba02a7ed", size = 17950613, upload-time = "2024-02-05T23:56:56.054Z" },
	242	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/4c/0c/9c603826b6465e82591e05ca230dfc13376da512b25ccd0894709b054ed0/numpy-1.26.4-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:ab47dbe5cc8210f55aa58e4805fe224dac469cde56b9f731a4c098b91917159a", size = 13572172, upload-time = "2024-02-05T23:57:21.56Z" },
	243	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/76/8c/2ba3902e1a0fc1c74962ea9bb33a534bb05984ad7ff9515bf8d07527cadd/numpy-1.26.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:1dda2e7b4ec9dd512f84935c5f126c8bd8b9f2fc001e9f54af255e8c5f16b0e0", size = 17786643, upload-time = "2024-02-05T23:57:56.585Z" },
	244	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/28/4a/46d9e65106879492374999e76eb85f87b15328e06bd1550668f79f7b18c6/numpy-1.26.4-cp312-cp312-win32.whl", hash = "sha256:50193e430acfc1346175fcbdaa28ffec49947a06918b7b92130744e81e640110", size = 5677803, upload-time = "2024-02-05T23:58:08.963Z" },
	245	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/16/2e/86f24451c2d530c88daf997cb8d6ac622c1d40d19f5a031ed68a4b73a374/numpy-1.26.4-cp312-cp312-win_amd64.whl", hash = "sha256:08beddf13648eb95f8d867350f6a018a4be2e5ad54c8d8caed89ebca558b2818", size = 15517754, upload-time = "2024-02-05T23:58:36.364Z" },
	246	+]
	247	+
	248	+[[package]]
	249	+name = "setuptools"
	250	+version = "80.9.0"
	251	+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" }
	252	+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/18/5d/3bf57dcd21979b887f014ea83c24ae194cfcd12b9e0fda66b957c69d1fca/setuptools-80.9.0.tar.gz", hash = "sha256:f36b47402ecde768dbfafc46e8e4207b4360c654f1f3bb84475f0a28628fb19c", size = 1319958, upload-time = "2025-05-27T00:56:51.443Z" }
	253	+wheels = [
	254	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/a3/dc/17031897dae0efacfea57dfd3a82fdd2a2aeb58e0ff71b77b87e44edc772/setuptools-80.9.0-py3-none-any.whl", hash = "sha256:062d34222ad13e0cc312a4c02d73f059e86a4acbfbdea8f8f76b28c99f306922", size = 1201486, upload-time = "2025-05-27T00:56:49.664Z" },
	255	+]
	256	+
	257	+[[package]]
	258	+name = "sympy"
	259	+version = "1.13.1"
	260	+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" }
	261	+dependencies = [
	262	+ { name = "mpmath" },
	263	+]
	264	+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/ca/99/5a5b6f19ff9f083671ddf7b9632028436167cd3d33e11015754e41b249a4/sympy-1.13.1.tar.gz", hash = "sha256:9cebf7e04ff162015ce31c9c6c9144daa34a93bd082f54fd8f12deca4f47515f", size = 7533040, upload-time = "2024-07-19T09:26:51.238Z" }
	265	+wheels = [
	266	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/b2/fe/81695a1aa331a842b582453b605175f419fe8540355886031328089d840a/sympy-1.13.1-py3-none-any.whl", hash = "sha256:db36cdc64bf61b9b24578b6f7bab1ecdd2452cf008f34faa33776680c26d66f8", size = 6189177, upload-time = "2024-07-19T09:26:48.863Z" },
	267	+]
	268	+
	269	+[[package]]
	270	+name = "torch"
	271	+version = "2.6.0"
	272	+source = { registry = "https://download.pytorch.org/whl/cpu" }
	273	+resolution-markers = [
	274	+ "python_full_version >= '3.13' and sys_platform == 'darwin'",
	275	+ "python_full_version == '3.12.*' and sys_platform == 'darwin'",
	276	+ "python_full_version < '3.12' and sys_platform == 'darwin'",
	277	+]
	278	+dependencies = [
	279	+ { name = "filelock", marker = "sys_platform == 'darwin'" },
	280	+ { name = "fsspec", marker = "sys_platform == 'darwin'" },
	281	+ { name = "jinja2", marker = "sys_platform == 'darwin'" },
	282	+ { name = "networkx", marker = "sys_platform == 'darwin'" },
	283	+ { name = "setuptools", marker = "python_full_version >= '3.12' and sys_platform == 'darwin'" },
	284	+ { name = "sympy", marker = "sys_platform == 'darwin'" },
	285	+ { name = "typing-extensions", marker = "sys_platform == 'darwin'" },
	286	+]
	287	+wheels = [
	288	+ { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:94fc63b3b4bedd327af588696559f68c264440e2503cc9e6954019473d74ae21" },
	289	+ { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:9a610afe216a85a8b9bc9f8365ed561535c93e804c2a317ef7fabcc5deda0989" },
	290	+ { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:ff96f4038f8af9f7ec4231710ed4549da1bdebad95923953a25045dcf6fd87e2" },
	291	+]
	292	+
	293	+[[package]]
	294	+name = "torch"
	295	+version = "2.6.0+cpu"
	296	+source = { registry = "https://download.pytorch.org/whl/cpu" }
	297	+resolution-markers = [
	298	+ "python_full_version >= '3.13' and sys_platform != 'darwin'",
	299	+ "python_full_version == '3.12.*' and sys_platform != 'darwin'",
	300	+ "python_full_version < '3.12' and sys_platform != 'darwin'",
	301	+]
	302	+dependencies = [
	303	+ { name = "filelock", marker = "sys_platform != 'darwin'" },
	304	+ { name = "fsspec", marker = "sys_platform != 'darwin'" },
	305	+ { name = "jinja2", marker = "sys_platform != 'darwin'" },
	306	+ { name = "networkx", marker = "sys_platform != 'darwin'" },
	307	+ { name = "setuptools", marker = "python_full_version >= '3.12' and sys_platform != 'darwin'" },
	308	+ { name = "sympy", marker = "sys_platform != 'darwin'" },
	309	+ { name = "typing-extensions", marker = "sys_platform != 'darwin'" },
	310	+]
	311	+wheels = [
	312	+ { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp311-cp311-linux_x86_64.whl", hash = "sha256:5b6ae523bfb67088a17ca7734d131548a2e60346c622621e4248ed09dd0790cc" },
	313	+ { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:d3dab9fb0294f268aec28e8aaba834e9d006b90a50db5bc2fe2191a9d48c6084" },
	314	+ { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp311-cp311-win_amd64.whl", hash = "sha256:24c9d3d13b9ea769dd7bd5c11cfa1fc463fd7391397156565484565ca685d908" },
	315	+ { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp312-cp312-linux_x86_64.whl", hash = "sha256:59e78aa0c690f70734e42670036d6b541930b8eabbaa18d94e090abf14cc4d91" },
	316	+ { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:318290e8924353c61b125cdc8768d15208704e279e7757c113b9620740deca98" },
	317	+ { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp312-cp312-win_amd64.whl", hash = "sha256:4027d982eb2781c93825ab9527f17fbbb12dbabf422298e4b954be60016f87d8" },
	318	+ { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp313-cp313-linux_x86_64.whl", hash = "sha256:e70ee2e37ad27a90201d101a41c2e10df7cf15a9ebd17c084f54cf2518c57bdf" },
	319	+ { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:b5e7e8d561b263b5ad8049736281cd12c78e51e7bc1a913fd4098fd0e0b96347" },
	320	+ { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp313-cp313-win_amd64.whl", hash = "sha256:b436a6c62d086dc5b32f5721b59f0ca8ad3bf9de09ee9b5b83dbf1e7a7e22c60" },
	321	+ { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp313-cp313t-linux_x86_64.whl", hash = "sha256:fb34d6cc4e6e20e66d74852c3d84e0301dc5e1a7c822076ef288886f978390f0" },
	322	+ { url = "https://download.pytorch.org/whl/cpu/torch-2.6.0%2Bcpu-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:7cac05af909ee1c5c2915e8f3efaa1ea015e7e414be0ff53071402b9e4f3c7df" },
	323	+]
	324	+
	325	+[[package]]
	326	+name = "typing-extensions"
	327	+version = "4.15.0"
	328	+source = { registry = "https://pypi.tuna.tsinghua.edu.cn/simple" }
	329	+sdist = { url = "https://pypi.tuna.tsinghua.edu.cn/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" }
	330	+wheels = [
	331	+ { url = "https://pypi.tuna.tsinghua.edu.cn/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" },
	332	+]

Mtests/run_all_test.sh+3-0

	@@ -51,6 +51,9 @@ fi
51	# optest	51	# optest
52	bash "$SCRIPT_PATH/run_optest.sh"	52	bash "$SCRIPT_PATH/run_optest.sh"
53		53
		54	+# catlass_cppgen
		55	+bash "$SCRIPT_PATH/run_cppgen.sh"
		56	+
54	# unittest	57	# unittest
55	bash "$BUILD_SCRIPT_PATH" --clean --tests catlass_unittest \|\| exit 1	58	bash "$BUILD_SCRIPT_PATH" --clean --tests catlass_unittest \|\| exit 1
56	$SCRIPT_PATH/../build/tests/unittest/catlass_unittest_"$CATLASS_ARCH"	59	$SCRIPT_PATH/../build/tests/unittest/catlass_unittest_"$CATLASS_ARCH"

Atests/run_cppgen.sh+81-0

@@ -0,0 +1,81 @@
	1	+#!/bin/bash
	2	+
	3	+set -e
	4	+
	5	+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
	6	+PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
	7	+TEST_DIR="$PROJECT_ROOT/python/catlass_cppgen/tests"
	8	+
	9	+# ------------------------------------------------------------------
	10	+# resolve python
	11	+# ------------------------------------------------------------------
	12	+if command -v pyenv &>/dev/null && pyenv which python &>/dev/null; then
	13	+ PYTHON="$(pyenv which python)"
	14	+else
	15	+ PYTHON="python3"
	16	+fi
	17	+
	18	+# 1. 编译 catlass_cppgen
	19	+echo ""
	20	+echo "============================================"
	21	+echo "Step 1: Building catlass_cppgen..."
	22	+echo "============================================"
	23	+cd "$PROJECT_ROOT/python/catlass_cppgen"
	24	+"$PYTHON" -m pip install build
	25	+"$PYTHON" -m build
	26	+whl_file=$(python -m build 2>&1 \| grep "Successfully built" \| grep -oE "[^ ]+\.whl")
	27	+
	28	+# 2. 安装 catlass_cppgen
	29	+echo ""
	30	+echo "============================================"
	31	+echo "Step 2: Installing catlass_cppgen from dist..."
	32	+echo "============================================"
	33	+"$PYTHON" -m pip install "dist/$whl_file" --force-reinstall --no-deps
	34	+
	35	+# 3. 执行所有catlass_cppgen的测试例
	36	+echo ""
	37	+echo "============================================"
	38	+echo "Step 3: Preparing test all testcases..."
	39	+echo "============================================"
	40	+
	41	+PASSED=0
	42	+FAILED=0
	43	+test_files=()
	44	+while IFS= read -r -d '' file; do
	45	+ test_files+=("$file")
	46	+done < <(find "$TEST_DIR" -name "test_*.py" -print0 \| sort -z)
	47	+
	48	+if [ ${#test_files[@]} -eq 0 ]; then
	49	+ echo "No test files found under $TEST_DIR"
	50	+ exit 1
	51	+fi
	52	+
	53	+for test_file in "${test_files[@]}"; do
	54	+ rel_path="${test_file#$TEST_DIR/}"
	55	+ echo "--- $rel_path ---"
	56	+ if "$PYTHON" "$test_file"; then
	57	+ PASSED=$((PASSED + 1))
	58	+ else
	59	+ FAILED=$((FAILED + 1))
	60	+ echo " [FAIL] $rel_path"
	61	+ fi
	62	+ echo ""
	63	+done
	64	+
	65	+# 4. 清理环境
	66	+echo ""
	67	+echo "============================================"
	68	+echo "Step 4: Cleaning and uninstalling catlass_cppgen..."
	69	+echo "============================================"
	70	+cd "$PROJECT_ROOT/python/catlass_cppgen"
	71	+rm -rf ./dist ./build ./catlass_cppgen.egg-info
	72	+"$PYTHON" -m pip uninstall catlass_cppgen -y
	73	+"$PYTHON" -m pip uninstall build -y
	74	+
	75	+echo ""
	76	+echo "============================================"
	77	+echo "All steps completed successfully!"
	78	+echo " SUMMARY: $PASSED passed, $FAILED failed"
	79	+echo "============================================"
	80	+
	81	+exit $FAILED