triton.language.extra.cann.extension.scope
- class triton.language.extra.cann.extension.scope(core_mode: str, _builder=None, _semantic=None, **kwargs)
Context manager for entering and exiting a scope, where operations within a scope shares some common characteristics.
import triton.language.extra.cann.extension as extension
@triton.jit def kernel(x_ptr, y_ptr, N):
# specify annotation with extension.scope(feature_a=True):
a = tl.load(x_ptr) b = tl.load(y_ptr) result = tl.dot(a, b)
- Reserved keywords:
core_mode: Allows explicitly specify which core type should be used for operations within a code block, helping the compiler generate appropriate code for cube or vector cores.
Example
@triton.jit def kernel_scope_vector(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) with al.scope(core_mode="vector"): x = tl.load(x_ptr + i, mask=i < n) y = tl.load(y_ptr + i, mask=i < n) result = x + y tl.store(out_ptr + i, result, mask=i < n)
Special Restrictions
core_mode: must be 'vector', 'cube', 'SIMT', or 'SIMD'.
Each kernel supports one cube scope and one vector scope; they execute in parallel.
Explicit synchronization (sync_block_set/sync_block_wait) required for cross-scope data dependencies.
- __init__(core_mode: str, _builder=None, _semantic=None, **kwargs)
- 参数:
core_mode -- Either "cube" or "vector" to specify the core type
_builder -- Internal builder object (set by code_generator)
_semantic -- Internal semantic object (set by code_generator)
kwargs -- Additional internal parameters
Methods
__init__(core_mode[, _builder, _semantic])