Runtime API

JIT decorator

@asc2.jit(fn: Callable[[P], T]) → JITFunction[P, T]

@asc2.jit(**options) → Callable[[Callable[[P], T]], JITFunction[P, T]]

Instantiate a JIT function using the default options. This should be used as a decorator for the kernel function:

@asc2.jit
def kernel(x, y):
    ...

JIT options may be provided as keyword arguments to be applied to the decorated kernel function. See asc.CodegenOptions, asc.CompileOptions, asc.LaunchOptions for the details.

JIT options

Attributes from the following classes can be used as keyword arguments in asc2.jit decorator.

Tip

The user can mix options from different sections. For example, to disable CodegenOptions.capture_exceptions and set CompileOptions.opt_level to 2, the kernel function may be decorated as the following:

@asc2.jit(capture_exceptions=False, opt_level=2)
def kernel(x, y):
    ...

class asc.CodegenOptions(capture_exceptions: bool = True, ir_multithreading: bool = True)

Code generation and AST traversal options

capture_exceptions: bool = True: Capture all exceptions raised by functions called by code generator internally and show the corresponding location in the user function code being traversed that caused an exception instead. Usually, it must always be enabled, but may be disabled for the language debugging purposes.

ir_multithreading: bool = True: Enable parallel processing for the current IR context if possible. Usually, it must always be enabled, but may be disabled for the IR debugging purposes.

class asc.CompileOptions(debug: bool = False, strip_loc: bool = False, verify_sync: bool = False, print_ir_before_all: bool = False, run_passes: bool = True, kernel_type: KernelType | None = None, auto_sync: bool | None = True, auto_sync_log: str | None = '', matmul_cube_only: bool = False, always_compile: bool = False, bisheng_options: Tuple[str] | None = None, opt_level: int = 3, run_asc2_passes: bool = False, densify_load_store: bool = False, insert_sync: bool | None = None, reuse_ub: bool = False, reuse_ub_in_out: bool = False, static_alloc: bool = False)

Binary compilation and IR transformation options

always_compile: bool = False: Always run full compilation pipeline instead of fetching a cached objects from the prevoius runs. This option may be useful for testing and parallel execution.

bisheng_options: Tuple[str] | None = None: Append extra arguments to the bisheng command line used to produce an object file for the kernel. Please, run bisheng --help to get the list of options supported by the compiler.

opt_level: int = 3: Optimization level for the Bisheng compiler. Supported values are 1, 2, 3. Typically, this parameter affects the -O argument of the command line for the compiler.

run_asc2_passes: bool = False: Enable PyAsc2 compilation pipeline. This option is enabled automatically when @asc2.jit decorator is used.

densify_load_store: bool = False: Densify asc2.load and asc2.store statements by grouping them together.

Warning

This is an experimental feature. It might or might not cause functional or performance regressions.

insert_sync: bool | None = None: Insert synchronization instructions automatically. This feature is enabled by default, which is usually a must, but may be disabled for the debugging purposes.

reuse_ub: bool = False: Try to reduce the UB memory usage by replacing the tiles with those allocated earlier but became unused. Having this feature enabled may help to avoid UB overflow but may introduce performance regressions.

reuse_ub_in_out: bool = False: Allow for input/output tiles to be reused alongside intermediate ones. This option only has effect if reuse_ub feature is enabled and densify_load_store feature is not.

Warning

This is an experimental feature. It might or might not cause functional or performance regressions.

static_alloc: bool = False: Perform static allocation for tiles instead of relying on Ascend C TPipe backend. The static allocation feature may help to reduce an overhead caused by scalar code.

class asc.LaunchOptions(core_num: int = 0, stream: c_void_p | None = None)

Kernel launch and device runtime options.

These options can also be used as positional arguments in [ brackets ] when launch JIT function:

@asc2.jit
def kernel(x_ptr, y_ptr):
    ...

def launch(x, y, core_num = 16):
    kernel[core_num](x, y)

core_num: int = 0: Number of active execution blocks (AI cores) used to launch the kernel. By default, all cores available on the current platform will be used.