Performance tips
Key Performance Principles
1. More Cores = More Performance
Ascend NPUs expose dozens of AI Cores (up to 56+ on C310). Distribute work across as many cores as possible. Under-utilizing cores is the single biggest performance mistake.
# Good: spread rows across all available cores
rows_per_block = asc2.ceildiv(input_num_rows, asc2.block_num())
block_offset = asc2.block_idx() * rows_per_block
# Bad: only use a few cores, leaving most idle
if asc2.block_idx() < 4:
...
Launch with the maximum core count the problem allows:
kernel[56](...) # use all 56 cores when workload is large enough
Each core processes its own slice of data in parallel — doubling active cores roughly doubles throughput.
2. Bigger Tiles = Better Performance, Less Transfer Overhead
Every asc2.copy_in / asc2.copy_out is a DMA transfer between global memory (HBM) and on-chip UB. Larger tiles amortize transfer latency and reduce the total number of transfers.
# Good: large tile, fewer transfers
tile_length = 10496 # ~40 KB for fp32 — fills UB well
for i in asc2.range(loop_count, ...):
xt = asc2.copy_in(in_gm, [tile_length], ...)
# Bad: small tile, many transfers, high overhead
tile_length = 128 # only 512 bytes — mostly waiting on DMA
for i in asc2.range(huge_loop_count, ...):
xt = asc2.copy_in(in_gm, [tile_length], ...)
Rule of thumb: make tiles as large as UB memory allows (typically ~256 KB total, shared among all live tiles). The test suite uses tiles of 10496–18752 elements for fp32 kernels (40–75 KB per tile).
3. Use ConstExpr for Static Code Generation
All tiling parameters, shapes, and loop bounds must be typed as asc2.ConstExpr. This tells the JIT compiler to treat them as compile-time constants, enabling:
Full loop unrolling (
unroll_factorbaked into generated code)Static UB allocation (
static_alloc=True) — no runtime memory management overheadDead-code elimination and constant folding in the MLIR pipeline
# Good: all tiling params are ConstExpr → static, optimized code
@asc2.jit(static_alloc=True, reuse_alloc=1)
def kernel(input_ptr: asc2.GlobalAddress, output_ptr: asc2.GlobalAddress,
input_length: asc2.ConstExpr, tile_length: asc2.ConstExpr,
unroll_factor: asc2.ConstExpr):
...
# Bad: plain int → dynamic code, no unrolling, runtime overhead
def kernel(input_ptr, output_ptr, input_length: int, tile_length: int):
...
ConstExpr values are fixed at JIT compile time. Changing them triggers recompilation, but the resulting binary is significantly faster.
JIT Options and Loop Control
JIT Decorator Options
Option |
Purpose |
Recommended Usage |
|---|---|---|
|
Static UB allocation at compile time |
Enabling recommended for most kernels. Leads to faster execution, no runtime memory management. Requires all local tensor shapes to be |
|
Reuse freed UB regions across iterations |
Setting to |
|
Enable vector fusion (experimental) |
Advanced option, when enabled lowering generated code to register level API. Experimental feature. May improve performance for element-wise chains. |
|
Cached kernels usage |
When enabled ignores cached kernels, provides recompilation. |
Loop Control: unroll_factor and parallel
These parameters control how loops are compiled and executed:
unroll_factor=N (passed to asc2.range() or range()):
Unrolls the loop N times at compile time
Reduces loop overhead and enables instruction-level parallelism
Recommended:
unroll_factor=2for most kernels,unroll_factor=1for very large tiles or memory-bound operationsMust be
ConstExprfor static unrolling
parallel=True (passed to asc2.range() or range()):
Enables parallel load/store optimization across loop iterations
Allows overlapping DMA transfers with computation
Recommended: Enable for outer loops that perform independent tile operations
Works best with
unroll_factor >= 2
# Recommended pattern: unroll + parallel for outer tile loop
for i in asc2.range(loop_count, unroll_factor=2, parallel=True):
xt = asc2.copy_in(in_gm, [tile_length], ...)
zt = xt + yt
asc2.copy_out(zt, out_gm, ...)
# For nested loops: parallel on outer, sequential on inner
for i in asc2.range(row_iters, unroll_factor=2, parallel=True):
for j in asc2.range(col_iters, parallel=False):
# inner loop operations
Vector function (VF) block fusion
For a sequence of basic elementwise or reduction vector operations, PyAsc can automatically fuse them into a single VF (vector function) block that executes at the register level, avoiding redundant UB reads/writes between intermediate steps. Enable this by passing vf_fusion=True to the JIT decorator:
@asc2.jit(vf_fusion=True)
def kernel(x_ptr, y_ptr, out_ptr, size: int, tile_size: asc2.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])
for i in asc2.range(asc2.ceildiv(size, tile_size)):
x = asc2.copy_in(x_gm, [i * tile_size], [tile_size])
y = asc2.copy_in(y_gm, [i * tile_size], [tile_size])
# These elementwise ops are fused into a single VF block
result = (x + y) * x - y
asc2.copy_out(result, out_gm, [i * tile_size])
The automatic fusion handles straightforward chains of built-in arithmetic and reduction operations. For more complex patterns — such as custom register-level logic, specialized masking, or operations not expressible through the standard API — use asc2.inline_vf() to embed raw Ascend C register code directly as a VF block:
# x * y + z — three inputs, processed in 64-element vector chunks
out = asc2.inline_vf(
"""
auto* out_ptr = reinterpret_cast<__ubuf__ float*>($0.GetPhyAddr());
auto* x_ptr = reinterpret_cast<__ubuf__ float*>($1.GetPhyAddr());
auto* y_ptr = reinterpret_cast<__ubuf__ float*>($2.GetPhyAddr());
auto* z_ptr = reinterpret_cast<__ubuf__ float*>($3.GetPhyAddr());
AscendC::MicroAPI::RegTensor<float> x_reg, y_reg, z_reg, xy_reg, result_reg;
uint32_t count = 256;
for (uint16_t i = 0; i < 4; i += 1) {
uint32_t offset = i * 64;
AscendC::MicroAPI::MaskReg mask = AscendC::MicroAPI::UpdateMask<float>(count);
AscendC::MicroAPI::DataCopy(x_reg, x_ptr + offset);
AscendC::MicroAPI::DataCopy(y_reg, y_ptr + offset);
AscendC::MicroAPI::Mul(xy_reg, x_reg, y_reg, mask);
AscendC::MicroAPI::DataCopy(z_reg, z_ptr + offset);
AscendC::MicroAPI::Add(result_reg, xy_reg, z_reg, mask);
AscendC::MicroAPI::DataCopy(out_ptr + offset, result_reg, mask);
}
""",
shape=x.shape, dtype=x.dtype, inputs=[x, y, z])
Inside the code string, $0 refers to the output tensor and $1, $2, … refer to input tensors in the order they appear in the inputs list. All input tensors must reside in UB memory.
Measuring Performance with the Profiler
The test suite provides a built-in profiler fixture that wraps kernel launches with NPU hardware profiling.
Enabling Profiling
Profiling only activates on real NPU hardware with the --profile flag:
# Profile on NPU backend (requires physical Ascend NPU)
pytest --backend NPU --profile python/test/asc2/target/test_vadd.py
# Profile with multiple runs per test (more stable median)
pytest --backend NPU --profile --runs 10 python/test/asc2/target/test_vadd.py
# Select specific platform
pytest --backend NPU --profile --platform Ascend950PR_9599 python/test/asc2/target/
Without --profile or when not using the NPU backend, a StubProfiler is used (no-op).
Using the Profiler in Tests
Wrap the kernel launch loop with profiler.profile():
def test_my_kernel(profiler, runs, block_num, ...):
in_tensor = torch.randn(...)
out_tensor = torch.zeros(...)
with profiler.profile():
for _ in range(runs):
my_kernel[block_num](in_tensor, out_tensor, ...)
expected = torch_ref(in_tensor)
torch.testing.assert_close(out_tensor, expected, atol=1e-3, rtol=1e-3)
The runs fixture (controlled by --runs, default=1) determines how many kernel launches are timed. Multiple runs produce a more stable median duration.
Output
After the test session, profiling results are printed in a summary table:
============================== Profiling results ==============================
python/test/asc2/target/test_vadd.py::test_add[...]: 12.34 μs
python/test/asc2/target/test_softmax.py::test_softmax[...]: 45.67 μs
The reported duration is the median task time across all runs (first run is skipped as warmup via skip=1 in task_time_median).
Profiling Tips
Use
--runs 10or more for stable measurements — single runs have high variance.Vary core number and tile sizes: run the same test with different
tile_lengthvalues to see transfer overhead impact. Varyblock_numto find the sweet spot (more cores = more parallelism, but decrease if work per core is too small).Use
always_compile=Falsefor fast iterations: by default, compiled kernels are cached. Setalways_compile=Trueonly when actively debugging compilation issues. Use--compile-onlyto verify tiling decisions compile without UB overflow before running on hardware.