feat: 修改AGMM pytorch接入方式，从customClass改为TORCH_LIBRARY，并更新对应skill-catccos-AtomGit

已合并

feat: 修改AGMM pytorch接入方式，从customClass改为TORCH_LIBRARY，并更新对应skill #51

i_soyabean创建于 6月10日从i_soyabean/catccos： feat/agmm-torch-library-registration 合入到 cann/catccos： master

feat: 修改AGMM pytorch接入方式，从customClass改为TORCH_LIBRARY，并更新对应skill #51

已合并

i_soyabean创建于 6月10日

从i_soyabean/catccos feat/agmm-torch-library-registration合入到cann/catccos master

共 13 个文件变更+604-685

Mexamples/CMakeLists.txt+7-0

	@@ -188,6 +188,11 @@ function(catccos_example_add_library TARGET)
188	list(APPEND COMPILE_DEFS -D${DEF})	188	list(APPEND COMPILE_DEFS -D${DEF})
189	endforeach()	189	endforeach()
190		190
		191	+ set(PROFILING_FLAGS "")
		192	+ if(ENABLE_TIMER)
		193	+ set(PROFILING_FLAGS -DENABLE_TIMER --cce-aicore-block-local-init)
		194	+ endif()
		195	+
191	# Create custom target for kernel library	196	# Create custom target for kernel library
192	add_custom_target(${TARGET}_kernel_build)	197	add_custom_target(${TARGET}_kernel_build)
193		198
	@@ -195,6 +200,7 @@ function(catccos_example_add_library TARGET)
