tilelang.jit.adapter.base

The profiler and convert to torch utils

Classes

BaseKernelAdapter

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

Module Contents

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.

参数:
func: Callable | None = None
mod
params
result_idx
static get_current_stream_functor()

Return a callable that reads Torch's current CUDA stream pointer.

The returned lambda yields the raw CUDA 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 CUDA is unavailable, it returns a lambda that yields 0.

返回类型:

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. On CPU or when CUDA is unavailable, returns torch.device('cpu').

返回类型:

Callable[[], torch.device]

__call__(*args, **kwds)
参数:

  • args (Any)

  • kwds (Any)

返回类型:

Any

get_kernel_source(kernel_only=True)
参数:

kernel_only (bool)

返回类型:

str