[AutoTune] Performing kernel with multiple configs may lead to silent results mismatch

[mieshkiwrk]

Describe the bug

Below example where each config writes first and second half of the result, when running with single config there's expected assert about 50% mismatched results, but when running with both configs (REPRO=1) it passes which may lead to some silent issues.

import os
import torch
import triton
import triton.language as tl

def get_configs():
    configs = [
        triton.Config({'OFFSET': 0, 'ELEMENTS': 512}),
        triton.Config({'OFFSET': 512, 'ELEMENTS': 512})
    ]
    repro = int(os.environ.get('REPRO', '0'))
    return configs if repro else configs[:1]

@triton.autotune(
    configs=get_configs(),
    key=[],
)
@triton.jit
def add_kernel(
    x_ptr,
    y_ptr,
    output_ptr,
    OFFSET: tl.constexpr,
    ELEMENTS: tl.constexpr    
):
    offsets = tl.arange(OFFSET, ELEMENTS + OFFSET)
    x = tl.load(x_ptr + offsets)
    y = tl.load(y_ptr + offsets)
    output = x + y
    tl.store(output_ptr + offsets, output)


def get_triton_result():
    x = torch.ones(1024, device='xpu')
    y = torch.ones(1024, device='xpu')
    output = torch.empty(1024, device='xpu')
    add_kernel[(1,)](x, y, output)
    return output

def get_torch_result():
    x = torch.ones(1024, device='xpu')
    y = torch.ones(1024, device='xpu')

    output = x + y
    return output

triton_output = get_triton_result()
torch_output = get_torch_result()
torch.testing.assert_close(triton_output, torch_output)

Environment details

Triton: latest

[mieshkiwrk]

Due to this I think we should use autotune with below options in our CI because some regression can potentially be missed.

@triton.autotune(
    configs=get_configs(),
    key=[],
    reset_to_zero=['output_ptr']
)

or

@triton.autotune(
    configs=get_configs(),
    key=[],
    restore_value=['output_ptr']
)

[Egor-Krivov]

I think it's a nice to have

[whitneywhtsang]

Please evaluate if there is any regressions (functional, performance, CI time) before enabling the change.

[mieshkiwrk]

Test PR with restore_value for xpu benchmarks: #5785

[mieshkiwrk]

Please evaluate if there is any regressions (functional, performance, CI time) before enabling the change.

In terms of functional/performance it seems the same.
CI time in general fluctuates a little for benchmarks but I don't see any spectacular difference so it looks good to go in my opinion

[wdziurdz]

PR adding restore_value to our benchmarks was opened by @mieshkiwrk and is currently under review, pending approval: #5785

Next steps: Waiting for Triton team feedback on whether adding restore_value to the upstream tutorials makes sense. Issue: triton-lang/triton#9230