triton-lang/triton

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

開放

#1,647 建立於 2023年5月10日

 (4 則留言) (0 個反應) (0 位負責人)MLIR (3,136 個分叉)github user discovery
help wanted

倉庫指標

星標
 (19,995 顆星)
PR 合併指標
 (平均合併 2天 18小時) (30 天內合併 185 個 PR)

描述

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}:

貢獻者指南