Triton

https://triton-lang.org/main/python-api/
Admin
Triton is a language and compiler for parallel programming that provides a Python-based environment...
Tokens:51,610
Snippets:482
Trust Score:8.1
Update:1 week ago
Show doc for...
Tokens:
Raw
### Example: Using make_tensor_descriptor in a Triton kernel

Source: https://triton-lang.org/main/python-api/generated/triton.language.make_tensor_descriptor.html

This example defines a Triton kernel that uses tl.make_tensor_descriptor to create a descriptor for a 2D tensor, then loads and stores blocks using the descriptor. It also shows the required global memory allocation setup with triton.set_allocator and the kernel launch.

```python
@triton.jit
def inplace_abs(in_out_ptr, M, N, M_BLOCK: tl.constexpr, N_BLOCK: tl.constexpr):
    desc = tl.make_tensor_descriptor(
        in_out_ptr,
        shape=[M, N],
        strides=[N, 1],
        block_shape=[M_BLOCK, N_BLOCK],
    )

    moffset = tl.program_id(0) * M_BLOCK
    noffset = tl.program_id(1) * N_BLOCK

    value = desc.load([moffset, noffset])
    desc.store([moffset, noffset], tl.abs(value))

# TMA descriptors require a global memory allocation
def alloc_fn(size: int, alignment: int, stream: Optional[int]):
    return torch.empty(size, device="cuda", dtype=torch.int8)

triton.set_allocator(alloc_fn)

M, N = 256, 256
x = torch.randn(M, N, device="cuda")
M_BLOCK, N_BLOCK = 32, 32
grid = (M / M_BLOCK, N / N_BLOCK)
inplace_abs[grid](x, M, N, M_BLOCK, N_BLOCK)
```

--------------------------------

### Correct variable scoping in Triton JIT

Source: https://triton-lang.org/main/python-api/triton-semantics.html

Example showing how to define a variable before a block to ensure it is accessible on all paths.

```python
@triton.jit
def kernel(out_ptr):
    value = 0.0
    for _ in range(0, 1):
        value = 1.0
    tl.store(out_ptr, value)
```

--------------------------------

### Execute PTX assembly with inline_asm_elementwise

Source: https://triton-lang.org/main/python-api/generated/triton.language.inline_asm_elementwise.html

This example demonstrates using PTX assembly to perform element-wise operations on two input tensors, processing four elements at a time.

```python
@triton.jit
def kernel(A, B, C, D, BLOCK: tl.constexpr):
    a = tl.load(A + tl.arange(0, BLOCK)) # uint8 tensor
    b = tl.load(B + tl.arange(0, BLOCK)) # float32 tensor

    # For each (a,b) in zip(a,b), perform the following:
    # - Let ai be `a` converted to int32.
    # - Let af be `a` converted to float.
    # - Let m be the max of ai and b.
    # - Return ai and mi.
    # Do the above 4 elements at a time.
    (c, d) = tl.inline_asm_elementwise(
        asm="""
        {
            // Unpack `a` into `ai`.
            .reg .b8 tmp<4>;
            mov.b32 {tmp0, tmp1, tmp2, tmp3}, $8;
            cvt.u32.u8 $0, tmp0;
            cvt.u32.u8 $1, tmp1;
            cvt.u32.u8 $2, tmp2;
            cvt.u32.u8 $3, tmp3;
        }
        // Convert `ai` to float.
        cvt.rn.f32.s32 $4, $0;
        cvt.rn.f32.s32 $5, $1;
        cvt.rn.f32.s32 $6, $2;
        cvt.rn.f32.s32 $7, $3;
        // Take max of `ai` and `b`.
        max.f32 $4, $4, $9;
        max.f32 $5, $5, $10;
        max.f32 $6, $6, $11;
        max.f32 $7, $7, $12;
        """,
        constraints=(
            # 8 output registers, namely
            #   $0=ai0, $1=ai1, $2=ai2, $3=ai3,
            #   $4=m0,  $5=m1,  $6=m2,  $7=m3.
            "=r,=r,=r,=r,=r,=r,=r,=r,"
            # 5 input registers, namely
            #   $8=ai,
            #   $9=b0, $10=b1, $11=b2, $12=b3.
            # The four elements from `a` are all packed into one register.
            "r,r,r,r,r"),
        args=[a, b],
        dtype=(tl.int32, tl.float32),
        is_pure=True,
        pack=4,
    )
    tl.store(C + tl.arange(0, BLOCK), c)
    tl.store(D + tl.arange(0, BLOCK), d)
```

--------------------------------

### Using static_assert for compile-time validation

Source: https://triton-lang.org/main/python-api/generated/triton.language.static_assert.html

This example demonstrates how to use tl.static_assert to verify that a block size constant meets a specific requirement at compile time.

```python
tl.static_assert(BLOCK_SIZE == 1024)
```

--------------------------------

### static_print usage example

Source: https://triton-lang.org/main/python-api/generated/triton.language.static_print.html

Use triton.language.static_print to print a value at compile time. The example prints the value of a variable named x using an f-string.

```python
static_print(x)
```

--------------------------------

### triton.language.tensor_descriptor.load

Source: https://triton-lang.org/main/python-api/generated/triton.language.tensor_descriptor.html

Load a block from the descriptor starting at the given element offsets.

