English | Chinese
PyPTO (pronounced: pai p-t-o) is a high-performance programming framework for AI accelerators, designed to simplify the development of complex fused operators and entire model networks while maintaining high-performance computing capabilities. The framework adopts an innovative PTO (Parallel Tensor/Tile Operation) programming paradigm, with a Tile-based programming model as its core design philosophy. Through a multi-level intermediate representation (IR) system, it compiles AI model applications built via APIs from high-level Tensor graphs step by step into hardware instructions, ultimately generating executable code that runs efficiently on target platforms.
- Tile-based Programming Model: All computations are based on Tiles (hardware-aware data blocks), fully leveraging hardware parallel computing capabilities and memory hierarchy
- Multi-level Computation Graph Transformation: Transforms Tensor Graphs into Tile Graphs, Block Graphs, and Execution Graphs through compilation passes, with each step including a series of pass optimization workflows
- Automated Code Generation: Compilation results generate low-level PTO virtual instruction code through CodeGen, which is then compiled into executable code for target platforms
- MPMD Execution Scheduling: Executable code is loaded to the device side and scheduled to processor cores on the device using MPMD (Multiple Program Multiple Data) approach
- Complete Toolchain Support: Full compilation artifacts and runtime performance data can be visualized through IDE-integrated toolchains to identify performance bottlenecks; developers can also control compilation and scheduling behavior through the toolchain
- Python-friendly API: Provides intuitive Tensor-level abstractions that align with algorithm developers' thinking patterns, supporting dynamic shapes and symbolic programming
- Layered Abstraction Design: Exposes different abstraction levels to different developers - algorithm developers use the Tensor level, performance experts use the Tile level, and system developers use the Block level
- Algorithm Developers: Primarily use Tensor-level programming for rapid algorithm implementation and validation, focusing on algorithm logic
- Performance Optimization Experts: Can use Tile or Block levels for deep performance tuning to achieve optimal performance
- System Developers: Can integrate with third-party frameworks or develop toolchains at Tensor/Tile/Block and PTO virtual instruction set levels
- Python: Version 3.10 or higher
- CMake: Version 3.15 or higher
- C++ Compiler: Supporting C++17 standard (GCC, Clang, or MSVC)
- nanobind: Version 2.0.0 or higher (automatically installed during build)
- scikit-build-core: Version 0.10.0 or higher (automatically installed during build)
-
Clone the repository:
git clone https://github.com/hw-native-sys/pypto.git cd pypto -
Install in development mode (recommended for development):
pip install -e .Or with development dependencies:
# Install CPU-only torch first (avoids pulling ~2GB CUDA dependencies) pip install torch --index-url https://download.pytorch.org/whl/cpu pip install -e ".[dev]"
-
Install in production mode:
pip install .
The build system uses scikit-build-core to automatically handle CMake configuration and C++ extension compilation.
-
Build type: The default build type is
RelWithDebInfo(optimized with debug symbols). You can override this:CMAKE_BUILD_TYPE=Release pip install . -
Enable ccache (optional, for faster rebuilds):
# ccache will be automatically detected and used if available brew install ccache # macOS sudo apt-get install ccache # Ubuntu/Debian
PyPTO includes examples organized by complexity:
python examples/hello_world.pypython examples/kernels/06_softmax.pypython examples/models/01_ffn.pyTo run unit tests:
# Install with development dependencies (CPU-only torch first)
pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install -e ".[dev]"
# Run unit tests in parallel
python -m pytest tests/ut -n auto --maxprocesses 8 -v
# Run a specific unit test file
python -m pytest tests/ut/core/test_error.py -n auto --maxprocesses 8 -vFor system tests, see tests/st/README.md.
This project is licensed under the CANN Open Software License Agreement Version 2.0.
This license grants you a limited, worldwide, royalty-free license to download, use, modify, integrate, and distribute the Software or its derivative works for the purpose of developing software solely for use in systems with Huawei AI Processors (including Ascend, Kirin, Yueying, and other Huawei-branded AI chipsets).
See the LICENSE file for the full license text.