An efficient way of reduce on the similarity matrix AB^T along two axes

[Yosshi999]

I want to calculate max(AB^T, axis=0) and max(AB^T, axis=1) in triton, where A is MxK and B is NxK matrix (M, N >> K).
I implemented a triton kernel to calculate only max(AB^T, axis=1) using tiling:

import torch

import triton
import triton.language as tl
from triton.runtime import driver

@triton.jit
def sim_kernel(
    A,  # Input matrix (M, K)
    B,  # Input matrix (N, K)
    Out, # Output vector (M,)
    stride_am, stride_ak,
    stride_bn, stride_bk,
    m, n,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    K: tl.constexpr
):
    start_m = tl.program_id(0)
    A_block_ptr = tl.make_block_ptr(
        base=A,
        shape=(m, K),
        strides=(stride_am, stride_ak),
        offsets=(start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, K),
        order=(1, 0))
    B_block_ptr = tl.make_block_ptr(
        base=B,
        shape=(K, n),
        strides=(stride_bk, stride_bn),
        offsets=(0, 0),
        block_shape=(K, BLOCK_N),
        order=(0, 1))
    Out_block_ptr = tl.make_block_ptr(
        base=Out,
        shape=(m,),
        strides=(1,),
        offsets=(start_m * BLOCK_M,),
        block_shape=(BLOCK_M,),
        order=(0,))

    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
    a = tl.load(A_block_ptr)
    for start_n in range(0, n, BLOCK_N):
        b = tl.load(B_block_ptr, boundary_check=(1,), padding_option="nan")
        ab = tl.dot(a, b)  # [BLOCK_M, BLOCK_N], nan if out-of-bounds
        m_i = tl.maximum(m_i, tl.max(ab, 1))  # max operation ignores nan
        B_block_ptr = tl.advance(B_block_ptr, (0, BLOCK_N))
    tl.store(Out_block_ptr, m_i.to(Out.type.element_ty))

def sim(a, b):
    a_rows, a_dim = a.shape
    b_rows, b_dim = b.shape
    assert a_dim == b_dim
    dim = a_dim
    assert dim == 64

    BLOCK_M = 32
    BLOCK_N = 32
    out = torch.zeros(a_rows, dtype=a.dtype, device=a.device)
    assert out.stride(0) == 1
    sim_kernel[(triton.cdiv(a_rows, BLOCK_M), 1, 1)](
        a, b, out,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        a_rows, b_rows,
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        K=dim,
        num_warps=4)
    return out

torch.manual_seed(0)
a = torch.randn(1200, 64, device="cuda").half()
b = torch.randn(1800, 64, device="cuda").half()
out_triton = sim(a, b)
out_torch = torch.max(torch.einsum("mk,nk->mn", a, b), dim=1).values
assert torch.allclose(out_triton, out_torch), (out_triton, out_torch)
print("assertion ok")

However, we need access A and B twice to calculate both max(AB^T, axis=0) and max(AB^T, axis=1), like below:

torch.manual_seed(0)
a = torch.randn(1200, 64, device="cuda").half()
b = torch.randn(1800, 64, device="cuda").half()
a2b = sim(a, b)  # max(ab^T, axis=1)
b2a = sim(b, a)  # max(ab^T, axis=0)

I think this is not I/O efficient.
Do you have any idea to improve this kernel?

[ThomasRaoux]

Questions about writing kernels are better asked on http://discord.com/invite/gpumode