```APIDOC
## `load`(self, offsets, _semantic)

### Description
Loads a block of data from the tensor descriptor starting at the given element offsets.

### Parameters
- **offsets** (tensor) - The element offsets to start loading from.
- **_semantic** (optional) - The memory semantic for the operation.

### Returns
The loaded block of data.
```

--------------------------------

### Mapping a scalar function with map_elementwise

Source: https://triton-lang.org/main/python-api/generated/triton.language.map_elementwise.html

This example demonstrates using map_elementwise to apply a scalar function containing control flow to a tensor.

```python
@triton.jit
def selu_scalar(x, alpha):
    if x > 0:
        return a
    else:
        return alpha * (tl.exp(x) - 1)

@triton.jit
def selu(x, alpha):
    return tl.map_elementwise(selu_scalar, x, alpha)
```

--------------------------------

### Using tl.device_print and print

Source: https://triton-lang.org/main/python-api/generated/triton.language.device_print.html

Shows how to call tl.device_print with a prefix and a value, and that the Python builtin print is equivalent.

```python
tl.device_print("pid", pid)
print("pid", pid)

```

--------------------------------

### Get top 4 elements from a 1D tensor using tl.topk

Source: https://triton-lang.org/main/python-api/generated/triton.language.topk.html

Example usage of triton.language.topk to retrieve the k largest elements from a 1D tensor. The example creates a tensor with values 0 to 15 and returns the top 4 elements, which are [15, 14, 13, 12].

```python
# Get top 4 elements from a 1D tensor
x = tl.arange(0, 16)
top4 = tl.topk(x, 4)  # Returns [15, 14, 13, 12]

```

--------------------------------

### Using triton.autotune with configs and key

Source: https://triton-lang.org/main/python-api/generated/triton.autotune.html

Decorates a triton.jit kernel with triton.autotune, providing a list of triton.Config objects and a key list to trigger re-evaluation when the value of x_size changes. Note that all configurations are evaluated, so the kernel runs multiple times; use reset_to_zero to avoid undesired side effects.

```python
@triton.autotune(configs=[
    triton.Config(kwargs={'BLOCK_SIZE': 128}, num_warps=4),
    triton.Config(kwargs={'BLOCK_SIZE': 1024}, num_warps=8),
  ],
  key=['x_size'] # the two above configs will be evaluated anytime
                 # the value of x_size changes
)
@triton.jit
def kernel(x_ptr, x_size, BLOCK_SIZE: tl.constexpr):
    ...

```

--------------------------------

### triton.language.arange(start, end, _semantic=None)

Source: https://triton-lang.org/main/python-api/generated/triton.language.arange.html

Returns contiguous values within the half-open interval [start, end). The difference end - start must be less than or equal to TRITON_MAX_TENSOR_NUMEL = 1048576. Both start and end must be powers of two, with end greater than start.

```APIDOC
## triton.language.arange(start, end, _semantic=None)

### Description
Returns contiguous values within the half-open interval `[start, end)`. The difference `end - start` must be less than or equal to `TRITON_MAX_TENSOR_NUMEL = 1048576`.

### Parameters
- **start** (int32) - Required - Start of the interval. Must be a power of two.
- **end** (int32) - Required - End of the interval. Must be a power of two greater than `start`.

### Constraints
- `end - start` must be less than or equal to `TRITON_MAX_TENSOR_NUMEL = 1048576`.
- `start` must be a power of two.
- `end` must be a power of two greater than `start`.
```

--------------------------------

### Benchmark.__init__()

Source: https://triton-lang.org/main/python-api/generated/triton.testing.Benchmark.html

Constructor for the Benchmark class, used to configure benchmark parameters for generating line plots.

```APIDOC
## Benchmark.__init__()

### Description
Constructor for the Benchmark class. Initializes a benchmark configuration for generating line plots, specifying x-axis arguments, line arguments, and fixed arguments.

### Method
Constructor

### Signature
`__init__(self, x_names: List[str], x_vals: List[Any], line_arg: str, line_vals: List[Any], line_names: List[str], plot_name: str, args: Dict[str, Any], xlabel: str = '', ylabel: str = '', x_log: bool = False, y_log: bool = False, styles=None)`

### Parameters
- **x_names** (List[str]) - Required - Name of the arguments that should appear on the x axis of the plot.
- **x_vals** (List[Any]) - Required - List of values to use for the arguments in `x_names`. Can be a list of scalars or a list of tuples/lists. If a list of scalars and multiple x_names, all arguments will have the same value. If a list of tuples/lists, each element should have the same length as x_names.
- **line_arg** (str) - Required - Argument name for which different values correspond to different lines in the plot.
- **line_vals** (List[Any]) - Required - List of values to use for the arguments in `line_arg`.
- **line_names** (List[str]) - Required - Label names for the different lines.
- **plot_name** (str) - Required - Name of the plot.
- **args** (Dict[str, Any]) - Required - Dictionary of keyword arguments to remain fixed throughout the benchmark.
- **xlabel** (str, optional) - Label for the x axis of the plot.
- **ylabel** (str, optional) - Label for the y axis of the plot.
- **x_log** (bool, optional) - Whether the x axis should be log scale.
- **y_log** (bool, optional) - Whether the y axis should be log scale.
- **styles** (list[tuple[str, str]]) - A list of tuples, where each tuple contains two elements: a color and a linestyle.

### Usage
This class is used by the `perf_report` function to generate line plots with a concise API.
```

