Basic GEMM Operator Example

【免费下载链接】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

Overview

This example demonstrates how to implement a basic GEMM operator using PTO and expose it as a PyTorch operator viatorch_npu.

Supported AI Processors

• A2/A3

Directory Layout

demos/baseline/gemm_basic/ ├── op_extension/ # Python package entry (module loader) ├── csrc/ │ ├── kernel/ # PTO kernel implementation │ └── host/ # Host-side PyTorch operator registration ├── test/ # Minimal Python test ├── CMakeLists.txt # Build configuration ├── setup.py # Wheel build script └── README.md # This document

Operator Description

Function

This example implements GEMM with fixed dimensions[m, k, n] = [512, 2048, 1536]:

$$ C = A \times B $$

Where:

• Ashape:[512, 2048](m × k)
• Bshape:[2048, 1536](k × n)
• Cshape:[512, 1536](m × n)

Specification

Tiling Parameters

The validation platform has 24 cores. The workload is split across cores (prioritizing splittingmandn) using a4 × 6grouping: splitminto 4 parts andninto 6 parts to fully utilize 24 cores.

Per-core shape:

• singleCoreM = 128,singleCoreK = 2048,singleCoreN = 256

Because the per-core tile still exceeds L0 capacity,kis further tiled into base blocks of size 64. The base block is:

• baseM = 128,baseK = 64,baseN = 256

Implementation Notes

Type definitions

The implementation defines matrix representations for GM, L1, and L0, then assigns backing storage for tiles. Example (simplified):

using NDValidShapeA = TileShape2D<U, baseM, baseK>; using NDsingleCoreShapeA = BaseShape2D<U, M, K>; using GlobalDataSrcA = GlobalTensor<U, NDValidShapeA, NDsingleCoreShapeA>; // A in GM (ND) using NDValidShapeB = TileShape2D<U, baseK, baseN, Layout::DN>; using NDsingleCoreShapeB = BaseShape2D<U, K, N, Layout::DN>; using GlobalDataSrcB = GlobalTensor<U, NDValidShapeB, NDsingleCoreShapeB, Layout::DN>; // B in GM (DN) using NDValidShapeC = TileShape2D<T, baseM, baseN>; using NDWholeShapeC = BaseShape2D<T, M, N>; using GlobalDataOut = GlobalTensor<T, NDValidShapeC, NDWholeShapeC>; // C in GM

Pipeline scheduling

This example overlaps data movement and compute using double buffering in L1 and L0 to improve utilization. Synchronization points ensure correct dependencies, including:

• Forward sync:MTE2 -> MTE1,MTE1 -> MMAD,MMAD -> FIXPIPE
• Reverse sync:MTE1 -> MTE2,MMAD -> MTE1

Pipeline overview:

Build and Run

1. Prepare the python environment

Create your own virtual environment and install the required python package.

python -m venv virEnv source virEnv/bin/activate python3 -m pip install -r requirements.txt

2. Configure environment and build the wheel

export 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_wheel

3. Install the wheel

pip install dist/*.whl

4. Run the example

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