Wrong values in broadcasted computation

[Edenhofer]

Triton produces wrong values in a broadcasted subtraction:

#!/usr/bin/env python3

import torch

import triton
import triton.language as tl


@triton.jit
def toy_kernel(
    dist_ptr,
    coo_bounds_ptr,
    output_ptr,
    block_size: tl.constexpr,
    n_elements: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    block_start = pid * block_size

    block_off = block_start + tl.arange(0, block_size)
    block_mask = block_off < n_elements
    dist = tl.load(dist_ptr + block_off, mask=block_mask)
    coo_bounds = tl.load(coo_bounds_ptr + block_off, mask=block_mask)

    coo_bounds = tl.reshape(coo_bounds, (1, n_elements))
    dist = tl.reshape(dist, (block_size, 1))
    nr = coo_bounds
    print(nr)
    nr -= dist
    print(nr)

    output = tl.sum(nr, axis=1)
    block_off = block_start + tl.arange(0, block_size)
    tl.store(output_ptr + block_off, output, mask=block_mask)


def toy(dist, coo_bounds):
    block_size = n_elements = 4
    grid = ((dist.shape[0] - 1) // block_size + 1, )
    output = torch.empty_like(dist)
    toy_kernel[grid](
        dist,
        coo_bounds,
        output,
        block_size=block_size,
        n_elements=n_elements,
        num_warps=2,
    )
    return output


dist = 1. + torch.arange(4, dtype=torch.float32, device='cuda')
coo_bounds = 2. + torch.arange(4, dtype=torch.float32, device='cuda')
print(toy(dist, coo_bounds))

# Reference implementation in torch
o = (coo_bounds.reshape(1, -1) - dist.reshape(-1, 1)).sum(axis=1)
print(o)

produces the following incompatible results

tensor([1., 0., 0., 0.], device='cuda:0')  # triton
tensor([10.,  6.,  2., -2.], device='cuda:0')  # reference implementation

[Jokeren]

The problem should have been fixed using triton master. Please verify.

[Edenhofer]

Yup, it works when building from master. Sorry for the noise.