--------------------------------

### triton.language.tensor_descriptor.store

Source: https://triton-lang.org/main/python-api/generated/triton.language.tensor_descriptor.html

Store a block from the descriptor starting at the given element offsets.

```APIDOC
## `store`(self, offsets, value, _semantic)

### Description
Stores a block of data to the tensor descriptor starting at the given element offsets.

### Parameters
- **offsets** (tensor) - The element offsets to start storing at.
- **value** (tensor) - The value to store.
- **_semantic** (optional) - The memory semantic for the operation.

### Returns
None
```

--------------------------------

### triton.jit(*, version=None, repr=None, launch_metadata=None, do_not_specialize=None, do_not_specialize_on_alignment=None, debug=None, noinline=None)

Source: https://triton-lang.org/main/python-api/generated/triton.jit.html

Decorator for JIT-compiling a function using the Triton compiler, with keyword-only parameters for advanced configuration. When a jit'd function is called, arguments are implicitly converted to pointers if they have a .data_ptr() method and a .dtype attribute. The function will be compiled and run on the GPU, with access only to Python primitives, builtins within the triton package, arguments, and other jit'd functions.

```APIDOC
## triton.jit(*, version=None, repr=None, launch_metadata=None, do_not_specialize=None, do_not_specialize_on_alignment=None, debug=None, noinline=None)

### Description
Decorator for JIT-compiling a function using the Triton compiler, with keyword-only parameters for advanced configuration. When a jit'd function is called, arguments are implicitly converted to pointers if they have a .data_ptr() method and a .dtype attribute. The function will be compiled and run on the GPU, with access only to Python primitives, builtins within the triton package, arguments, and other jit'd functions.

### Signature
triton.jit(_*_ , _version =None_, _repr : Callable | None = None_, _launch_metadata : Callable | None = None_, _do_not_specialize : Iterable[int | str] | None = None_, _do_not_specialize_on_alignment : Iterable[int | str] | None = None_, _debug : bool | None = None_, _noinline : bool | None = None_) -> Callable[[T], JITFunction[T]]

### Parameters
- **fn** (Callable) - Required - The function to be jit-compiled.
- **version** (optional) - Optional - Version parameter.
- **repr** (Callable | None) - Optional - Representation callable.
- **launch_metadata** (Callable | None) - Optional - Launch metadata callable.
- **do_not_specialize** (Iterable[int | str] | None) - Optional - Indices or names of arguments to not specialize.
- **do_not_specialize_on_alignment** (Iterable[int | str] | None) - Optional - Indices or names of arguments to not specialize on alignment.
- **debug** (bool | None) - Optional - Debug flag.
- **noinline** (bool | None) - Optional - Noinline flag.

### Notes
- Arguments with a .data_ptr() method and a .dtype attribute are implicitly converted to pointers.
- The function will be compiled and run on the GPU.
- The function will only have access to: python primitives, builtins within the triton package, arguments to this function, and other jit'd functions.
```

--------------------------------

### triton.jit

Source: https://triton-lang.org/main/python-api/triton.html

Decorator for JIT-compiling a function using the Triton compiler.

```APIDOC
## triton.jit

### Description
Decorator for JIT-compiling a function using the Triton compiler.

### Usage
```python
@triton.jit
def kernel(...):
    ...
```
```

--------------------------------

### print() equivalent

Source: https://triton-lang.org/main/python-api/generated/triton.language.device_print.html

The Python builtin `print` is equivalent to `triton.language.device_print` in this context.

```APIDOC
## Python builtin print

### Description
Calling the Python builtin `print` is the same as calling `triton.language.device_print`, and the requirements for the arguments will match this function (not the normal requirements for `print`).
```

--------------------------------

### arange

Source: https://triton-lang.org/main/python-api/triton.language.html

Returns contiguous values within the half-open interval [start, end).

```APIDOC
## arange

### Description
Returns contiguous values within the half-open interval [start, end).

### Parameters
- **start** (int) - Required - The start of the interval.
- **end** (int) - Required - The end of the interval (exclusive).

### Returns
- A tensor containing contiguous values from start to end-1.
```

--------------------------------

### histogram

Source: https://triton-lang.org/main/python-api/triton.language.html

Computes a histogram based on input tensor with num_bins bins, the bins have a width of 1 and start at 0.

```APIDOC
## histogram

### Description
Computes a histogram based on input tensor with num_bins bins, the bins have a width of 1 and start at 0.

### Parameters
- **input** (tensor) - The input tensor.
- **num_bins** (int) - The number of bins.

### Returns
- (tensor) - A tensor containing the histogram.
```

--------------------------------

### Using static_range in a JIT kernel

Source: https://triton-lang.org/main/python-api/generated/triton.language.static_range.html

Demonstrates the basic syntax for using static_range within a JIT-compiled Triton kernel.

```python
@triton.jit
def kernel(...):
    for i in tl.static_range(10):
        ...
```

--------------------------------

### triton.language.load

Source: https://triton-lang.org/main/python-api/generated/triton.language.load.html

Documentation for the triton.language.load function, which loads data from memory.

