tilelang.autotuner.capture module#
- class tilelang.autotuner.capture.CaptureStack#
Bases:
objectA simple stack implementation for capturing items in a thread-local context. Used to manage a stack of items (e.g., input tensors) for auto-tuning capture.
- pop()#
Pop and return the top item from the stack.
- Returns:
The item at the top of the stack.
- Raises:
IndexError – If the stack is empty.
- push(item)#
Push an item onto the top of the stack.
- Parameters:
item – The item to be pushed onto the stack.
- size()#
Return the number of items in the stack.
- Returns:
The size of the stack.
- Return type:
int
- top()#
Return the item at the top of the stack without removing it.
- Returns:
The item at the top of the stack.
- Raises:
IndexError – If the stack is empty.
- tilelang.autotuner.capture.get_autotune_inputs() Optional[List[Any]]#
Get the current autotune inputs from the stack.
- tilelang.autotuner.capture.set_autotune_inputs(*args) AutotuneInputsCapture#
Set input tensors for auto-tuning.
This function creates a context manager for capturing input tensors during the auto-tuning process. It supports both:
set_autotune_inputs(a, b, c) set_autotune_inputs([a, b, c])
- Parameters:
*args – Either a single list/tuple of tensors, or multiple tensor arguments.
- Returns:
A context manager for auto-tuning inputs.
- Return type: