triton-operator-code-review

Static code review of Triton operators (Ascend NPU)

Review principles

• Ascend-specific constraints first: Focus on hardware differences, do not check general Triton knowledge
• Zero tolerance for Mask: Ascend has zero tolerance for out-of-bounds access

Severity classification

Level	Meaning	Typical issue
P0	Certain crash or wrong result	Missing Mask, core type mismatch, Atomic loop deadlock
P1	High-probability precision/functional issue	Reduction without precision promotion, Softmax without subtracting max
P2	Performance/maintainability	Redundant memory access, BLOCK not aligned

Reference resource loading

Stage	Load	Do not load
Phase 1: Host side	ascend-triton-api-constraints.md	dtype-matrix
Phase 2: Device side	ascend-api-dtype-matrix.md	test-patterns
Check item by item	code-review-checklist.md	—
Reference official implementation	ascend-test-patterns.md	—
Confirm API signature/parameters	triton-api-reference.md	—

Load trigger: When it is necessary to confirm the signature, parameters, or available data types of a certain tl/libdevice API, read in full triton-api-reference.md。

Review workflow

Phase 1: Host side

Check item	Identification method	Level
Hardcoded core count	grid = (20,) Literals such as	P0
Core type mismatch	Contains tl.dot kernel uses num_vectorcore	P0
Matrix multiplication degraded to element-wise	No tl.dot, implement matmul/GEMV using Vector Core element-wise multiply-add	P0
BLOCK_SIZE not tl.constexpr	Declaration check	P1
Matrix operation BLOCK not a multiple of 16	Numeric check	P2

Core type: Contains tl.dot → AI Core (num_aicore); element-wise/reduction/activation → Vector Core (num_vectorcore). Note: matrix multiplication operators (including GEMV) must use tl.dot + AI Core; Cube Core throughput is far higher than Vector Core.

Phase 2: Device side

Mask completeness (P0): All tl.load/tl.store must have mask= or use make_block_ptr。

Data types (P0-P1)：tl.dot input only supports int8/fp16/fp32/bf16;dot_scaled conditionally supported (bf16/fp16 lhs/rhs, int8 scale, fp32 out);permute/trans does not support int64.

Precision handling (P1): Before FP16/BF16 reduction .to(tl.float32); Softmax must subtract max.

Single-pass check for reduction operators (P1): For operators such as GroupNorm/LayerNorm/RMSNorm, check whether the same data is loaded multiple times (first compute statistics, then load for normalization). When D ≤ UB, it should be single-pass: after one load, tl.sum(x, 1) completes all computation within UB. Double-pass can be 5x slower.

Control flow (P0): Triton IR uses MLIR structured control flow (scf.for/scf.while), does not support early exit.

• ❌ for/while Within the loop return — Compilation error: "Cannot have return statements inside while or for" (including return in sub-functions)
• ❌ for Inside the loop break — Compilation error: "unsupported AST node type: Break"

Tensor indexing (P0): Triton tensors do not support Python-style [] subscript operations.

• ❌ tensor[i] = val — AssertionError
• ❌ val = tensor[i] — AssertionError
• ❌ tensor[i:j] slicing — Compilation error

Code patterns：

• ❌ for ... : tl.atomic_cas/or/xor/and/xchg(...) — possible deadlock (P0)
• ❌ Using tl.atomic_add return values in multi-core kernels (P0)
• ❌ Inside the kernel, import numpy（P0）
• ⚠️ tensor[i].item() on the host hot path — triggers CPU-NPU synchronization (P2)
• ⚠️ Reduction-type operators have multiple for ... : tl.load(x_ptr + ...) loops — double-pass anti-pattern (P1)
• ⚠️ Post-dot operations do not use tl.parallel(bind_sub_block=True) — tl.dot is followed by element-wise operations but has no parallel sharding (P2)
• ⚠️ Integer type conversion overflow — tl.cast to int8/int16 without specifying overflow_mode（P1）

Phase 3: Performance pitfalls (P2)

• Multiple tl.load for the same ptr → redundant GM access
• tl.arange(0, N) * stride(stride > 1) → non-contiguous memory access
• pid directly maps to a block with no loop → uneven load
• Inside the kernel, use broadcast stride to compute offsets for auxiliary tensors (cos/sin, etc.) (b * stride0 + n * stride1 + ...) → recommend switching to host-side expand+contiguous, unify row * D + col offsets
• Reduction-type operators with weight/bias not pre-expanded, using weight_ptr + ch_idx gather inside the kernel → recommend expanding to a contiguous layout on the host side

Anti-pattern list (NEVER)

Host

• ❌ Hardcoded core count / using num_vectorcore for matrix multiplication / BLOCK_SIZE not tl.constexpr
• ❌ Matrix multiplication (including GEMV) implemented with Vector Core element-wise multiply-add — must use tl.dot + AI Core

Device

• ❌ tl.load/tl.store without mask / tl.dot input int32/int16/int64 / dot_scaled
• ❌ atomic_or/xor/and/xchg/cas inside a for loop / third-party libraries inside the kernel
• ❌ FP16/BF16 reduction without promoting to FP32 / Softmax without subtracting max
• ❌ Integer truncation cast without specifying overflow_mode
• ❌ for/while Inside the loop return / break — Triton structured control flow does not support early exit (including return in sub-functions)
• ❌ tensor[i] index operations (read/assign/slice) — use tl.where / tl.gather / tl.extract_slice instead

Output

By code-review-report-template.md format output the report.