Repository: pytorch
Issue: [triton 3.2] test_convolution_as_mm failure on A100
Labels: triaged,upstream triton

### 🐛 Describe the bug

**Update**: Tracked in triton in https://github.com/triton-lang/triton/issues/5204. As a workaround, a PR was reverted in 3.2 (but not the main triton branch). We also filed an nvbug.

This is uncovered during the triton 3.2 pin update. @embg and I are looking into this, but I'm filing an issue to provide additional context.

The issue can be reproed without the inductor test with this script:

```python
import torch


import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties

'''
@triton_heuristics.template(
    num_stages=5,
    num_warps=4,
    triton_meta={'signature': {'in_ptr0': '*fp32', 'arg_A': '*fp32', 'arg_B': '*fp32', 'out_ptr0': '*fp32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108, warp_size=32), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3), equal_to_1=())]},
    inductor_meta={'kernel_name': 'Placeholder.DESCRIPTIVE_NAME', 'backend_hash': '5C8FD749058C1A642DD2829E30FB5522FD8965A7E9F6B2F8F8AB489A76569BB4', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': True, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
)
'''
@triton.jit
def triton_mm(in_ptr0, arg_A, arg_B, out_ptr0, bias_ptr):
    GROUP_M : tl.constexpr = 8
    EVEN_K : tl.constexpr = False
    ALLOW_TF32 : tl.constexpr = False
    # ALLOW_TF32 : tl.constexpr = True
    ACC_TYPE : tl.constexpr = tl.float32
    B_PROLOGUE_CAST_TYPE : tl.constexpr = None
    BLOCK_M : tl.constexpr = 64
    BLOCK_N : tl.constexpr = 64
    BLOCK_K : tl.constexpr = 64
    A = arg_A
    B = arg_B

    M = 512
    N = 34
    K = 33
    if M * N == 0:
        # early exit due to zero-size input(s)
        return
    stride_am = 33
    stride_ak = 1
    stride_bk = 1
    stride_bn = 33

    # based on triton.ops.matmul
    pid = tl.program_id(0)
    grid_m = (M + BLOCK_M - 1) // BLOCK_M
    grid_n = (N + BLOCK_N - 1) // BLOCK_N

    # re-order program ID for better L2 performance
    width = GROUP_M * grid_n
    group_id = pid // width
    group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
    pid_m = group_id * GROUP_M + (pid % group_size)
    pid_n = (pid % width) // (group_size)

    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    if (stride_am == 1 and stride_ak == M) or (stride_am == K and stride_ak == 1):
        ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)
    else:
        ram = rm % M
    if (stride_bk == 1 and stride_bn == K) or (stride_bk == N and stride_bn == 1):
        rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)
    else:
        rbn = rn % N
    rk = tl.arange(0, BLOCK_K)
    A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak)
    B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn)

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
    for k in range(K, 0, -BLOCK_K):
        if EVEN_K:
            a = tl.load(A)
            b = tl.load(B)
        else:
            a = tl.load(A, mask=rk[None, :] < k, other=0.)
            b = tl.load(B, mask=rk[:, None] < k, other=0.)
        if B_PROLOGUE_CAST_TYPE is not None:
            b = b.to(B_PROLOGUE_CAST_TYPE)
        acc += tl.dot(a, b, allow_tf32=ALLOW_TF32)
        A += BLOCK_K * stride_ak
        B += BLOCK_K * stride_bk

    # rematerialize rm and rn to save registers
    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    idx_m = rm[:, None]
    idx_n = rn[None, :]
    mask = (idx_m < M) & (idx_n < N)

    # inductor generates a suffix
    xindex = idx_n + (34*idx_m)
    tmp0 = tl.load(in_ptr0 + (tl.broadcast_to(idx_n, acc.shape)), mask, eviction_policy='evict_last')
    # tl.store(bias_ptr + tl.broadcast_to(xindex, acc.shape), tmp0, mask)
    tmp1 = acc + tmp0
    tl.store(out_ptr0 + (tl.broadcast_to(xindex, acc.shape)), tmp1, mask)
    tl.store(bias_ptr + tl.broadcast_to(xindex, acc.shape), tmp0, mask)

BLOCK_M : tl.constexpr = 64
BLOCK_N : tl.constexpr = 64
BLOCK_K : tl.constexpr = 64
dtype = torch.float32

A = torch.ones(512, 33, device="cuda", dtype=dtype)
B = torch.ones(33, 34, device="cuda", dtype=dtype)
bias = torch.ones(34, device="cuda", dtype=dtype)
out = torch.empty(512, 34, device="cuda", dtype=dtype)
bias_out = torch.empty(512, 34, device="cuda", dtype=dtype)

# ret = triton_mm[(512 // 64,)](bias, A, B, out, bias_out, debug=True)
ret = triton_mm[(512 // 64,)](bias, A, B, out, bias_out)

'''
from triton.tools.disasm import get_sass
sass = get_sass(ret.asm['cubin'])
with open("sass.sass", "w") as f:
    f.write(sass)
'''

expect = A @ B + bias
# breakpoint()
print((expect-out).abs().max())
assert (expect-out).abs().max().item() < 0.01
# v = ((expect-out).abs() > 0.01).to(torch.int32).tolist()
# print("\n".join([str(x) for x in v]))

```

Bisect finds https://github.com/triton-lang/triton/pull/4582 and although this blame seems reliable, it's unclear how the PR causes the issue

### Versions

A100, triton main, pytorch viable/strict

cc @bertmaher @int3 @nmacchioni @chenyang78 @embg @peterbell10 @aakhundov