CANN / pto-isa：基础 GEMM 算子示例

基础 GEMM 算子示例【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa概览本示例展示如何使用 PTO 实现一个基础 GEMM 算子并通过torch_npu将其暴露为 PyTorch 算子。支持的 AI 处理器A2/A3目录结构demos/baseline/gemm_basic/ ├── op_extension/ # Python 包入口模块加载 ├── csrc/ │ ├── kernel/ # PTO kernel 实现 │ └── host/ # Host 侧 PyTorch 算子注册 ├── test/ # 最小化 Python 测试 ├── CMakeLists.txt # 构建配置 ├── setup.py # Wheel 构建脚本 └── README.md # 本文档算子说明计算功能本示例实现固定维度[m, k, n] [512, 2048, 1536]的 GEMM$$ C A \times B $$其中A形状[512, 2048]m × kB形状[2048, 1536]k × nC形状[512, 1536]m × n规格项目值OpTypegemm输入a:m×k,float16,ND;b:k×n,float16,DN输出c:m×n,float,NDKernel 名称gemm_basic_customTiling 参数验证平台有 24 个核。本示例采用4 × 6分组方式将工作量切分到核上优先切m与n将m分为 4 份n分为 6 份以充分利用 24 核。每核形状singleCoreM 128、singleCoreK 2048、singleCoreN 256由于每核 tile 仍然超过 L0 容量进一步将k按 64 的 base block 切分。base block 为baseM 128、baseK 64、baseN 256参数值m512k2048n1536singleCoreM128singleCoreK2048singleCoreN256baseM128baseK64baseN256实现说明类型定义实现中定义了 GM、L1、L0 的矩阵表示并为 tiles 绑定底层存储。示例简化using NDValidShapeA TileShape2DU, baseM, baseK; using NDsingleCoreShapeA BaseShape2DU, M, K; using GlobalDataSrcA GlobalTensorU, NDValidShapeA, NDsingleCoreShapeA; // A in GM (ND) using NDValidShapeB TileShape2DU, baseK, baseN, Layout::DN; using NDsingleCoreShapeB BaseShape2DU, K, N, Layout::DN; using GlobalDataSrcB GlobalTensorU, NDValidShapeB, NDsingleCoreShapeB, Layout::DN; // B in GM (DN) using NDValidShapeC TileShape2DT, baseM, baseN; using NDWholeShapeC BaseShape2DT, M, N; using GlobalDataOut GlobalTensorT, NDValidShapeC, NDWholeShapeC; // C in GM流水线调度本示例通过在 L1 与 L0 使用双缓冲来重叠数据搬运与计算以提高利用率。同步点用于确保依赖关系正确包括正向同步MTE2 - MTE1、MTE1 - MMAD、MMAD - FIXPIPE反向同步MTE1 - MTE2、MMAD - MTE1流水线示意图构建与运行1. 准备Python环境创建python虚拟环境并安装需要的python包python -m venv virEnv source virEnv/bin/activate python3 -m pip install -r requirements.txt2. 构建 wheel配置 Ascend CANN 环境、 PTO Tile Lib 路径并构建 wheelexport ASCEND_HOME_PATH/usr/local/Ascend/ source ${ASCEND_INSTALL_PATH}/bin/setenv.bash export PTO_LIB_PATH[YOUR_PATH]/pto-isa rm -rf build op_extension.egg-info python3 setup.py bdist_wheel3. 安装 wheelpip install dist/*.whl4. 运行测试cd test python3 test.py【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考