```APIDOC
## triton.language.load

### Description
Loads a tensor of data from memory at the location defined by the pointer. If the pointer is a single element pointer, a scalar is loaded; if the pointer is an N-dimensional tensor of pointers, an N-dimensional tensor is loaded. The mask and other parameters are implicitly broadcast to the pointer shape when the pointer is a tensor.

### Method
Function call

### Signature
`triton.language.load(pointer, mask=None, other=None, *, cache_modifier='', eviction_policy='', volatile=False, _semantic=None)`

### Parameters
- **pointer** (triton.PointerType, or block of dtype=triton.PointerType) - Required - Pointer to the data to be loaded.
- **mask** (Block of triton.int1, optional) - If mask[idx] is false, do not load the data at address pointer[idx].
- **other** (Block, optional) - If mask[idx] is false, return other[idx]. If other is None, the masked-out value is undefined.
- **cache_modifier** (str, optional) - Changes cache option in NVIDIA PTX. Should be one of {"", ".ca", ".cg", ".cv"}, where ".ca" stands for cache at all levels, ".cg" stands for cache at global level (cache in L2 and below, not L1), and ".cv" means don't cache and fetch again.
- **eviction_policy** (str, optional) - Changes eviction policy in NVIDIA PTX.
- **volatile** (bool, optional) - Changes volatile option in NVIDIA PTX.

### Returns
A tensor of data loaded from memory. If pointer is a single element pointer, a scalar is returned; if pointer is an N-dimensional tensor of pointers, an N-dimensional tensor is returned.

### Notes
- When pointer is a single element pointer, mask and other must also be scalars, and other is implicitly typecast to pointer.dtype.element_ty.
- When pointer is an N-dimensional tensor of pointers, mask and other are implicitly broadcast to pointer.shape, and other is implicitly typecast to pointer.dtype.element_ty.
- The `*` in the signature indicates keyword-only parameters separator (PEP 3102).
```

--------------------------------

### Variable scoping error in Triton JIT

Source: https://triton-lang.org/main/python-api/triton-semantics.html

Example demonstrating a NameError when a variable is only assigned inside a loop body and accessed outside of it.

```python
@triton.jit
def kernel(out_ptr):
    for _ in range(0, 1):
        value = 1.0
    tl.store(out_ptr, value)  # `value` is not defined here
```

--------------------------------

### triton.testing.do_bench(fn, warmup=25, rep=100, grad_to_none=None, quantiles=None, return_mode='mean')

Source: https://triton-lang.org/main/python-api/generated/triton.testing.do_bench.html

Benchmarks the runtime of the provided function. By default, returns the mean runtime in milliseconds as a single float. Supports returning quantiles or other summary statistics.

```APIDOC
## triton.testing.do_bench

### Description
Benchmark the runtime of the provided function. By default, returns the mean runtime (in milliseconds) of `fn` as a single float.

### Method
Function call

### Signature
`triton.testing.do_bench(fn, warmup=25, rep=100, grad_to_none=None, quantiles=None, return_mode='mean')`

### Parameters
- **fn** (_Callable_) - Required - Function to benchmark
- **warmup** (_int_) - Optional - Warmup time (in ms). Controls how long the function is run before timing begins. Default: 25
- **rep** (_int_) - Optional - Repetition time (in ms). Controls the total duration of the timed runs. Default: 100
- **grad_to_none** (_list[torch.Tensor]_, optional) - Optional - Reset the gradient of the provided tensors to None before each run, to avoid gradient accumulation affecting timing.
- **quantiles** (_list[float]_, optional) - Optional - If provided, return these quantiles of the runtime distribution instead of a summary statistic. For example, `[0.2, 0.5, 0.8]` returns the 20th, 50th, and 80th percentile runtimes in ms. When set, `return_mode` is ignored.
- **return_mode** (_str_) - Optional - The summary statistic to return when `quantiles` is not set. `"mean"` and `"median"` return a single float. `"min"` and `"max"` return the fastest and slowest run respectively. `"all"` returns the raw list of per-run timings in ms. Default is `"mean"`.

### Returns
The runtime(s) in milliseconds: a single float for a scalar `return_mode`, or a list of floats if `quantiles` is set or `return_mode="all"`.

### Return Type
float | list[float]
```

--------------------------------

### triton.testing.do_bench_proton(fn, warmup=25, rep=100, grad_to_none=None, quantiles=None, return_mode='mean')

Source: https://triton-lang.org/main/python-api/generated/triton.testing.do_bench_proton.html

This code snippet documents the triton.testing.do_bench_proton function, which benchmarks the runtime of kernels invoked by a provided function using the Proton profiler.

```APIDOC
## triton.testing.do_bench_proton(fn, warmup=25, rep=100, grad_to_none=None, quantiles=None, return_mode='mean')

### Description
Benchmark the runtime of kernels invoked by the provided function using the Proton profiler. The measured runtime is generally more accurate than `do_bench` for short kernels that are affected by CPU overhead. Note that this function has several constraints compared to `do_bench`: it does not measure GPU operations other than kernels (e.g., memory copies, synchronization, etc.) and it supports only AMD and NVIDIA GPUs.

### Method
Function call (Python API)

### Endpoint
triton.testing.do_bench_proton

### Parameters
- **fn** (Callable) - Required - Function to benchmark.
- **warmup** (int) - Optional - Warmup time (in ms). Default is 25.
- **rep** (int) - Optional - Repetition time (in ms). Default is 100.
- **grad_to_none** (torch.Tensor, optional) - Optional - Reset the gradient of the provided tensor(s) to None.
- **quantiles** (list[float], optional) - Optional - Performance percentiles to return in addition to the median.
- **return_mode** (str) - Optional - The statistical measure to return. Options are \"min\", \"max\", \"mean\", \"median\", or \"all\". Default is \"mean\".

### Returns
The runtime(s) in milliseconds: a single float for a scalar `return_mode`, or a list of floats if `quantiles` is set or `return_mode=\"all\"`.

