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.

Example: ```python

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])