tilelang.jit.adapter.base¶

The profiler and convert to torch utils

Attributes¶

CachedTextSourceLike

Classes¶

CachedTextSource

BaseKernelAdapter

Helper class that provides a standard way to create an ABC using

Functions¶

output_allocation_stream(stream, device)

Return the caller-supplied launch stream when outputs must be allocated on it.

Module Contents¶

class tilelang.jit.adapter.base.CachedTextSource¶
text: str | None = None¶
path: str | None = None¶
tilelang.jit.adapter.base.CachedTextSourceLike¶
tilelang.jit.adapter.base.output_allocation_stream(stream, device)¶

Return the caller-supplied launch stream when outputs must be allocated on it.

torch’s caching allocator orders reuse of a block only after earlier work on the stream the block was allocated on, so tensors written by a kernel launched on stream are allocated on stream. None means the launch stream is the current stream, where a plain torch.empty is already correct.

Parameters:

  • stream (int)

  • device (torch.device)

Return type:

torch.cuda.Stream | None

class tilelang.jit.adapter.base.BaseKernelAdapter(mod, params, result_idx)¶

Bases: abc.ABC

Helper class that provides a standard way to create an ABC using inheritance.

Parameters:
func: collections.abc.Callable | None = None¶
mod¶
params¶
result_idx¶
static get_current_stream_functor(device=None)¶

Return a callable that reads Torch’s current device stream pointer.

The returned lambda yields the raw stream handle of the current PyTorch stream on the active device. It’s a thunk (evaluated at call time) so that any upstream stream guards are respected. If no supported device runtime is available, it returns a lambda that yields 0. An explicit Torch device takes precedence over runtime auto-detection.

Parameters:

device (torch.device | None)

Return type:

collections.abc.Callable[[], int]

static get_current_device_functor()¶

Return a callable that yields Torch’s current device.

Similar to the stream functor, we capture a callable that, when called, fetches the current device according to PyTorch, independently of the compilation target. If neither NPU nor CUDA is available, returns torch.device('cpu').

Return type:

collections.abc.Callable[[], torch.device]

__call__(*args, **kwds)¶
Parameters:

  • args (Any)

  • kwds (Any)

Return type:

Any

get_kernel_source(kernel_only=True)¶
Parameters:

kernel_only (bool)

Return type:

str