### Return Type
float | list[float]
```

--------------------------------

### triton.language.histogram(input, num_bins, mask=None, _semantic=None, _generator=None)

Source: https://triton-lang.org/main/python-api/generated/triton.language.histogram.html

Computes a histogram based on an input tensor with num_bins bins, each of width 1 starting at 0. Optionally accepts a mask to exclude certain elements.

```APIDOC
## triton.language.histogram(input, num_bins, mask=None, _semantic=None, _generator=None)

### Description
Computes a histogram based on an input tensor with num_bins bins, each of width 1 starting at 0. Optionally accepts a mask to exclude certain elements.

### Parameters
- **input** (Tensor) - Required - The input tensor.
- **num_bins** (int) - Required - Number of histogram bins.
- **mask** (Block of triton.int1, optional) - Optional - If mask[idx] is false, exclude input[idx] from histogram.

### Usage
This function can also be called as a member function on `tensor`, as `x.histogram(...)` instead of `histogram(x, ...)`.
```

--------------------------------

### triton.language.join(a, b)

Source: https://triton-lang.org/main/python-api/generated/triton.language.join.html

Joins the given tensors in a new, minor dimension. For example, two tensors of shape (4,8) produce a tensor of shape (4,8,2), and two scalars produce a tensor of shape (2). The inputs are broadcasted to the same shape. To join more than two elements, call this function multiple times, reflecting Triton's constraint on power-of-two sizes. `join` is the inverse of `split`.

```APIDOC
## `triton.language.join(a, b, _semantic=None)`

### Description
Joins the given tensors into a new, minor dimension. For example, joining two tensors of shape (4,8) produces a tensor of shape (4,8,2). Joining two scalars returns a tensor of shape (2). The inputs are broadcasted to the same shape. To join more than two elements, call this function multiple times, reflecting Triton's constraint on power-of-two sizes. `join` is the inverse of `split`.

### Method
Function call

### Parameters
- **a** (Tensor) – The first input tensor.
- **b** (Tensor) – The second input tensor.

### Return Value
Returns a new tensor with an additional minor dimension.
```

--------------------------------

### store(self, value[, mask, cache_modifier, ...])

Source: https://triton-lang.org/main/python-api/generated/triton.language.tensor.html

Forwards to the store() free function.

```APIDOC
## store(self, value[, mask, cache_modifier, ...])

### Description
Forwards to the store() free function.

### Method
store

### Parameters
- **self** (tensor) - Required - The destination tensor.
- **value** (tensor) - Required - The value to store.
- **mask** (tensor) - Optional - The mask for conditional store.
- **cache_modifier** (str) - Optional - Cache modifier.
- **...** - Optional - Additional arguments.

### Returns
None
```

--------------------------------

### do_bench

Source: https://triton-lang.org/main/python-api/triton.testing.html

Benchmark the runtime of the provided function.

```APIDOC
## do_bench

### Description
Benchmark the runtime of the provided function.

### Usage
Call this function with the function to benchmark to measure its execution time.
```

--------------------------------

### triton.jit(fn)

Source: https://triton-lang.org/main/python-api/generated/triton.jit.html

Decorator for JIT-compiling a function using the Triton compiler. When called, arguments with a .data_ptr() method and a .dtype attribute are implicitly converted to pointers. The function will be compiled and run on the GPU, with access only to Python primitives, builtins within the triton package, arguments, and other jit'd functions.

```APIDOC
## triton.jit(fn)

### Description
Decorator for JIT-compiling a function using the Triton compiler. When a jit'd function is called, arguments are implicitly converted to pointers if they have a .data_ptr() method and a .dtype attribute. The function will be compiled and run on the GPU, with access only to Python primitives, builtins within the triton package, arguments, and other jit'd functions.

### Signature
triton.jit(_fn_ : T) -> JITFunction[T]

### Parameters
- **fn** (Callable) - Required - The function to be jit-compiled.

### Notes
- Arguments with a .data_ptr() method and a .dtype attribute are implicitly converted to pointers.
- The function will be compiled and run on the GPU.
- The function will only have access to: python primitives, builtins within the triton package, arguments to this function, and other jit'd functions.
```

--------------------------------

### triton.autotune

Source: https://triton-lang.org/main/python-api/triton.html

Decorator for auto-tuning a `triton.jit`'d function.

```APIDOC
## triton.autotune

### Description
Decorator for auto-tuning a `triton.jit`'d function.

### Usage
```python
@triton.autotune(configs=[...])
@triton.jit
def kernel(...):
    ...
```
```

--------------------------------

### device_assert

Source: https://triton-lang.org/main/python-api/triton.language.html

Assert the condition at runtime from the device.

```APIDOC
## device_assert

### Description
Assert the condition at runtime from the device.

```

--------------------------------

### ravel

Source: https://triton-lang.org/main/python-api/triton.language.html

Returns a contiguous flattened view of x.

```APIDOC
## ravel

### Description
Returns a contiguous flattened view of x.