195	add_custom_command(	200	add_custom_command(
196	OUTPUT ${CMAKE_BINARY_DIR}/lib/lib${TARGET}_kernel.so	201	OUTPUT ${CMAKE_BINARY_DIR}/lib/lib${TARGET}_kernel.so
197	COMMAND ${CCEC} -O2 -shared -fPIC -std=c++17 -xcce --cce-aicore-arch=${CCE_AICORE_ARCH_BASE} ${CATLASS_ARCH_FLAGS}	202	COMMAND ${CCEC} -O2 -shared -fPIC -std=c++17 -xcce --cce-aicore-arch=${CCE_AICORE_ARCH_BASE} ${CATLASS_ARCH_FLAGS}
		203	+ ${PROFILING_FLAGS}
198	-mllvm -cce-aicore-stack-size=0x8000	204	-mllvm -cce-aicore-stack-size=0x8000
199	-mllvm -cce-aicore-function-stack-size=0x8000	205	-mllvm -cce-aicore-function-stack-size=0x8000
200	-mllvm -cce-aicore-record-overflow=true	206	-mllvm -cce-aicore-record-overflow=true
	@@ -207,6 +213,7 @@ function(catccos_example_add_library TARGET)
207	-I${CMAKE_SOURCE_DIR}/3rdparty/catlass/include	213	-I${CMAKE_SOURCE_DIR}/3rdparty/catlass/include
208	-I${CMAKE_SOURCE_DIR}/examples	214	-I${CMAKE_SOURCE_DIR}/examples
209	-I${CMAKE_SOURCE_DIR}/utils	215	-I${CMAKE_SOURCE_DIR}/utils
		216	+ -I${CMAKE_SOURCE_DIR}/tools
210	-I${ASCEND_HOME_PATH}/include	217	-I${ASCEND_HOME_PATH}/include
211	-I${ASCEND_HOME_PATH}/compiler/tikcpp/tikcfw	218	-I${ASCEND_HOME_PATH}/compiler/tikcpp/tikcfw
212	-I${ASCEND_HOME_PATH}/compiler/tikcpp/tikcfw/impl	219	-I${ASCEND_HOME_PATH}/compiler/tikcpp/tikcfw/impl

Mexamples/allgather_matmul/allgather_matmul_device.h+2-1文件内容审核中，请稍后刷新重试

Mexamples/allgather_matmul/allgather_matmul_host.h+20-15文件内容审核中，请稍后刷新重试

Mexamples/allgather_matmul/allgather_matmul_wrapper.cpp+81-14文件内容审核中，请稍后刷新重试

Mexamples/allgather_matmul/main.cpp+3-2文件内容审核中，请稍后刷新重试

Mexamples/allgather_matmul/scripts/allgather_matmul.py+37-36

	@@ -22,7 +22,6 @@ from glob import glob
22	import argparse	22	import argparse
23	import torch	23	import torch
24	import torch_npu	24	import torch_npu
25	-import fcntl
26		25
27		26
28	def load_torch_library(lib_name):	27	def load_torch_library(lib_name):
	@@ -89,50 +88,52 @@ def worker(args):
89	torch_npu.npu.set_device(rank_id)	88	torch_npu.npu.set_device(rank_id)
90	local_mem_size = 1024 * 1024 * 1024	89	local_mem_size = 1024 * 1024 * 1024
91		90
92	- manager = torch.classes.CatccosOps.Manager()	91	+ status = torch.ops.catccos.init(rank_id, rank_size, local_mem_size, ip_port)
93	- manager.attr_init(rank_id, rank_size, local_mem_size, ip_port)	92	+ if status != 0:
		93	+ raise RuntimeError(f"Rank {rank_id}: torch.ops.catccos.init failed with status {status}")
94	print(f"Rank {rank_id}: ACLCATCCOS init success!")	94	print(f"Rank {rank_id}: ACLCATCCOS init success!")
95		95
96	- # Create AllGatherMatmul operator	96	+ try:
97	- ag_mm = torch.classes.CatccosOps.AllGatherMatmul()	97	+ # Load input tensors from files
		98	+ a_file = os.path.join(data_dir, f"rank_{rank_id}_a.bin")
		99	+ b_file = os.path.join(data_dir, f"rank_{rank_id}_b.bin")
98		100
99	- # Load input tensors from files	101	+ # Each rank has input shape (m, k) and (k, n)
100	- a_file = os.path.join(data_dir, f"rank_{rank_id}_a.bin")	102	+ # Output shape is (m * rankSize, n) for allgather_matmul
101	- b_file = os.path.join(data_dir, f"rank_{rank_id}_b.bin")	103	+ a_shape = (m, k)
		104	+ b_shape = (k, n)
102		105
103	- # Each rank has input shape (m, k) and (k, n)	106	+ tensor_a = load_tensor_from_file(a_file, a_shape, torch.float16)
104	- # Output shape is (m * rankSize, n) for allgather_matmul	107	+ tensor_b = load_tensor_from_file(b_file, b_shape, torch.float16)
105	- a_shape = (m, k)
106	- b_shape = (k, n)
107	- c_shape = (m * rank_size, n)
108		108
109	- tensor_a = load_tensor_from_file(a_file, a_shape, torch.float16)	109	+ # Move to NPU
110	- tensor_b = load_tensor_from_file(b_file, b_shape, torch.float16)	110	+ tensor_a_npu = tensor_a.npu()
		111	+ tensor_b_npu = tensor_b.npu()
111		112
112	- # Move to NPU	113	+ # Execute kernel
113	- tensor_a_npu = tensor_a.npu()	114	+ tensor_c_npu = torch.ops.catccos.allgather_matmul(tensor_a_npu, tensor_b_npu, rank_size)
114	- tensor_b_npu = tensor_b.npu()
115		115
116	- # Allocate output buffer	116	+ # Synchronize
117	- tensor_c_npu = torch.zeros(c_shape, dtype=torch.float16).npu()	117	+ torch_npu.npu.synchronize()
118		118
119	- # Execute kernel	119	+ # Write output to file (only rank 0 writes)
120	- ag_mm.compute(tensor_c_npu, tensor_a_npu, tensor_b_npu)	120	+ if rank_id == 0:
		121	+ c_output_file = os.path.join(data_dir, "output.bin")
		122	+ tensor_c_cpu = tensor_c_npu.cpu()
121		123
122	- # Synchronize	124	+ with open(c_output_file, "wb") as f:
123	- torch_npu.npu.synchronize()	125	+ f.write(tensor_c_cpu.numpy().tobytes())
124		126
125	- # Write output to file (only rank 0 writes)	127	+ print(f"Rank {rank_id}: Done! Output written to {c_output_file}")
126	- if rank_id == 0:	128	+ else:
127	- c_output_file = os.path.join(data_dir, "output.bin")	129	+ print(f"Rank {rank_id}: Done!")
128	- tensor_c_cpu = tensor_c_npu.cpu()	130	+ finally:
129	-	131	+ torch_npu.npu.synchronize()
130	- with open(c_output_file, "wb") as f:	132	+ finalize_status = torch.ops.catccos.finalize()
131	- f.write(tensor_c_cpu.numpy().tobytes())	133	+ if finalize_status != 0:
132	-	134	+ raise RuntimeError(
133	- print(f"Rank {rank_id}: Done! Output written to {c_output_file}")	135	+ f"Rank {rank_id}: torch.ops.catccos.finalize failed with status {finalize_status}"
134	- else:	136	+ )
135	- print(f"Rank {rank_id}: Done!")
136		137
137		138
138	if __name__ == "__main__":	139	if __name__ == "__main__":
	@@ -151,4 +152,4 @@ if __name__ == "__main__":
151		152
152	args = parser.parse_args()	153	args = parser.parse_args()
153		154
154	- worker(args)	155	+ worker(args)

Mexamples/allgather_matmul/scripts/build_python.sh+12-2

	@@ -15,8 +15,18 @@ PROJECT_ROOT=$( dirname $( dirname $(dirname "$SCRIPT_DIR")))
15	SOURCE_DIR=$PROJECT_ROOT	15	SOURCE_DIR=$PROJECT_ROOT
16	BUILD_DIR=$PROJECT_ROOT/build	16	BUILD_DIR=$PROJECT_ROOT/build
17		17
		18	+if ! command -v python3 >/dev/null 2>&1; then
		19	+ echo "[ERROR] python3 is required but not found in PATH"
		20	+ exit 1
		21	+fi
		22	+
		23	+if ! command -v cmake >/dev/null 2>&1; then
		24	+ echo "[ERROR] cmake is required but not found in PATH"
		25	+ exit 1
		26	+fi
		27	+
18	# Get Torch's CMake prefix path	28	# Get Torch's CMake prefix path
19	-TORCH_PREFIX_PATH="$(python -c 'import torch; print(torch.utils.cmake_prefix_path)')"	29	+TORCH_PREFIX_PATH="$(python3 -c 'import torch; print(torch.utils.cmake_prefix_path)')"
20		30
21	echo "Using TORCH_PREFIX_PATH=${TORCH_PREFIX_PATH}"	31	echo "Using TORCH_PREFIX_PATH=${TORCH_PREFIX_PATH}"
22	mkdir -p $BUILD_DIR	32	mkdir -p $BUILD_DIR
	@@ -27,4 +37,4 @@ cmake -B $BUILD_DIR -S $SOURCE_DIR \
27		37
28		38
29	echo "== Build target allgather_matmul =="	39	echo "== Build target allgather_matmul =="
30	-cmake --build $BUILD_DIR --target allgather_matmul_kernel_build catccos_torch -j	40	+cmake --build $BUILD_DIR --target allgather_matmul_kernel_build catccos_torch -j

Mexamples/allgather_matmul/scripts/run_python.sh+71-26

	@@ -27,44 +27,89 @@ if [ "$RANK_SIZE" -gt 8 ]; then
27	fi	27	fi
28		28
29	cd ${CURRENT_DIR}	29	cd ${CURRENT_DIR}
		30	+mkdir -p ./out
30	DATA_DIR=$(realpath ./out)	31	DATA_DIR=$(realpath ./out)
31	echo "DATA_DIR: $DATA_DIR"	32	echo "DATA_DIR: $DATA_DIR"
32	PYTHON_SCRIPT=${SCRIPT_DIR}/allgather_matmul.py	33	PYTHON_SCRIPT=${SCRIPT_DIR}/allgather_matmul.py
33		34
		35	+TOTAL_CASES=0
		36	+PASSED_CASES=0
		37	+FAILED_CASES=0
		38	+
34	# Read MNK from CSV	39	# Read MNK from CSV
35	-tail -n +2 "$CSV_FILE" \| while IFS=',' read -r M K N; do	40	+while IFS=',' read -r M K N; do
36	- echo "Processing test case: M=${M}, K=${K}, N=${N}"	41	+ if [ -z "$M" ]; then
		42	+ continue
		43	+ fi
		44	+
		45	+ TOTAL_CASES=$((TOTAL_CASES + 1))
		46	+ CASE_PASSED=1
		47	+ echo "Processing test case #${TOTAL_CASES}: M=${M}, K=${K}, N=${N}"
37		48
38	# Generate golden data	49	# Generate golden data
39	- rm -rf ./out/*.bin	50	+ rm -rf "${DATA_DIR}"/*.bin
40	- python3 ${UTILS_PATH}/gen_data.py "agmm" 1 ${RANK_SIZE} ${M} ${N} ${K} 0 0 ${DATA_DIR}	51	+ if ! python3 ${UTILS_PATH}/gen_data.py "agmm" 1 ${RANK_SIZE} ${M} ${N} ${K} 0 0 ${DATA_DIR}; then
		52	+ echo "[FAIL] Generate data failed for M=${M}, K=${K}, N=${N}"
		53	+ CASE_PASSED=0
		54	+ fi
41		55
42	# Set necessary parameters	56	# Set necessary parameters
43	IPPORT="tcp://127.0.0.1:8735"	57	IPPORT="tcp://127.0.0.1:8735"
44		58
45	- # Generate output file (m * rankSize * n * 2 bytes for fp16)	59	+ if [ ${CASE_PASSED} -eq 1 ]; then
46	- FILE_SIZE=$((M * RANK_SIZE * N * 2))	60	+ # Generate output file (m * rankSize * n * 2 bytes for fp16)
47	- dd if=/dev/zero of="${DATA_DIR}/output.bin" bs=${FILE_SIZE} count=1	61	+ FILE_SIZE=$((M * RANK_SIZE * N * 2))
48	- echo "Output File Created!"	62	+ if ! dd if=/dev/zero of="${DATA_DIR}/output.bin" bs=${FILE_SIZE} count=1; then
		63	+ echo "[FAIL] Create output file failed for M=${M}, K=${K}, N=${N}"
		64	+ CASE_PASSED=0
		65	+ else
		66	+ echo "Output File Created!"
		67	+ fi
		68	+ fi
49		69
50	- # Start Python processes	70	+ if [ ${CASE_PASSED} -eq 1 ]; then
51	- for (( idx = 0; idx < ${RANK_SIZE}; idx = idx + 1 )); do	71	+ # Start Python processes
52	- echo "Rank ${idx} started!"	72	+ PIDS=()
53	- python3 ${PYTHON_SCRIPT} \	73	+ for (( idx = 0; idx < ${RANK_SIZE}; idx = idx + 1 )); do
54	- --rank_size ${RANK_SIZE} \	74	+ echo "Rank ${idx} started!"
55	- --rank_id ${idx} \	75	+ python3 ${PYTHON_SCRIPT} \
56	- --m ${M} \	76	+ --rank_size ${RANK_SIZE} \
57	- --n ${N} \	77	+ --rank_id ${idx} \
58	- --k ${K} \	78	+ --m ${M} \
59	- --data_dir "${DATA_DIR}" \	79	+ --n ${N} \
60	- --ip_port "${IPPORT}" &	80	+ --k ${K} \
61	- done	81	+ --data_dir "${DATA_DIR}" \
		82	+ --ip_port "${IPPORT}" &
		83	+ PIDS+=($!)
		84	+ done
62		85
63	- # Wait until all process exit	86	+ # Wait until all process exit
64	- wait	87	+ for pid in "${PIDS[@]}"; do
		88	+ if ! wait "${pid}"; then
		89	+ CASE_PASSED=0
		90	+ fi
		91	+ done
		92	+ fi
65		93
66	- # Verify output	94	+ if [ ${CASE_PASSED} -eq 1 ]; then
67	- python3 ${UTILS_PATH}/verify_result.py ${DATA_DIR}/output.bin ${DATA_DIR}/golden.bin 1 $((M * RANK_SIZE)) ${N} ${K}	95	+ # Verify output
68	-done	96	+ if ! python3 ${UTILS_PATH}/verify_result.py ${DATA_DIR}/output.bin ${DATA_DIR}/golden.bin 1 $((M * RANK_SIZE)) ${N} ${K}; then
		97	+ CASE_PASSED=0
		98	+ fi
		99	+ fi
69		100
70	-cd ${CURRENT_DIR}	101	+ if [ ${CASE_PASSED} -eq 1 ]; then
		102	+ PASSED_CASES=$((PASSED_CASES + 1))
		103	+ echo "[PASS] M=${M}, K=${K}, N=${N}"
		104	+ else
		105	+ FAILED_CASES=$((FAILED_CASES + 1))
		106	+ echo "[FAIL] M=${M}, K=${K}, N=${N}"
		107	+ fi
		108	+done < <(tail -n +2 "$CSV_FILE")
		109	+
		110	+echo "AGMM Python test summary: total=${TOTAL_CASES}, pass=${PASSED_CASES}, fail=${FAILED_CASES}"
		111	+
		112	+cd ${CURRENT_DIR}
		113	+if [ ${FAILED_CASES} -ne 0 ]; then
		114	+ exit 1
		115	+fi

Mexamples/torch_binding/CMakeLists.txt+1-0

	@@ -42,6 +42,7 @@ message("Start processing torch_binding")
42	# Compile torch_binding as shared library	42	# Compile torch_binding as shared library
43	add_library(catccos_torch SHARED	43	add_library(catccos_torch SHARED
44	src/torch_bindings.cpp	44	src/torch_bindings.cpp
		45	+ src/torch_bindings_meta.cpp
45	)	46	)
46		47
47	set_target_properties(catccos_torch PROPERTIES	48	set_target_properties(catccos_torch PROPERTIES

Dexamples/torch_binding/include/torch_register.h+0-57

	@@ -1,57 +0,0 @@
1	-/**
2	- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3	- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4	- * 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,
7	- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8	- * See LICENSE in the root of the software repository for the full text of the License.
9	- */
10	-#ifndef CATCCOS_TORCH_REGISTER_H
11	-#define CATCCOS_TORCH_REGISTER_H
12	-
13	-#define REGISTER_CATCCOS_OPS_CLASS(CLASS_NAME, ...) \
14	- static auto registry_##CLASS_NAME = \
15	- torch::jit::class_<CatccosOps::CLASS_NAME>("CatccosOps", #CLASS_NAME) \
16	- .def(torch::jit::init<>()) \
17	- REGISTER_CATCCOS_OPS_FUNCS_HELPER(CLASS_NAME, ##__VA_ARGS__)
18	-
19	-#define REGISTER_CATCCOS_OPS_FUNCS_HELPER(CLASS, ...) \
20	- REGISTER_CATCCOS_OPS_FUNCS_CHOOSER(__VA_ARGS__, REGISTER_CATCCOS_OPS_FUNCS_6, REGISTER_CATCCOS_OPS_FUNCS_5, REGISTER_CATCCOS_OPS_FUNCS_4, REGISTER_CATCCOS_OPS_FUNCS_3, REGISTER_CATCCOS_OPS_FUNCS_2, REGISTER_CATCCOS_OPS_FUNCS_1)(CLASS, ##__VA_ARGS__)
21	-
22	-#define REGISTER_CATCCOS_OPS_FUNCS_CHOOSER(_1, _2, _3, _4, _5, _6, FUNC, ...) FUNC
23	-
24	-#define REGISTER_CATCCOS_OPS_FUNCS_1(CLASS, FUNC1) \
25	- .def(#FUNC1, &CatccosOps::CLASS::FUNC1)
26	-
27	-#define REGISTER_CATCCOS_OPS_FUNCS_2(CLASS, FUNC1, FUNC2) \
28	- .def(#FUNC1, &CatccosOps::CLASS::FUNC1) \
29	- .def(#FUNC2, &CatccosOps::CLASS::FUNC2)
30	-
31	-#define REGISTER_CATCCOS_OPS_FUNCS_3(CLASS, FUNC1, FUNC2, FUNC3) \
32	- .def(#FUNC1, &CatccosOps::CLASS::FUNC1) \
33	- .def(#FUNC2, &CatccosOps::CLASS::FUNC2) \
34	- .def(#FUNC3, &CatccosOps::CLASS::FUNC3)
35	-
36	-#define REGISTER_CATCCOS_OPS_FUNCS_4(CLASS, FUNC1, FUNC2, FUNC3, FUNC4) \
37	- .def(#FUNC1, &CatccosOps::CLASS::FUNC1) \
38	- .def(#FUNC2, &CatccosOps::CLASS::FUNC2) \
39	- .def(#FUNC3, &CatccosOps::CLASS::FUNC3) \
40	- .def(#FUNC4, &CatccosOps::CLASS::FUNC4)
41	-
42	-#define REGISTER_CATCCOS_OPS_FUNCS_5(CLASS, FUNC1, FUNC2, FUNC3, FUNC4, FUNC5) \
43	- .def(#FUNC1, &CatccosOps::CLASS::FUNC1) \
44	- .def(#FUNC2, &CatccosOps::CLASS::FUNC2) \
45	- .def(#FUNC3, &CatccosOps::CLASS::FUNC3) \
46	- .def(#FUNC4, &CatccosOps::CLASS::FUNC4) \
47	- .def(#FUNC5, &CatccosOps::CLASS::FUNC5)
48	-
49	-#define REGISTER_CATCCOS_OPS_FUNCS_6(CLASS, FUNC1, FUNC2, FUNC3, FUNC4, FUNC5, FUNC6) \
50	- .def(#FUNC1, &CatccosOps::CLASS::FUNC1) \
51	- .def(#FUNC2, &CatccosOps::CLASS::FUNC2) \
52	- .def(#FUNC3, &CatccosOps::CLASS::FUNC3) \
53	- .def(#FUNC4, &CatccosOps::CLASS::FUNC4) \
54	- .def(#FUNC5, &CatccosOps::CLASS::FUNC5) \
55	- .def(#FUNC6, &CatccosOps::CLASS::FUNC6)
56	-
57	-#endif // CATCCOS_TORCH_REGISTER_H

Mexamples/torch_binding/src/torch_bindings.cpp+180-154文件内容审核中，请稍后刷新重试

Aexamples/torch_binding/src/torch_bindings_meta.cpp+36-0

@@ -0,0 +1,36 @@
	1	+/**
	2	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	4	+ * 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,
	7	+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	8	+ * See LICENSE in the root of the software repository for the full text of the License.
	9	+ */
	10	+
	11	+#include <torch/torch.h>
	12	+#include <torch/library.h>
	13	+
	14	+#include <vector>
	15	+
	16	+namespace catccos::meta {
	17	+
	18	+at::Tensor allgather_matmul_meta(
	19	+ const at::Tensor& a,
	20	+ const at::Tensor& b,
	21	+ int64_t rank_size)
	22	+{
	23	+ auto a_sizes = a.sym_sizes();
	24	+ auto b_sizes = b.sym_sizes();
	25	+ std::vector<c10::SymInt> out_shape = {
	26	+ a_sizes[0] * c10::SymInt(rank_size),
	27	+ b_sizes[1],
	28	+ };
	29	+ return at::empty_symint(out_shape, a.options());
	30	+}
	31	+
	32	+} // namespace catccos::meta
	33	+
	34	+TORCH_LIBRARY_IMPL(catccos, Meta, m) {
	35	+ m.impl("allgather_matmul", &catccos::meta::allgather_matmul_meta);
	36	+}

Mskills/torch-binding/SKILL.md+154-378

	@@ -1,419 +1,195 @@
1	---	1	---
2	name: torch-binding	2	name: torch-binding
3	-description: Use when integrating catccos operators into PyTorch, implementing wrapper functions, custom classes, and binding Python to C++ kernels.	3	+description: Use when integrating catccos operators into PyTorch through TORCH_LIBRARY, adding wrapper functions, kernel shared libraries, Meta kernels, torch.ops.catccos runtime bindings, and Python verification scripts.
4	---	4	---
5		5
6	-# 操作步骤	6	+# catccos PyTorch TORCH_LIBRARY 接入
7		7
8	-以 allgather_matmul 的 pytorch 算子接入为例子。整体的实现步骤如下：	8	+这个 skill 用于把 catccos 算子接入为 PyTorch 自定义算子。当前默认路线是 `TORCH_LIBRARY`，Python 调用面是 `torch.ops.catccos.*`。
9		9
10	-1. 使用 `Config::Device` + `DeviceOp::Run()` 模式为算子实现 wrapper 函数。	10	+不要再使用旧 custom class 路线：不要新增 `torch::jit::CustomClassHolder`、`REGISTER_CATCCOS_OPS_CLASS`、`torch_register.h` 或 `torch.classes.CatccosOps.*`。
11	-2. 编译算子和 wrapper 文件为动态链接库 .so。
12	-3. 在 `torch_bindings.cpp` 中创建算子类，继承 `torch::jit::CustomClassHolder`。
13	-4. 把算子注册成 torch 算子，编译 `torch_binding` 为 .so，并链接之前的算子 .so。
14	-5. Python 中加载动态链接库，调用并运行算子。
15	-6. 测试数据生成和验证的脚本编写和调用。
16		11
17	-以下是涉及的需要修改的关键文件的目录结构。	12	+## 核心结构
18		13
19	-```	14	+新增或改造一个算子时，通常涉及这些位置：
20	-└── catccos/	15	+
21	- ├── examples/	16	+```text
22	- │ ├── <算子名>/	17	+catccos/
23	- │ │ ├── scripts/	18	+├── examples/
24	- │ │ │ ├── build_python.sh	19	+│ ├── <op>/
25	- │ │ │ ├── run_python.sh	20	+│ │ ├── <op>_wrapper.cpp
26	- │ │ │ └── <算子名>.py	21	+│ │ ├── <op>_device.h
27	- │ │ ├── <算子名>_wrapper.cpp # 新增，wrapper 函数	22	+│ │ ├── CMakeLists.txt
28	- │ │ ├── <算子名>_device.h # 包含 Config 结构体	23	+│ │ └── scripts/
29	- │ │ ├── main.cpp	24	+│ │ ├── build_python.sh
30	- │ │ └── CMakeLists.txt	25	+│ │ ├── run_python.sh
31	- │ ├── torch_binding/	26	+│ │ ├── <op>.py
32	- │ │ ├── include/	27	+│ │ └── test_shapes.csv
33	- │ │ │ ├── catccos_torch_kernel.h # 新增 wrapper 函数声明	28	+│ ├── torch_binding/
34	- │ │ │ └── torch_register.h	29	+│ │ ├── include/
35	- │ │ ├── src/	30	+│ │ │ └── catccos_torch_kernel.h
36	- │ │ │ └── torch_bindings.cpp # 算子类 + torch 注册	31	+│ │ ├── src/
37	- │ │ └── CMakeLists.txt	32	+│ │ │ ├── torch_bindings.cpp
38	- │ ├── utils/	33	+│ │ │ └── torch_bindings_meta.cpp
39	- │ │ └── shmem_init.h # 包含 set_attr 函数	34	+│ │ └── CMakeLists.txt
40	- │ └── CMakeLists.txt # catccos_example_add_library 函数	35	+│ └── CMakeLists.txt
41	- └── CMakeLists.txt	36	+└── include/
42	```	37	```
43		38
44	----	39	+AGMM 的现有实现是可参考样例：
45		40
46	-## 前置检查：确认 main.cpp 使用新 API	41	+- kernel wrapper: `examples/allgather_matmul/allgather_matmul_wrapper.cpp`
		42	+- wrapper 声明: `examples/torch_binding/include/catccos_torch_kernel.h`
		43	+- runtime 注册: `examples/torch_binding/src/torch_bindings.cpp`
		44	+- Meta 注册: `examples/torch_binding/src/torch_bindings_meta.cpp`
		45	+- Python 入口: `examples/allgather_matmul/scripts/allgather_matmul.py`
47		46
48	-在写 wrapper 之前，务必先检查该算子的 `main.cpp` 是否已经升级：	47	+## 接入步骤
49		48
50	-旧 API（不应再使用）：	49	+1. 在算子目录实现 wrapper
51	-```cpp	50	+ - 在 `examples/<op>/<op>_wrapper.cpp` 中实现 `CatccosKernel::catccos_<op>_wrapper(...)`。
52	-AllGatherMatmul<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC>	51	+ - wrapper 负责设置 `CocTilingParams`、构造 device op arguments、选择 kernel config、申请必要 workspace，并调用 device op。
53	- <<<BLOCK_NUM, nullptr, stream>>>(fftsAddr, aPtr, bPtr, cPtr, gmSymmetric, cocTiling);	52	+ - wrapper 和 native example 的 tiling 参数必须逐项对齐。不要从其他算子复制 `commBlockShape`、`commTileShape` 等参数后直接假设通用。
54	-```
55		53
56	-新 API（以此为模板写 wrapper）：	54	+2. 在统一头文件声明 wrapper
57	-```cpp	55	+ - 在 `examples/torch_binding/include/catccos_torch_kernel.h` 中添加 wrapper 函数声明。
58	-using Config = AllGatherMatmulConfig_M0_128<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC>;	56	+ - 声明、定义、`torch_bindings.cpp` 调用点的签名必须一致。
59	-using DeviceOp = Config::Device;	57	+ - 常见参数包括 `uint32_t blockDim`、`aclrtStream stream`、`uint64_t fftsAddr`、输入/输出 device pointer、symmetric workspace pointer、shape 和 rank 信息。
60		58
61	-DeviceOp::Arguments args{ problemShape, rankId, rankSize, commInterval,	59	+3. 编译 kernel wrapper so
62	- aPtr, bPtr, cPtr, gmSymmetric,	60	+ - 在 `examples/<op>/CMakeLists.txt` 中使用：
63	- commCoreSplit, commBlockShape, commTileShape };
64	-
65	-DeviceOp deviceOp;
66	-deviceOp.Initialize(args);
67	-deviceOp.Run(stream, BLOCK_NUM, fftsAddr);
68	-```
69	-
70	-如果 main.cpp 还用的是旧 API，说明 kernel 头文件还没升级，wrapper 也先不要动。
71	-
72	----
73	-
74	-# 算子接入需要修改的文件清单
75	-
76	-\| 序号 \| 文件路径 \| 修改内容 \|
77	-\|------\|----------\|----------\|
78	-\| 1 \| `examples/<算子名>/<算子名>_wrapper.cpp` \| 新建，Config::Device 模式的 wrapper 实现 \|
79	-\| 2 \| `examples/torch_binding/include/catccos_torch_kernel.h` \| 添加 wrapper 函数声明（含 rankId，uint8_t*） \|
80	-\| 3 \| `examples/torch_binding/src/torch_bindings.cpp` \| 添加算子类实现 + 注册 \|
81	-\| 4 \| `examples/CMakeLists.txt` \| 在 `CATCCOS_TORCH_SUPPORTED` 中添加算子名 \|
82	-\| 5 \| `examples/<算子名>/CMakeLists.txt` \| 添加 `catccos_example_add_library` + `COMPILE_DEFINITIONS` \|
83	-\| 6 \| `examples/<算子名>/scripts/<算子名>.py` \| 新建，Python 调用脚本 \|
84	-\| 7 \| `examples/<算子名>/scripts/build_python.sh` \| 新建，编译脚本 \|
85	-\| 8 \| `examples/<算子名>/scripts/run_python.sh` \| 新建，运行脚本 \|
86	-
87	----
88	-
89	-## 各文件修改详情
90	-
91	-### 1. <算子名>_wrapper.cpp — 新建 wrapper 实现
92	-
93	-使用 `Config::Device` + `DeviceOp::Run()` 模式，参考 main.cpp 中对应该算子的调用方式。
94	-
95	-*所有指针参数类型必须为 `uint8_t`**，和 main.cpp 保持一致。
96	-
97	-必须新增 `rankId` 参数，传递给 `DeviceOp::Arguments`。
98	-
99	-```cpp
100	-#include "<算子名>_device.h"
101	-
102	-using namespace AscendC;
103	-using namespace Catccos;
104	-
105	-using LayoutA = Catlass::layout::RowMajor;
106	-using LayoutB = Catlass::layout::RowMajor;
107	-using LayoutC = Catlass::layout::RowMajor;
108	-
109	-using ElementA = half;
110	-using ElementB = half;
111	-using ElementC = half;
112	-
113	-namespace CatccosKernel {
114	-
115	-void catccos_<算子名>_wrapper(
116	- uint32_t blockDim,
117	- aclrtStream stream,
118	- uint64_t fftsAddr,
119	- uint8_t* aPtr,
120	- uint8_t* bPtr,
121	- uint8_t* cPtr,
122	- uint8_t* gmSymmetric,
123	- uint32_t m,
124	- uint32_t n,
125	- uint32_t k,
126	- int rankId,
127	- int rankSize)
128	-{
129	- using Config = <算子名>Config_M0_128<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC>;
130	- using DeviceOp = Config::Device;
131	-
132	- CocTilingParams cocTiling;
133	- cocTiling.m = m;
134	- cocTiling.n = n;
135	- cocTiling.k = k;
136	- cocTiling.m0 = 128;
137	- cocTiling.n0 = 256;
138	- cocTiling.k0 = 256;
139	- cocTiling.commTileM = 64;
140	- cocTiling.commInterval = 3;
141	- cocTiling.commNpuSplit = 1;
142	- cocTiling.commDataSplit = 20;
143	- cocTiling.commBlockM = 64;
144	- cocTiling.rankSize = rankSize;
145	-
146	- Catlass::GemmCoord problemShape{m, n, k};
147	- Catlass::MatrixCoord commCoreSplit{cocTiling.commDataSplit, cocTiling.commNpuSplit};
148	- Catlass::MatrixCoord commBlockShape{cocTiling.commBlockM, UINT_MAX / 2};
149	- Catlass::MatrixCoord commTileShape{cocTiling.commTileM / 2, cocTiling.n0};
150	-
151	- DeviceOp::Arguments args{
152	- problemShape,
153	- static_cast<uint32_t>(rankId), static_cast<uint32_t>(rankSize),
154	- cocTiling.commInterval,
155	- aPtr, bPtr, cPtr, gmSymmetric,
156	- commCoreSplit, commBlockShape, commTileShape
157	- };
158	-
159	- DeviceOp deviceOp;
160	- deviceOp.Initialize(args);
161	- deviceOp.Run(stream, blockDim, fftsAddr);
162	-}
163	-
164	-} // namespace CatccosKernel
165	-```
166	-
167	-Config 类型名规则：查看对应算子的 `<算子名>_device.h`，命名规则为 `<算子名>Config_M0_<tile大小>`。例如：
168	-- `AllGatherMatmulConfig_M0_128`
169	-- `MatmulReduceScatterConfig_M0_128`
170	-- `MatmulAllReduceConfig_M0_128`
171	-
172	-### 2. catccos_torch_kernel.h — 添加 wrapper 声明
173	-
174	-在命名空间 `CatccosKernel` 中添加：
175	-
176	-```cpp
177	-/**
178	- * @brief <算子名> kernel wrapper
179	- * @param blockDim number of AICore blocks
180	- * @param stream ACL stream
181	- * @param fftsAddr FFTS configuration address
182	- * @param aPtr pointer to matrix A on device
183	- * @param bPtr pointer to matrix B on device
184	- * @param cPtr pointer to matrix C on device (output)
185	- * @param gmSymmetric pointer to symmetric memory
186	- * @param m rows of matrix A
187	- * @param n columns of matrix B
188	- * @param k columns of matrix A / rows of matrix B
189	- * @param rankId current rank index
190	- * @param rankSize number of ranks
191	- */
192	-void catccos_<算子名>_wrapper(
193	- uint32_t blockDim,
194	- aclrtStream stream,
195	- uint64_t fftsAddr,
196	- uint8_t* aPtr,
197	- uint8_t* bPtr,
198	- uint8_t* cPtr,
199	- uint8_t* gmSymmetric,
200	- uint32_t m,
201	- uint32_t n,
202	- uint32_t k,
203	- int rankId,
204	- int rankSize
205	-);
206	-```
207	-
208	-重要：所有指针参数类型为 `uint8_t`，不是 `void`。
209	-
210	-### 3. torch_bindings.cpp — 添加算子类
211	-
212	-#### 3.1 在 namespace CatccosOps 中添加算子类
213	-
214	-```cpp
215	-class <算子类名> : public torch::jit::CustomClassHolder {
216	-public:
217	- <算子类名>() : name_("<算子名>"), count_(0), fftsAddr_(shmemx_get_ffts_config()), symmPtr_(nullptr)
218	- {
219	- // 分配对称内存（shmem_malloc，用于跨 rank 通信数据交换）
220	- symmPtr_ = static_cast<uint8_t*>(shmem_malloc(SHMEM_BUFF_BYTES));
221	- aclrtMemset(symmPtr_, SHMEM_BUFF_BYTES, 0, SHMEM_BUFF_BYTES);
222	- }
223	-
224	- ~<算子类名>()
225	- {
226	- if (symmPtr_ != nullptr) {
227	- shmem_free(symmPtr_);
228	- symmPtr_ = nullptr;
229	- }
230	- }
231	-
232	- std::string get_name() const
233	- {
234	- return name_;
235	- }
236	-
237	- void compute(const at::Tensor& c_tensor,
238	- const at::Tensor& a_tensor,
239	- const at::Tensor& b_tensor)
240	- {
241	- // 1. 参数校验
242	- TORCH_CHECK(a_tensor.dtype() == at::kHalf,
243	- "Compute Error: Only float16 (half) is supported! ");
244	- TORCH_CHECK(b_tensor.dtype() == at::kHalf,
245	- "Compute Error: Only float16 (half) is supported! ");
246	- TORCH_CHECK(c_tensor.dtype() == at::kHalf,
247	- "Compute Error: Only float16 (half) is supported! ");
248	- TORCH_CHECK(a_tensor.device().type() == c10::DeviceType::PrivateUse1,
249	- "Compute Error: Only NPU device is supported! ");
250	- TORCH_CHECK(b_tensor.device().type() == c10::DeviceType::PrivateUse1,
251	- "Compute Error: Only NPU device is supported! ");
252	- TORCH_CHECK(c_tensor.device().type() == c10::DeviceType::PrivateUse1,
253	- "Compute Error: Only NPU device is supported! ");
254	-
255	- // 2. 形状校验
256	- TORCH_CHECK(a_tensor.dim() == 2, "A tensor must be 2D!");
257	- TORCH_CHECK(b_tensor.dim() == 2, "B tensor must be 2D!");
258	- TORCH_CHECK(c_tensor.dim() == 2, "C tensor must be 2D!");
259	-
260	- int64_t m = a_tensor.size(0);
261	- int64_t k = a_tensor.size(1);
262	- int64_t n = b_tensor.size(1);
263	- TORCH_CHECK(b_tensor.size(0) == k, "A/K mismatch!");
264	-
265	- int32_t n_pes = shmem_n_pes();
266	-
267	- // 根据算子类型校验输出形状：
268	- // - MatmulReduceScatter: c_shape = (m / rankSize, n)
269	- // - AllGatherMatmul: c_shape = (m * rankSize, n)
270	- // - MatmulAllReduce: c_shape = (m, n)
271	-
272	- // 3. Make tensors contiguous
273	- at::Tensor a_contig = a_tensor.contiguous();
274	- at::Tensor b_contig = b_tensor.contiguous();
275	- at::Tensor c_contig = c_tensor.contiguous();
276	-
277	- // 4. Get device pointers
278	- // 注意：必须使用 storage().data()，不能使用 data_ptr<T>()（对 NPU tensor 不兼容）
279	- uint8_t* aPtr = static_cast<uint8_t>(const_cast<void>(a_contig.storage().data()));
280	- uint8_t* bPtr = static_cast<uint8_t>(const_cast<void>(b_contig.storage().data()));
281	- uint8_t* cPtr = static_cast<uint8_t>(const_cast<void>(c_contig.storage().data()));
282	-
283	- // 5. Get NPU stream and rank info
284	- aclrtStream stream = c10_npu::getCurrentNPUStream().stream(false);
285	- int32_t my_pe = shmem_my_pe(); // 获取当前 rank id
286	- count_++;
287	-
288	- // 6. Call wrapper
289	- CatccosKernel::catccos_<算子名>_wrapper(
290	- BLOCK_NUM,
291	- stream,
292	- fftsAddr_,
293	- aPtr,
294	- bPtr,
295	- cPtr,
296	- symmPtr_,
297	- m,
298	- n,
299	- k,
300	- my_pe, // rankId
301	- n_pes // rankSize
302	- );
303	- }
304	-
305	-private:
306	- std::string name_;
307	- int32_t count_;
308	- uint64_t fftsAddr_;
309	- uint8_t* symmPtr_;
310	- uint32_t BLOCK_NUM = 20;
311	-};
312	-```
313	-
314	-#### 3.2 注册算子
315	-
316	-```cpp
317	-REGISTER_CATCCOS_OPS_CLASS(<算子类名>, compute, get_name);
318	-```
319	-
320	-#### 3.3 其他注意事项
321	-
322	-- `set_attr` 函数：不要在 torch_bindings.cpp 中内联定义。使用 `examples/utils/shmem_init.h` 中的版本（`#include "utils/shmem_init.h"`）。
323	-- Manager::attr_init：必须添加错误检查（见下方 Code Review 清单）。
324	-- `symmPtr_` vs `workspace_ptr_`：`symmPtr_` 用 `shmem_malloc` 分配，用于跨 rank 通信数据交换（所有通信算子都需要）；`workspace_ptr_` 用 `aclrtMalloc` 分配，用于算子内部中间计算缓冲（仅部分算子如 dequant、dispatch 需要）。两者独立，不可混淆。
325	-- `SHMEM_BUFF_BYTES`：定义在 `utils/info.h` 中，建议用它替代硬编码的大小。
326	-
327	-### 4. examples/CMakeLists.txt — 添加到白名单
328	-
329	-在 `CATCCOS_TORCH_SUPPORTED` 中添加算子名：
330		61
331	```cmake	62	```cmake
332	-set(CATCCOS_TORCH_SUPPORTED
333	- allgather_matmul
334	- matmul_reduce_scatter # 新增
335	-)
336	-```
337	-
338	-### 5. <算子名>/CMakeLists.txt — 添加编译
339	-
340	-```cmake
341	-# Compile wrapper to shared library for PyTorch extension
342	if(CATCCOS_TORCH_EXTENSION)	63	if(CATCCOS_TORCH_EXTENSION)
343	catccos_example_add_library(	64	catccos_example_add_library(
344	- <算子名>	65	+ <op>
345	- SOURCES <算子名>_wrapper.cpp	66	+ SOURCES <op>_wrapper.cpp
346	- COMPILE_DEFINITIONS CATLASS_ARCH=2201 # 必须加！	67	+ COMPILE_DEFINITIONS CATLASS_ARCH=2201
347	)	68	)
348	endif()	69	endif()
349	```	70	```
350		71
351	-### 6. Python 脚本	72	+ - 在 `examples/CMakeLists.txt` 的 `CATCCOS_TORCH_SUPPORTED` 中加入 `<op>`。
		73	+ - `catccos_example_add_library(<op>)` 会生成 `build/lib/lib<op>_kernel.so`，`catccos_torch` 会链接这些 kernel so。
352		74
353	-根据算子类型确定输出规模：	75	+4. 添加 runtime PyTorch op
		76	+ - 在 `examples/torch_binding/src/torch_bindings.cpp` 中添加 C++ runtime 函数，例如：
354		77
355	-\| 算子类型 \| 输出形状 \| 输出文件写入 \|	78	+```cpp
356	-\|----------\|----------\|-------------\|	79	+namespace catccos {
357	-\| matmul_reduce_scatter \| (m/rankSize, n) \| 所有 rank 按偏移写入 \|
358	-\| allgather_matmul \| (m*rankSize, n) \| rank 0 写入 \|
359	-\| matmul_allreduce \| (m, n) \| rank 0 写入 \|
360		80
361	-### 7. 编译脚本编写规范	81	+at::Tensor <op>(const at::Tensor& a, const at::Tensor& b, int64_t rank_size)
		82	+{
		83	+ c10_npu::OptionalNPUGuard guard(a.device());
		84	+ // 校验 dtype/device/shape/init state
		85	+ // 创建输出 tensor
		86	+ // 获取 aclrtStream、FFTS 地址、symmetric workspace
362		87
363	-编译脚本参考 `allgather_matmul` 编译脚本 `build_python.sh` 的实现。	88	+ at_npu::native::OpCommand cmd;
		89	+ cmd.Name("catccos_<op>");
		90	+ cmd.Input(a_contig);
		91	+ cmd.Input(b_contig);
		92	+ cmd.Output(out);
		93	+ cmd.SetCustomHandler([...]() -> int {
		94	+ CatccosKernel::catccos_<op>_wrapper(...);
		95	+ return 0;
		96	+ });
		97	+ cmd.Run();
		98	+ return out;
		99	+}
364		100
365	-### 8. 运行脚本编写规范	101	+} // namespace catccos
		102	+```
366		103
367	-测试参数读取：必须从 CSV 读取测试参数，参考现有实现。	104	+ - lifecycle 通过 `torch.ops.catccos.init(...)` 和 `torch.ops.catccos.finalize()` 管理。
		105	+ - 不要把 symmetric workspace 绑到 custom class 构造/析构；`TORCH_LIBRARY` 路线没有 custom class 实例生命周期。
368		106
369	-测试数据生成：生成测试数据时，应参考对应算子目录下的 `run.sh` 中调用的脚本。	107	+5. 注册 schema 和 backend impl
		108	+ - 在 `TORCH_LIBRARY(catccos, m)` 中定义 schema 并绑定实现：
370		109
371	-如果 utils 中没有对应的脚本：	110	+```cpp
372	-- 需要开发者手动实现测试数据生成脚本	111	+TORCH_LIBRARY(catccos, m) {
373	-- 参考其他 gen_*.py 脚本的格式	112	+ m.def("init(int rank_id, int rank_size, int local_mem_size, str ip_port) -> int");
374	-- 路径: `examples/utils/gen_<算子名>_data.py`	113	+ m.def("<op>(Tensor a, Tensor b, int rank_size) -> Tensor");
		114	+ m.def("finalize() -> int");
375		115
376	----	116	+ m.impl("init", &catccos::init);
		117	+ m.impl("<op>", torch::kPrivateUse1, &catccos::<op>);
		118	+ m.impl("finalize", &catccos::finalize);
		119	+}
		120	+```
377		121
378	-## Code Review 检查清单	122	+ - 多个算子共享同一个 `catccos` namespace 时，不要重复定义已有 schema；按当前文件组织合并到同一个注册块。
379		123
380	-接入新算子后，按以下清单进行检查：	124	+6. 添加 Meta kernel
		125	+ - 在 `examples/torch_binding/src/torch_bindings_meta.cpp` 中为算子添加 Meta 实现。
		126	+ - 对需要入图或符号 shape 的算子，优先使用 `sym_sizes()` 和 `at::empty_symint(...)`：
381		127
382	-### 1. wrapper 函数	128	+```cpp
383	-- [ ] 使用 `Config::Device` + `DeviceOp::Run()` 模式（不是旧 `<<<>>>` 直接调用）	129	+at::Tensor <op>_meta(const at::Tensor& a, const at::Tensor& b, int64_t rank_size)
384	-- [ ] 包含 `rankId` 参数，传递给 `DeviceOp::Arguments`	130	+{
385	-- [ ] 所有指针参数类型为 `uint8_t`（不是 `void`）	131	+ auto a_sizes = a.sym_sizes();
386	-- [ ] CocTilingParams 参数和 main.cpp 中一致	132	+ auto b_sizes = b.sym_sizes();
387	-- [ ] 头文件声明与实现签名一致	133	+ std::vector<c10::SymInt> out_shape = {
		134	+ a_sizes[0] * c10::SymInt(rank_size),
		135	+ b_sizes[1],
		136	+ };
		137	+ return at::empty_symint(out_shape, a.options());
		138	+}
388		139
389	-### 2. 算子类（torch_bindings.cpp）	140	+TORCH_LIBRARY_IMPL(catccos, Meta, m) {
390	-- [ ] 继承 `torch::jit::CustomClassHolder`	141	+ m.impl("<op>", &catccos::meta::<op>_meta);
391	-- [ ] `symmPtr_` 类型为 `uint8_t`，`shmem_malloc` 后 `static_cast<uint8_t>`	142	+}
392	-- [ ] 构造函数中用 `shmem_malloc(SHMEM_BUFF_BYTES)` 分配对称内存，析构函数中 `shmem_free`	143	+```
393	-- [ ] 指针提取使用 `static_cast<uint8_t>(const_cast<void>(tensor.storage().data()))`，不使用 `data_ptr<T>()`
394	-- [ ] compute() 中调用 `shmem_my_pe()` 获取 rankId 并传给 wrapper
395	-- [ ] 输出形状校验符合算子语义
396		144
397	-### 3. torch_bindings.cpp 全局	145	+7. 更新 Python 验证入口
398	-- [ ] 使用 `#include "utils/shmem_init.h"` 中的 `set_attr`，不在文件中内联定义	146	+ - 正式脚本放在 `examples/<op>/scripts/` 下。
399	-- [ ] Manager::attr_init 有错误检查（`set_conf_store_tls` 和 `init_attr` 都检查返回值）	147	+ - Python 中加载 `libcatccos_torch.so` 后，只通过 `torch.ops.catccos.*` 调用：
400		148
401	-### 4. CMake 配置	149	+```python
402	-- [ ] 算子名在 `CATCCOS_TORCH_SUPPORTED` 中	150	+status = torch.ops.catccos.init(rank_id, rank_size, local_mem_size, ip_port)
403	-- [ ] `catccos_example_add_library` 调用了 `COMPILE_DEFINITIONS CATLASS_ARCH=2201`	151	+out = torch.ops.catccos.<op>(tensor_a_npu, tensor_b_npu, rank_size)
404	-- [ ] `examples/CMakeLists.txt` 的 `add_library` 函数支持 `COMPILE_DEFINITIONS` 参数	152	+finalize_status = torch.ops.catccos.finalize()
405	-- [ ] `add_library` 函数有 `ARCH` 检测逻辑（a5 → dav-c310, 其他 → dav-c220）	153	+```
406	-- [ ] `add_library` 链接了 `-lm`
407		154
408	-### 5. 指针类型一致性	155	+ - 输出 tensor 推荐由 runtime op 返回，不再要求 Python 预先构造输出 tensor 传给 custom class。
409	-- [ ] `catccos_torch_kernel.h` 声明中指针为 `uint8_t*`
410	-- [ ] `wrapper.cpp` 实现中指针为 `uint8_t*`
411	-- [ ] `torch_bindings.cpp` 中成员变量 `symmPtr_`、`workspace_ptr_` 为 `uint8_t*`
412		156
413	-### 6. Python 脚本	157	+8. 更新 build/run 脚本
414	-- [ ] 输出规模与算子语义匹配	158	+ - `build_python.sh` 应配置 `-DCATCCOS_TORCH_EXTENSION=ON`，并构建 `<op>_kernel_build catccos_torch`。
415	-- [ ] 数据文件路径正确	159	+ - `run_python.sh` 复用对应算子的 CSV、数据生成和 `verify_result.py`。
		160	+ - 临时验证脚本只放在算子目录 `scripts/` 下，正式提交前删除；正式保留 `<op>.py`、`run_python.sh`、`build_python.sh`、`test_shapes.csv`。
416		161
417	-### 7. 运行脚本	162	+## 数据生成和验证
418	-- [ ] 使用正确的 gen_data 脚本	163	+
419	-- [ ] 从 CSV 读取 MNK，不能硬编码	164	+先查看算子原有 `scripts/run.sh`，确认它调用哪个 `examples/utils` 数据生成脚本。不要凭名字猜。
		165	+
		166	+常见形式：
		167	+
		168	+```bash
		169	+python3 ${UTILS_PATH}/gen_data.py "agmm" 1 ${RANK_SIZE} ${M} ${N} ${K} 0 0 ${DATA_DIR}
		170	+python3 ${UTILS_PATH}/verify_result.py ${DATA_DIR}/output.bin ${DATA_DIR}/golden.bin 1 $((M * RANK_SIZE)) ${N} ${K}
		171	+```
		172	+
		173	+如果算子已有正式 Python 脚本，优先复用，只把 lifecycle 和核心 op 调用切到 `torch.ops.catccos.*`。
		174	+
		175	+## 检查清单
		176	+
		177	+- wrapper 声明、定义、调用三处签名一致。
		178	+- `catccos_torch` 能链接对应 `lib<op>_kernel.so`。
		179	+- `TORCH_LIBRARY(catccos, m)` schema 与 Python 调用一致。
		180	+- backend impl 注册到 `torch::kPrivateUse1`。
		181	+- Meta impl 注册到 `TORCH_LIBRARY_IMPL(catccos, Meta, m)`。
		182	+- Meta 输出 shape 保留符号维度，使用 `sym_sizes()` / `empty_symint(...)`。
		183	+- runtime 使用 `at_npu::native::OpCommand` + `SetCustomHandler` 调 wrapper。
		184	+- Python 验证入口不再出现旧 custom class 调用。
		185	+- 正式脚本统一放在 `examples/<op>/scripts/` 下。
		186	+
		187	+## 禁用旧路线
		188	+
		189	+下面这些只属于旧 custom class 接入方式，新增或清理算子时不要再引入：
		190	+
		191	+- `torch::jit::CustomClassHolder`
		192	+- `REGISTER_CATCCOS_OPS_CLASS`
		193	+- `torch_register.h`
		194	+- `torch.classes.CatccosOps.*`
		195	+- 通过 custom class 构造函数申请 symmetric workspace、析构函数释放 workspace 的生命周期设计