tilelang.jit.abi¶

Shared helpers for preparing TVM-FFI callable ABIs.

Attributes¶

OUT_IDX_ATTR

CALLEE_ALLOCATED_OUTPUTS_ATTR

Functions¶

prepare_tvm_ffi_callee_allocated_outputs(func, out_idx, *)

Resolve output indices and expose them to TVM-FFI lowering.

Module Contents¶

tilelang.jit.abi.OUT_IDX_ATTR = 'tilelang_out_idx'¶
tilelang.jit.abi.CALLEE_ALLOCATED_OUTPUTS_ATTR = 'tilelang_callee_allocated_outputs'¶
tilelang.jit.abi.prepare_tvm_ffi_callee_allocated_outputs(func, out_idx, *, supports_callee_allocated_outputs=False)¶

Resolve output indices and expose them to TVM-FFI lowering.

When the resolved execution backend declares supports_callee_allocated_outputs, the ABI decision is recorded on the derived PrimFunc so lowering and the runtime adapter both honor it. Otherwise tilelang_out_idx stays purely declarative and the kernel keeps the caller-preallocated output path.

Parameters:

  • func (tvm.tirx.PrimFunc)

  • out_idx (list[int] | int | None)

  • supports_callee_allocated_outputs (bool)

Return type:

tuple[tvm.tirx.PrimFunc, list[int] | None]