### Parameters
- **x** (tensor) - Required - The input tensor.

### Returns
- A contiguous flattened view of the input tensor.
```

--------------------------------

### view(self, *shape[, _semantic])

Source: https://triton-lang.org/main/python-api/generated/triton.language.tensor.html

Forwards to the view() free function.

```APIDOC
## view(self, *shape[, _semantic])

### Description
Forwards to the view() free function.

### Method
view

### Parameters
- **self** (tensor) - Required - The input tensor.
- **shape** (int) - Required - The new shape.
- **_semantic** (optional) - Optional - Semantic information.

### Returns
The viewed tensor.
```

--------------------------------

### Broadcast to shape syntax

Source: https://triton-lang.org/main/python-api/generated/triton.language.broadcast_to.html

Demonstrates that the shape parameter can be provided either as a tuple or as individual arguments.

```python
# These are equivalent
broadcast_to(x, (32, 32))
broadcast_to(x, 32, 32)
```

--------------------------------

### triton.language.device_print()

Source: https://triton-lang.org/main/python-api/generated/triton.language.device_print.html

Print values at runtime from the device. The first argument must be a string literal, and subsequent arguments can be scalars or tensors.

```APIDOC
## triton.language.device_print(prefix, *args, hex=False, _semantic=None)

### Description
Print the values at runtime from the device. String formatting does not work for runtime values, so you should provide the values you want to print as arguments. The first value must be a string, all following values must be scalars or tensors.

### Method
`device_print`

### Parameters
- **prefix** (string) - Required - A prefix to print before the values. This is required to be a string literal.
- **args** (scalars or tensors) - Optional - The values to print. They can be any tensor or scalar.
- **hex** (bool) - Optional - Print integers in hexadecimal and floating-point values in hexadecimal floating-point notation.
```

--------------------------------

### triton.autotune

Source: https://triton-lang.org/main/python-api/generated/triton.autotune.html

Decorator for auto-tuning a triton.jit'd function. It evaluates a list of configurations and selects the best one based on benchmarking.

```APIDOC
## triton.autotune

### Description
Decorator for auto-tuning a `triton.jit`'d function. It evaluates a list of configurations and selects the best one based on benchmarking.

### Method
Decorator

### Signature
`triton.autotune(configs, key, prune_configs_by=None, reset_to_zero=None, restore_value=None, pre_hook=None, post_hook=None, warmup=None, rep=None, use_cuda_graph=False, do_bench=None, cache_results=False)`

### Parameters
- **configs** (list[triton.Config]) - Required - A list of `triton.Config` objects.
- **key** (list[str]) - Required - A list of argument names whose change in value will trigger the evaluation of all provided configs.
- **prune_configs_by** (dict) - Optional - A dict of functions used to prune configs. Fields: 'perf_model', 'top_k', 'early_config_prune'.
- **reset_to_zero** (list[str]) - Optional - A list of argument names whose value will be reset to zero before evaluating any configs.
- **restore_value** (list[str]) - Optional - A list of argument names whose value will be restored after evaluating any configs.
- **pre_hook** (lambda args, reset_only) - Optional - A function called before the kernel is called. Overrides the default pre_hook for 'reset_to_zero' and 'restore_value'.
- **post_hook** (lambda args, exception) - Optional - A function called after the kernel is called. Overrides the default post_hook for 'restore_value'.
- **warmup** (int) - Optional - Warmup time (in ms) to pass to benchmarking (deprecated).
- **rep** (int) - Optional - Repetition time (in ms) to pass to benchmarking (deprecated).
- **use_cuda_graph** (bool) - Optional - Whether to use CUDA graphs. Defaults to False.
- **do_bench** (lambda fn, quantiles) - Optional - A benchmark function to measure the time of each run.
- **cache_results** (bool) - Optional - Whether to cache autotune timings to disk. Defaults to False.

### Example
```python
@triton.autotune(configs=[
    triton.Config(kwargs={'BLOCK_SIZE': 128}, num_warps=4),
    triton.Config(kwargs={'BLOCK_SIZE': 1024}, num_warps=8),
  ],
  key=['x_size'] # the two above configs will be evaluated anytime
                 # the value of x_size changes
)
@triton.jit
def kernel(x_ptr, x_size, BLOCK_SIZE: tl.constexpr):
    ...
```

### Notes
- When all configurations are evaluated, the kernel will run multiple times. To avoid undesired behavior, use the `reset_to_zero` argument to reset the value of the provided tensor to zero before running any configuration.
- If the environment variable `TRITON_PRINT_AUTOTUNING` is set to `"1"`, Triton will print a message to stdout after autotuning each kernel, including the time spent autotuning and the best configuration.
```

--------------------------------

### device_print

Source: https://triton-lang.org/main/python-api/triton.language.html

Print the values at runtime from the device.

```APIDOC
## device_print

### Description
Print the values at runtime from the device.

```

--------------------------------

### triton.language.tensor_descriptor.__init__

Source: https://triton-lang.org/main/python-api/generated/triton.language.tensor_descriptor.html

Constructor for the tensor_descriptor class. Not called by user code.

```APIDOC
## `__init__`(self, handle, shape, strides, block_type)

### Description
Initializes a tensor descriptor. This method is not called by user code.

