triton-lang/triton

Calculate tl.dot, the output is the matrix of col major, but the result is wrong

Open

#1,647 opened on May 10, 2023

 (4 comments) (0 reactions) (0 assignees)MLIR (3,126 forks)github user discovery
help wanted

Repository metrics

Stars
 (19,982 stars)
PR merge metrics
 (Avg merge 2d 18h) (185 merged PRs in 30d)

Description

My test code is following:

def test_dot(M_, N_, K_, num_warps_): 

      @triton.autotune(
             configs=[triton.Config({'BLOCK_M': M_, 'BLOCK_N': N_, 'BLOCK_K': K_}, num_stages=1, num_warps=num_warps_)],
             key=[]
       )
      @triton.jit
       def matmul_kernel(
             a_ptr, b_ptr, c_ptr, M, N, K,
             BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr
       ):
            offs_am = tl.arange(0, BLOCK_M)
            offs_bn = tl.arange(0, BLOCK_N)
            offs_k = tl.arange(0, BLOCK_K)

            a_ptrs = a_ptr + (offs_am[:, None] * K + offs_k[None, :])
            b_ptrs = b_ptr + (offs_k[:, None] * N + offs_bn[None, :])
 
            a = tl.load(a_ptrs)
            b = tl.load(b_ptrs)

            c = tl.dot(a, b, allow_tf32 = False)
            
            c_ptrs = c_ptr + 1 * offs_am[:, None] + M * offs_bn[None, :]
            tl.store(c_ptrs, c)

     def  matmul(a, b):
            c = torch.empty((M_, N_), device = a.device, dtype = a.dtype)
            c = torch.as_strided(c, (M_, N_), c.stride()[::-1])
            grid = lambda META: (
                  triton.cdiv(M_, M_) * triton.cdiv(N_, N_),
            )
            matmul_kernel[grid](
                     a, b, c,
                     M_, N_, K_,
            )
           return c

    torch.manual_seed(0)
    a = torch.randn(M_, K_, dtype = torch.float).cuda()
    b = torch.randn(K_, N_, dtype = torch.float).cuda()
    triton_output = matmul(a, b).cpu()
    torch_output = torch.matmul(a, b).cpu()
    assert triton.testing.allclose(triton_output, torch_output), (triton_output, torch_output)

But when M!=N, the triton output will be different with torch_output. If c_ptrs = c_ptr + N * offs_am[:, None] + 1 * offs_bn[None, :], which means the output of tl.dot is row major, the results are equal.

I checked the generated ttgir, to confirm that the output of tl.dot is col major or the order is {0 ,1}:

Contributor guide