Design overview
1. Goals & Scope
PyAsc2 is a tile-based programming model for writing Ascend NPU kernels in Python. It sits at a higher abstraction level than the original PyAsc (asc) model, which mapped Python APIs 1:1 to Ascend C intrinsics.
Core goals (in priority order):
Let developers express kernels in terms of tensors (ND-arrays in global memory) and tiles (fixed-shape chunks in on-chip memory). Buffer addresses,
TPipe/TQuelifecycles, and synchronization barriers are not exposed to the user.Reach ≈90% of the performance of hand-optimized Ascend C operators for representative workloads.
Provide a NumPy-like operation set: arithmetic, reductions, shape manipulation, masking, atomics.
[Main engineering challenge] Automate synchronization insertion, UB memory allocation, and ping-pong optimization through compiler passes and compiler hints.
Support Ascend A2 (910B), A3 (910C — two 910B dies in one package), and A5 (950) hardware families.
Relationship to PyAsc1 (asc):
PyAsc2 is implemented on top of the existing PyAsc infrastructure. The asc2.jit decorator re-uses JITFunction, the same AST-to-IR codegen, the same MLIR pass manager, and the same Bisheng compilation step. What is new is the asctile MLIR dialect and a dedicated lowering pipeline that converts tile-level IR down to the ascendc dialect before the existing backend takes over.
Out of scope:
Direct access to Ascend C intrinsics (use
ascfor that).Dynamic shapes at IR level — tile shapes must be statically known at JIT time; tensor shapes may be dynamic (runtime values).
2. Key Challenges for Ascend NPU
Ascend A2 AI Core has two types of compute cores:
AIC (Cube) — matrix / tensor operations.
AIV (Vector) — vector operations; 1 AIC : 2 AIV ratio on A2.
Each core has its own MTE (DMA), Scalar Unit, and Compute Unit. AIC and AIV have no direct interconnect — data exchange goes through L2 / Global Memory.
2.1 Synchronization Insertion
MTE and Compute Unit within each core run in parallel. Explicit set_flag / wait_flag barriers must be inserted between them. Missing barriers cause silent data hazards; redundant barriers cause stalls. Must be fully automated by the compiler.
2.2 Ping-Pong (Double Buffering)
MTE loads for tile N+1 must overlap with Compute on tile N. Requires the compiler to split the loop body, hoist loads, and insert correct barriers — without user annotation.
2.3 UB Memory Allocation and Reuse
On-chip UB (Unified Buffer) is ~256 KB per Da Vinci core (A2). After loop unrolling, many tile SSA values can be live simultaneously. The compiler must compute liveness and reuse freed UB regions. Tile shapes are statically known at JIT time; total live footprint must be validated at compile time.
2.4 Hardware Differences: A2/A3 vs A5
A3 (910C) is two A2 (910B) dies in one package — same Da Vinci architecture, same AIC / AIV model, same challenges as A2. This section therefore focuses on A5.
A5 is not an incremental refresh of A2. It introduces new hardware capabilities (register-addressable SIMD, kernel-side printf) and two new Ascend C programming levels (covered in §9 and the companion Programming Models Insights). Findings below are derived from inspection of the CANN 9 SDK preview. [58]
2.4.1 Hardware Capability Deltas
A5 adds the following hardware capabilities relative to A2: [58]
Register file is first-class. The vector register file is now exposed as a planned address space for tile computation.
Low-precision matmul first-class. Bit-mode matmul is now exposed as a dedicated hardware path, consistent with the new MX-family data formats (below).
New data formats. MXFP4, MXFP8, HiF8 (in addition to FP16 / BF16 / INT8); the type-conversion matrix grew roughly 3× to cover the new dtypes, including saturating casts.
AIC↔AIV communication moves on-chip. On A2/A3 there is no dedicated hardware channel between Cube and Vector cores — they exchange data through Global Memory (DMA in and out). On A5, a dedicated tag-based dual-channel FIFO lives in the consumer’s on-chip SRAM to enable zero-copy. The PyPTO framework’s TPUSH / TPOP protocol is the software layer built on top of this hardware. [4]
Cross-core address resolution. On-chip ring-buffer placement means the producer must know the consumer’s SRAM base. Resolved via per-function constants plus an allocator-reserved region in the consumer’s SRAM.
Pipeline sync unchanged at the hardware level. Handshake between memory engines and vector cores remains explicit — no new hardware barriers.
2.4.2 Implications for the Three DSL Challenges
Sync insertion. Three levels now coexist with different sync models: basic_api keeps explicit
set_flag/wait_flagand TPipe events; MicroAPI addsMaskRegpredication at the register-tile granularity (masks are not barriers — barriers remain basic_api’s responsibility); SIMT-API adds warp-level primitives for fine-grained sync. A DSL must decide which model to expose (or hide) and stay coherent across levels.Ping-pong. The on-chip TPUSH / TPOP ring buffer on A5 removes the GM round-trip for Cube ↔ Vector handoff — the cost model of a pipelined stage shifts substantially. On MicroAPI, pipelining is a concern at register-tile granularity, not UB-tile granularity, because loops iterate over
RegTensorchunks.UB memory allocation. Two new address-spaces to plan: the register file (first-class in MicroAPI) and the reserved consumer SRAM segment for the TPUSH / TPOP ring buffer (on A5, a fixed exclusion zone inside UB or L1). The DSL allocator must model both.
Portability. Targeting only basic_api is the conservative choice but forfeits A5’s register-file throughput. Targeting MicroAPI reaches the register file but requires c310 (A5-only builds). A DSL that claims to target A2 + A5 must lower to basic_api only, or implement per-target lowering.
3. Architecture Overview
┌──────────────────────────────────────────────────────────┐
│ User Python Code │
│ @asc2.jit │
└──────────────────────┬───────────────────────────────────┘
│
┌──────────────────────▼───────────────────────────────────┐
│ PyAsc2 Frontend (python/asc2/) │
│ │
│ GlobalTensor (GM) LocalTensor (UB/L0/…) asc2.range │
│ copy_in / copy_out arithmetic reductions shape ops │
│ atomics matmul unary creation indexing │
└──────────────────────┬───────────────────────────────────┘
│ FunctionVisitor (AST → IR)
┌──────────────────────▼───────────────────────────────────┐
│ asctile MLIR Dialect │
│ GlobalTensorType / LocalTensorType LoadOp / StoreOp │
│ BinaryOps UnaryOps ReductionOps ShapeOps │
│ AtomicRMWOp MatmulOp SoftmaxOp SelectOp │
│ CountMaskOp BitwiseMaskOp │
└──────────────────────┬───────────────────────────────────┘
│ AscTile passes (lib/Dialect/AscTile/Transforms/)
│ + AscLower passes (lib/Conversion/LowerToAsc/)
┌──────────────────────▼───────────────────────────────────┐
│ ascendc MLIR Dialect (existing backend) │
│ LocalTensor / GlobalTensor TPipe / TQue │
│ set_flag / wait_flag Ascend C API ops │
└──────────────────────┬───────────────────────────────────┘
│ lib/Target/AscendC/ (CodeEmitter)
┌──────────────────────▼───────────────────────────────────┐
│ Ascend C source (.cce) │
└──────────────────────┬───────────────────────────────────┘
│ Bisheng compiler
▼
NPU kernel binary (.o)
End-to-end compilation flow
@asc2.jit kernel invocation
→ argument type specialization (same as asc1)
→ two-level cache lookup (same as asc1)
→ [miss]
AST walk → asctile MLIR (FunctionVisitor + asctile builder APIs)
↓
AscTile passes (unrolling, loop transforms, math specialization)
↓
AscLower passes (lower asctile → ascendc dialect)
↓
ascendc passes (UB allocation, sync insertion, boilerplate gen)
↓
CodeEmitter (ascendc → Ascend C source)
↓
Bisheng (Ascend C → .o binary)
↓
cache store
→ kernel launch via runtime rt library calls (same as asc1)
4. Programming Model
4.1 The two types: GlobalTensor and LocalTensor
|
|
|
|---|---|---|
Memory |
Global Memory (HBM, GM) |
On-chip: UB, L0A, L0B, L0C, L1 |
Shape |
ND, may be dynamic (RuntimeInt dims) |
ND, must be static (known at JIT time) |
Creation |
|
Returned by |
Mutability |
Passed to/from kernel; never computed in-kernel |
Value semantics; produced by ops, consumed once |
IR type |
|
|
GlobalTensor is a descriptor — it holds a pointer and a shape. LocalTensor carries ValueSemantics in MLIR, meaning instances are value types — self-contained and safe to CSE, copy, and move. Each op produces a new tile SSA value. The compiler is free to map multiple logical tiles to the same physical UB region.
4.2 Memory hierarchy (TensorLocation)
Enum value |
Hardware unit |
Typical use |
|---|---|---|
|
Unified Buffer |
Vector operations (default) |
|
L0 matrix input buffers |
Matmul inputs |
|
L0 matrix output buffer |
Matmul accumulator |
|
L1 cache |
Intermediate staging |
copy_in defaults to TensorLocation.UB. Matmul implicitly creates L0A/L0B/L0C tiles; the LowerToL0 pass handles the layout conversion.
4.3 Load and store
# Load a local tensor from a global tensor:
tile = asc2.copy_in(tensor, [base], [128]) # explicit element offsets
# Load a scalar (no shape → scalar load)
scalar = asc2.copy_in(tensor, [i])
# Optional padding for partial tiles
tile = asc2.copy_in(tensor, [base], [128], pad_value=0.0)
# Copy a local tensor (or sub-region) into a fresh local tensor, optionally on another location
clone = asc2.copy(tile)
sub = asc2.copy(tile, [0], [64])
# Store local tensor or scalar back
asc2.copy_out(tile, tensor, [base])
asc2.copy_out(scalar_value, tensor, [i])
The last dimension of every local tensor shape must be aligned to ub_block_size (32 bytes). This is enforced at JIT time via check_data_alignment.
4.4 Operations
All operations produce new LocalTensor or scalar values; no mutation.
Arithmetic / comparison (tile ⊕ tile, tile ⊕ scalar, scalar ⊕ tile):
add, sub, mul, div, maximum, minimum, left_shift, right_shift,
equal, not_equal, greater, greater_equal, less, less_equal
Operator overloads on LocalTensor (+, -, *, /, >, ==, …) call the same functions. Dtypes are promoted automatically via infer_common_dtype.
Unary:
abs, ceil, floor, negative, relu,
sqrt, rsqrt, exp, exp2, log, log2,
sin, cos, tan, sinh, cosh, tanh, erf, softmax, rms_norm
Reductions (one or more axes, optional keep_dims):
reduce_sum, reduce_max, reduce_min, reduce_prod
— with no axes: returns a scalar PlainValue.
— also bound as methods: tile.sum(), tile.max(), etc.
Shape manipulation:
reshape, ravel, broadcast_to, expand_dims, squeeze
Creation:
full(shape, value), zeros(shape), full_like, zeros_like,
zeros_acc(shape, dtype) — zeros tile in L0C (matmul accumulator),
concat(*tiles) — concatenate along the first dimension.
Matrix multiply:
matmul(a, b)— 2-D float tiles; result is alwaysfloat32.matmul_acc(acc, a, b)— fused matmul that accumulates into an existingacctile (typically created withzeros_acc).
Conditional / masking:
where(mask, src0, src1)— element-wise select.mask(count=…)/mask(bits=(hi, lo), other=…)— context manager that wraps the enclosed operations in a conditional region (CountMaskOp/BitwiseMaskOp).
Atomics:
atomic_add, atomic_max, atomic_min — write back to a global GlobalTensor.
4.5 Programming model operations
i = asc2.block_idx() # current NPU block index (PlainValue)
n = asc2.block_num() # total number of blocks (PlainValue)
4.6 asc2.range
asc2.range extends the base range with two attributes that control compiler behaviour:
for i in asc2.range(start, stop, step, unroll_factor=4, parallel=False):
...
unroll_factor— how many iterations to unroll. Tagged on the MLIRforop and processed byTagUnrollGroups+UnrollLooppasses.parallel=True— marks the loop body soParallelLoadStorepass can emit stores and loads in parallel.
5. IR Design
5.1 asctile dialect
Defined in include/ascir/Dialect/AscTile/IR/ using TableGen.
Types:
Type |
MLIR syntax |
Python proxy |
|---|---|---|
|
ND global descriptor |
|
|
Fixed-shape local chunk |
|
LocalTensor carries ValueSemantics, meaning instances are value types — self-contained and safe to CSE, copy, and move. Both LocalTensor and GlobalTensor implement ShapedTypeInterface.
Key operations (from Ops.td):
Category |
Op |
Description |
|---|---|---|
Tensor creation |
|
Create a |
|
Read a dynamic dimension from a |
|
|
Fill a tile with a scalar constant |
|
|
Create a zeros tile in L0C (matmul accumulator) |
|
Memory transfer |
|
Load tile from |
|
Store |
|
|
Store an L0C tile to a |
|
|
Copy a tile (or sub-region) into a new tile |
|
|
Copy an L0C tile into another tile via fixpipe |
|
|
Scalar load from |
|
|
Scalar store to |
|
Element-wise compute |
|
Type conversion between tile dtypes |
|
Fused ReLU |
|
|
Element-wise conditional select ( |
|
|
Element-wise compare (tile/tile and tile/scalar) |
|
|
Vector-scalar binary ops |
|
Reduction / norm |
|
Reduce tile along one or more axes |
|
Reduce a tile to a 1-element tile or a scalar |
|
|
Row-wise softmax on a 1D / 2D tile |
|
|
RMS normalisation |
|
Matmul |
|
2-D matrix multiply |
|
2-D matrix multiply accumulating into an L0C tile |
|
Shape |
|
Reshape a tile |
|
Broadcast tile to a wider shape |
|
|
Concatenate tiles along the first dimension |
|
Atomic |
|
Atomic read-modify-write (add / max / min) |
Region / control |
|
Region: conditional on element count |
|
Region: conditional on bitmask |
|
|
Region terminator |
Standard dialect ops (arith, scf, math) are emitted directly by the AST walker for scalar arithmetic and control flow; the AscLower passes later handle them.
5.2 IR generation
FunctionVisitor (Python AST walker, python/asc/codegen/function_visitor.py) calls pybind11-exposed builder methods from python/src/OpBuilder.cpp to create asctile.* ops. Every asc2.* Python API call produces exactly one (or a small fixed number of) asctile operations. No analysis is needed at this stage.
5.3 Lowering chain
asctile dialect
↓ AscTile passes (loop unrolling, math specialization, masking)
↓ AscLower passes (lower each asctile op → ascendc ops + scf + arith)
ascendc dialect
↓ ascendc passes (UB allocation, sync insertion, boilerplate)
Ascend C source
6. Optimization Passes
The full pass pipeline is scheduled by Compiler._schedule_passes() inside python/asc/runtime/compiler.py. The run_asc2_passes=True flag (set automatically by asc2.jit) activates the AscTile and AscLower phases.
6.1 AscTile passes (lib/Dialect/AscTile/Transforms/)
These operate on the asctile dialect before lowering begins.
Pass |
Purpose |
|---|---|
|
Scan |
|
Hoist pure (side-effect-free) ops — e.g., shape computations — out of loop bodies |
|
Cluster load/store ops within an unroll group so they appear adjacent, enabling more efficient pipeline scheduling |
|
Split loads destined for L0A/L0B into the legal Cube data path |
|
Rewrite |
|
Specialise generic |
|
Convert eligible |
|
Replace scalar-tail reductions with vector forms |
|
Physically unroll loops tagged by |
|
Standard MLIR cleanup between stages |
6.2 AscLower passes (lib/Conversion/LowerToAsc/)
These lower asctile.* operations into ascendc.* + standard MLIR ops.
Pass |
Purpose |
|---|---|
|
Replace |
|
Widen boolean (i1) tensors to |
|
Lower |
|
Lower tile-scalar and scalar-tile arithmetic |
|
Lower |
|
Lower |
|
Main lowering: |
|
Handle boolean tile lowering special cases |
|
Lower remaining |
|
Lower |
|
Replace |
|
Materialise any pending unrealized casts from the lowering |
|
Lower |
6.3 ascendc passes (after lowering, lib/Dialect/Asc/Transforms/)
Once all asctile ops are gone, the existing ascendc pipeline takes over:
Phase |
Key passes |
|---|---|
Lowering glue |
|
Memory allocation |
|
Synchronization |
|
Optimization |
|
Postprocessing |
|
ComputeMemoryConsumption is only active in asc2 mode and records per-TPosition UB usage as a module attribute; it raises a compile-time error on overflow.
7. Memory & Resource Model
7.1 LocalTensor-level view
From the user’s perspective, memory management is invisible. The user:
Creates
GlobalTensordescriptors pointing to HBM buffers.Calls
copy_in(tensor, shape, offsets=…)to bring data into aLocalTensor.Applies operations to produce new tiles.
Calls
copy_out(tile, tensor, offsets=…)to write results back.
The compiler is responsible for mapping each LocalTensor SSA value to a physical UB region.
7.2 Compiler-side UB allocation
Two strategies, selected per compilation:
Strategy |
Option |
How it works |
|---|---|---|
TPipe-managed |
|
|
Static allocation |
|
|
static_alloc is Optional[bool]; when left at None, Compiler.__init__ resolves it to True on 910_95 and False otherwise.
UB pressure reduction (both modes):
HoistUBAllocation— move allocations above loops so one allocation covers all iterations.ReuseUBAllocation(reuse_alloc=1) — when a tile’s lifetime ends before a new one is needed, reuse its UB region.ComputeMemoryConsumption— sums all live tile sizes perTPositionand fails the compilation if UB limits are exceeded.
7.3 Synchronization
Because asc2 users don’t write set_flag/wait_flag, all synchronization is inserted by the compiler:
insert_sync=Trueis set automatically byasc2.jit.On 910B / 910_93:
InsertSyncuses the classic queue-position–based algorithm.On 910_95:
InsertBufIdSyncuses a newer algorithm that reduces sync overhead by tracking buffer IDs rather than queue positions;FuseBufIdSyncmerges adjacent sync ops.verify_sync=Trueruns an additionalVerifySyncpass after sync insertion as a sanity check.
asc2.range(parallel=True) annotates a loop with an MLIR attribute that is intended to relax sync insertion for parallel-safe load/store groups. Today the attribute is only set; the consuming pass is not yet wired up.
8. Usage Examples
Vector addition
import asc2
@asc2.jit
def vadd(x_ptr, y_ptr, out_ptr, size: int, tiles_per_block: int, TILE: asc.ConstExpr[int]):
x_gm = asc2.global_tensor(x_ptr, [size])
y_gm = asc2.global_tensor(y_ptr, [size])
out_gm = asc2.global_tensor(out_ptr, [size])
base = asc2.block_idx() * tiles_per_block * TILE
for i in asc2.range(tiles_per_block):
off = base + i * TILE
x = asc2.copy_in(x_gm, [off], [TILE])
y = asc2.copy_in(y_gm, [off], [TILE])
asc2.copy_out(x + y, out_gm, [off])
vadd[8](x, y, out, n, asc2.ceildiv(n // 256, 8), TILE=256)
Softmax (row-wise)
@asc2.jit
def softmax(x_ptr, out_ptr, rows: int, cols: int, TILE: asc.ConstExpr[int]):
x_gm = asc2.global_tensor(x_ptr, [rows, cols])
out_gm = asc2.global_tensor(out_ptr, [rows, cols])
for row in asc2.range(asc2.block_idx(), rows, asc2.block_num()):
x = asc2.copy_in(x_gm, [row, 0], [1, TILE])
exp_x = asc2.exp(x - x.max())
asc2.copy_out(exp_x / exp_x.sum(), out_gm, [row, 0])
9. Insights from Existing Programming Models
A comparative analysis of existing tile-based and kernel DSLs relevant to Ascend NPU (AscendC, Triton, cuTile, TileLang-Ascend, Triton-Ascend, PyPTO, Pallas, Mojo) is maintained in a companion document:
The observations from that document inform the design decisions in §10.
10. Key Design Decisions
TODO: to be filled in.
11. References
[4] PTO ISA — TPUSH/TPOP protocol. https://github.com/hw-native-sys/pypto/blob/main/docs/en/reference/pto-isa/01-tpush_tpop.md
[58] CANN 9 SDK preview — 950 / A5 programming model findings (internal notes from SDK source inspection).
Appendix: Key CompileOptions for PyAsc2
Option |
Default in |
Effect |
|---|---|---|
|
|
Enable AscTile + AscLower pipeline |
|
|
Auto-insert sync barriers |
|
|
Static vs TPipe-managed UB allocation |
|
|
Reuse freed UB regions |
|
|
Densify load/store groups (experimental) |
|
|
Fuse consecutive vector ops into Ascend C MicroAPI VF blocks |
|
|
Run |
|
|
Emit cube-only kernels (drives |
|
|
Pass |
|
|
Bisheng |
|
|
Bypass cache, recompile every call |