### Parameters
- **handle** (object) - The handle to the tensor descriptor.
- **shape** (List[tensor]) - The shape of the tensor.
- **strides** (List[tensor]) - The strides of the tensor.
- **block_type** (block_type) - The block type of the descriptor.

### Returns
None
```

--------------------------------

### ravel(x[, can_reorder])

Source: https://triton-lang.org/main/python-api/generated/triton.language.tensor.html

Returns a contiguous flattened view of x.

```APIDOC
## ravel(x[, can_reorder])

### Description
Returns a contiguous flattened view of x.

### Method
ravel

### Parameters
- **x** (tensor) - Required - The input tensor.
- **can_reorder** (bool) - Optional - Whether reordering is allowed.

### Returns
A flattened view of the input tensor.
```

--------------------------------

### do_bench_proton

Source: https://triton-lang.org/main/python-api/triton.testing.html

Benchmark the runtime of kernels invoked by the provided function using the Proton profiler.

```APIDOC
## do_bench_proton

### Description
Benchmark the runtime of kernels invoked by the provided function using the Proton profiler.

### Usage
Call this function with the function to benchmark to measure kernel execution times using Proton.
```

--------------------------------

### inline_asm_elementwise

Source: https://triton-lang.org/main/python-api/triton.language.html

Execute inline assembly over a tensor.

```APIDOC
## inline_asm_elementwise

### Description
Execute inline assembly over a tensor.

```

--------------------------------

### static_print

Source: https://triton-lang.org/main/python-api/triton.language.html

Print the values at compile time.

```APIDOC
## static_print

### Description
Print the values at compile time.

```

--------------------------------

### assume

Source: https://triton-lang.org/main/python-api/triton.language.html

Allow compiler to assume the cond is True.

```APIDOC
## assume

### Description
Allow compiler to assume the cond is True.

```

--------------------------------

### reduce(self, axis, combine_fn[, keep_dims, ...])

Source: https://triton-lang.org/main/python-api/generated/triton.language.tensor.html

Forwards to the reduce() free function.

```APIDOC
## reduce(self, axis, combine_fn[, keep_dims, ...])

### Description
Forwards to the reduce() free function.

### Method
reduce

### Parameters
- **self** (tensor) - Required - The input tensor.
- **axis** (int) - Required - The axis along which to reduce.
- **combine_fn** (callable) - Required - The reduction function.
- **keep_dims** (bool) - Optional - Whether to keep the reduced dimension.
- **...** - Optional - Additional arguments.

### Returns
The reduced tensor.
```

--------------------------------

### rsqrt(self[, _semantic])

Source: https://triton-lang.org/main/python-api/generated/triton.language.tensor.html

Forwards to the rsqrt() free function.

```APIDOC
## rsqrt(self[, _semantic])

### Description
Forwards to the rsqrt() free function.

### Method
rsqrt

### Parameters
- **self** (tensor) - Required - The input tensor.
- **_semantic** (optional) - Optional - Semantic information.

### Returns
The reciprocal square root of the input tensor.
```

--------------------------------

### triton.Config

Source: https://triton-lang.org/main/python-api/triton.html

An object that represents a possible kernel configuration for the auto-tuner to try.

```APIDOC
## triton.Config

### Description
An object that represents a possible kernel configuration for the auto-tuner to try.

### Usage
```python
config = triton.Config({'BLOCK_SIZE': 128}, num_warps=4)
```
```

--------------------------------

### triton.language.device_assert(cond, msg='', mask=None, _semantic=None)

Source: https://triton-lang.org/main/python-api/generated/triton.language.device_assert.html

Asserts a condition at runtime from the device. Requires the environment variable TRITON_DEBUG to be set to a value besides '0' for the assertion to have any effect. The condition must be a boolean tensor, and the message must be a string literal.

```APIDOC
## triton.language.device_assert(cond, msg='', mask=None, _semantic=None)

### Description
Asserts a condition at runtime from the device. Requires the environment variable `TRITON_DEBUG` to be set to a value besides `0` for the assertion to have any effect. Using the Python `assert` statement is equivalent, except that the second argument must be provided and must be a string.

### Parameters
- **cond** (boolean tensor) - Required - The condition to assert.
- **msg** (string literal) - Optional - The message to print if the assertion fails. Defaults to an empty string.
- **mask** (tensor) - Optional - A mask to conditionally apply the assertion.
- **_semantic** - Optional - Internal semantic parameter.

### Example
```python
tl.device_assert(pid == 0)
assert pid == 0, f"pid != 0"
```
```

--------------------------------

### num_programs

Source: https://triton-lang.org/main/python-api/triton.language.html

Returns the number of program instances launched along the given axis.

```APIDOC
## num_programs

### Description
Returns the number of program instances launched along the given axis.

### Parameters
- **axis** (int) - Required - The axis along which to get the number of programs.

### Returns
- The number of program instances along the specified axis.
```

--------------------------------

### triton.testing.perf_report(benchmarks)

Source: https://triton-lang.org/main/python-api/generated/triton.testing.perf_report.html

Marks a function for benchmarking. The benchmark can then be executed by using the .run method on the return value.

```APIDOC
## triton.testing.perf_report(benchmarks)

### Description
Marks a function for benchmarking. The benchmark can then be executed by using the `.run` method on the return value.

### Parameters
- **benchmarks** (List of `Benchmark`) - Required - Benchmarking configurations.

