tilelang.jit.adapter.base

The profiler and convert to torch utils

類別

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