tilelang.contrib.pto_toolchain¶

Utilities to invoke the PTO toolchain (ptodsl -> ptoas -> bisheng) for Ascend NPU targets.

compile_ptodsl_to_aibin lowers one TileLang-generated PTODSL module into an executable CCE ELF (aibin): the PTODSL frontend script (ptodsl_compile.py, run as a subprocess) lowers the source to MLIR, ptoas assembles a fat object, bisheng links it, and the .aicore_binary section of the linked image is extracted as the device binary consumed by the Ascend runtime module.

Attributes¶

logger

BISHENG_FATOBJ_LINK_FLAGS

Functions¶

resolve_pto_options()

Resolve the ptoas arch and extra flags from the environment.

parse_pto_kernel_names(ptodsl_source)

Return the top-level @pto.jit kernel names in definition order.

compile_ptodsl_to_pto(ptodsl_source, kernel_names, ...)

Lower one PTODSL module to MLIR text.

run_ptoas(pto_path, fatobj_path, pto_arch, pto_flags)

Assemble the MLIR text into a device fat object.

bisheng_fatobj_link(fatobj_path, out_so[, verbose])

Link the fat object alone into a shared image carrying the device ELF.

extract_aicore_binary(so_path)

Extract the .aicore_binary section bytes from a linked image.

toolchain_stamp()

Identity of the tools that produce a PTO aibin, for cache keys.

compile_ptodsl_to_aibin(ptodsl_source, pto_arch, pto_flags)

Compile one PTODSL module into an executable CCE ELF (aibin).

Module Contents¶

tilelang.contrib.pto_toolchain.logger¶
tilelang.contrib.pto_toolchain.resolve_pto_options()¶

Resolve the ptoas arch and extra flags from the environment.

Return type:

tuple[str, list[str]]

tilelang.contrib.pto_toolchain.parse_pto_kernel_names(ptodsl_source)¶

Return the top-level @pto.jit kernel names in definition order.

Parameters:

ptodsl_source (str)

Return type:

list[str]

tilelang.contrib.pto_toolchain.compile_ptodsl_to_pto(ptodsl_source, kernel_names, src_path, out_path, verbose=False)¶

Lower one PTODSL module to MLIR text.

ptodsl_compile.py runs in a subprocess so the DSL’s tracing globals (pto.init_core and friends) never enter this interpreter.

Parameters:

  • ptodsl_source (str)

  • kernel_names (list[str])

  • src_path (str | pathlib.Path)

  • out_path (str | pathlib.Path)

  • verbose (bool)

Return type:

None

tilelang.contrib.pto_toolchain.run_ptoas(pto_path, fatobj_path, pto_arch, pto_flags, verbose=False)¶

Assemble the MLIR text into a device fat object.

Parameters:

  • pto_path (str | pathlib.Path)

  • fatobj_path (str | pathlib.Path)

  • pto_arch (str)

  • pto_flags (list[str])

  • verbose (bool)

Return type:

None

Link the fat object alone into a shared image carrying the device ELF.

Parameters:

  • fatobj_path (str | pathlib.Path)

  • out_so (str | pathlib.Path)

  • verbose (bool)

Return type:

None

tilelang.contrib.pto_toolchain.extract_aicore_binary(so_path)¶

Extract the .aicore_binary section bytes from a linked image.

Parses the ELF64 section table directly instead of shelling out to readelf: tokenizing its text output can silently truncate the extract, and a bad extract is only rejected later by aclrtBinaryLoadFromData.

Parameters:

so_path (str | pathlib.Path)

Return type:

bytearray

tilelang.contrib.pto_toolchain.toolchain_stamp()¶

Identity of the tools that produce a PTO aibin, for cache keys.

ptoas compiles the device object with the toolchain it resolves from ASCEND_HOME_PATH, so the CANN identity (which includes the bisheng and the CCE linker inside that install, and the version ptoas gates its ABI on) is recorded.

Return type:

dict[str, str]

tilelang.contrib.pto_toolchain.compile_ptodsl_to_aibin(ptodsl_source, pto_arch, pto_flags, verbose=False)¶

Compile one PTODSL module into an executable CCE ELF (aibin).

Parameters:

  • ptodsl_source (str)

  • pto_arch (str)

  • pto_flags (list[str])

  • verbose (bool)

Return type:

bytearray