### Returns
The return value has a `.run` method to execute the benchmark.
```

--------------------------------

### perf_report

Source: https://triton-lang.org/main/python-api/triton.testing.html

Mark a function for benchmarking.

```APIDOC
## perf_report

### Description
Mark a function for benchmarking.

### Usage
Decorate a function with `@perf_report` to enable benchmarking and plot generation.
```

--------------------------------

### triton.language.atomic_min(pointer, val, mask=None, sem=None, scope=None, _semantic=None)

Source: https://triton-lang.org/main/python-api/generated/triton.language.atomic_min.html

Performs an atomic min at the memory location specified by pointer. Returns the data stored at pointer before the atomic operation.

```APIDOC
## triton.language.atomic_min(pointer, val, mask=None, sem=None, scope=None, _semantic=None)

### Description
Performs an atomic min at the memory location specified by `pointer`. Returns the data stored at `pointer` before the atomic operation.

### Parameters
- **pointer** (Block of dtype=triton.PointerDType) - Required - The memory locations to operate on
- **val** (Block of dtype=pointer.dtype.element_ty) - Required - The values with which to perform the atomic operation
- **mask** (optional) - Optional - Mask for the operation (not detailed in source)
- **sem** (str, optional) - Optional - Specifies the memory semantics for the operation. Acceptable values are "acquire", "release", "acq_rel" (stands for "ACQUIRE_RELEASE"), and "relaxed". If not provided, defaults to "acq_rel" semantics.
- **scope** (str, optional) - Optional - Defines the scope of threads that observe the synchronizing effect of the atomic operation. Acceptable values are "gpu" (default), "cta" (cooperative thread array, thread block), or "sys" (stands for "SYSTEM"). The default value is "gpu".
- **_semantic** (optional) - Optional - Internal parameter (not detailed in source)

### Return Value
Returns the data stored at `pointer` before the atomic operation.

### Usage Notes
This function can also be called as a member function on `tensor`, as `x.atomic_min(...)` instead of `atomic_min(x, ...)`.
```

--------------------------------

### reshape(self, *shape[, can_reorder, ...])

Source: https://triton-lang.org/main/python-api/generated/triton.language.tensor.html

Forwards to the reshape() free function.

```APIDOC
## reshape(self, *shape[, can_reorder, ...])

### Description
Forwards to the reshape() free function.

### Method
reshape

### Parameters
- **self** (tensor) - Required - The input tensor.
- **shape** (int) - Required - The new shape.
- **can_reorder** (bool) - Optional - Whether reordering is allowed.
- **...** - Optional - Additional arguments.

### Returns
The reshaped tensor.
```

--------------------------------

### triton.language.all

Source: https://triton-lang.org/main/python-api/generated/triton.language.all.html

Documentation for the triton.language.all function, which reduces a tensor to a boolean indicating whether all elements are nonzero along a specified axis.

```APIDOC
## triton.language.all

### Description
Returns whether all elements in `input` are nonzero along the provided `axis`. The result has dtype `int1`. NaN and infinity are treated as nonzero.

### Method
Function call

### Signature
`triton.language.all(input, axis=None, keep_dims=False)`

### Parameters
- **input** (Tensor) - Required - The input values
- **axis** (int | None) - Optional - The dimension to reduce. If None, reduce all dimensions
- **keep_dims** (bool) - Optional - If true, keep the reduced dimensions with length 1

### Returns
Tensor with dtype `int1`

### Usage
This function can also be called as a member function on `tensor`, as `x.all(...)` instead of `all(x, ...)`.
```

--------------------------------

### permute(self, *dims[, _semantic])

Source: https://triton-lang.org/main/python-api/generated/triton.language.tensor.html

Forwards to the permute() free function.

```APIDOC
## permute(self, *dims[, _semantic])

### Description
Forwards to the permute() free function.

### Method
permute

### Parameters
- **self** (tensor) - Required - The input tensor.
- **dims** (int) - Required - The new dimension order.
- **_semantic** (optional) - Optional - Semantic information.

### Returns
The permuted tensor.
```

--------------------------------

### triton.testing.do_bench_cudagraph(fn, rep=20, grad_to_none=None, quantiles=None, return_mode='mean')

Source: https://triton-lang.org/main/python-api/generated/triton.testing.do_bench_cudagraph.html

Benchmarks the runtime of the provided function. The function can be configured with repetition time, gradient reset, quantiles, and return mode. Returns the runtime in milliseconds.

```APIDOC
## triton.testing.do_bench_cudagraph

### Description
Benchmark the runtime of the provided function.

### Method
Function call

### Signature
`triton.testing.do_bench_cudagraph(fn, rep=20, grad_to_none=None, quantiles=None, return_mode='mean')`

### Parameters
- **fn** (_Callable_) - Required - Function to benchmark
- **rep** (_int_) - Optional - Repetition time (in ms). Default is 20.
- **grad_to_none** (_torch.tensor_, _optional_) - Optional - Reset the gradient of the provided tensor to None.
- **quantiles** (_optional_) - Optional - Quantiles for return mode. If set, returns a list of floats.
- **return_mode** (_str_) - Optional - The statistical measure to return. Options are "min", "max", "mean", "median", or "all". Default is "mean".

### Returns
The runtime(s) in milliseconds: a single float for a scalar `return_mode`, or a list of floats if `quantiles` is set or `return_mode="all"`.

### Return Type
float | list[float]
```
Tokens:
Raw