Login
Triton.Lang How to handle Block sizes
I am trying to use triton-lang to perform a simple element-wise dot product between a column vector and a matrix that both have complex value. I can make the code work if I don't specify block_sizes but I can't figure out how to cut my grid and how to handle my pointers. I somewhat understand the theory on how it should work but I can't make it work.
def cdot(x: torch.Tensor, y: torch.Tensor):
return x * y
def cdot_triton(x: torch.Tensor, y: torch.Tensor, BLOCK_SIZE):
# preallocate the output
z = torch.empty_like(y)
# check arguments
assert x.is_cuda and y.is_cuda and z.is_cuda
# get vector size
N = z.numel()
# 1D launch kernel where each block gets its own program
grid = lambda meta: (N // BLOCK_SIZE, N // BLOCK_SIZE)
# launch the kernel
cdot_kernel[grid](x.real, x.imag, y.real, y.imag, z.real, z.imag, N, BLOCK_SIZE)
return z
@triton.jit
def cdot_kernel(
x_real_ptr,
x_imag_ptr,
y_real_ptr,
y_imag_ptr,
z_real_ptr,
z_imag_ptr,
N: tl.constexpr, # Size of the vector
BLOCK_SIZE: tl.constexpr, # Number of elements each program should process
):
row = tl.program_id(0)
col = tl.arange(0, 2*BLOCK_SIZE)
if row < BLOCK_SIZE:
idx = row * BLOCK_SIZE + col
x_real = tl.load(x_real_ptr + 2*row)
x_imag = tl.load(x_imag_ptr + 2*row)
y_real = tl.load(y_real_ptr + 2*idx, mask=col<BLOCK_SIZE, other=0)
y_imag = tl.load(y_imag_ptr + 2*idx, mask=col<BLOCK_SIZE, other=0)
z_real = x_real * y_real - x_imag * y_imag
z_imag = x_real * y_imag + x_imag * y_real
tl.store(z_real_ptr + 2*idx, z_real, mask=col<BLOCK_SIZE)
tl.store(z_imag_ptr + 2*idx, z_imag, mask=col<BLOCK_SIZE)
# ===========================================
# Test kernel
# ===========================================
size = 4
dtype = torch.complex64
x = torch.rand((size, 1), device='cuda', dtype=dtype)
y = torch.rand((size, size), device='cuda', dtype=dtype)
out_dot = cdot(x,y)
out_kernel = cdot_triton(x,y, BLOCK_SIZE=2)
This is the output:
tensor([[-0.1322+1.1461j, -0.1098+0.8015j, 0.2948+1.2155j, -0.1326+0.6076j],
[-0.3687+0.4646j, 0.2349+0.5802j, 0.0568+0.9461j, -0.0457+0.3213j],
[ 0.0523+0.9351j, 0.4409+0.5076j, 0.3956+0.4018j, 0.6230+0.9270j],
[-0.3503+0.7194j, -0.3742+0.2311j, -0.3353+0.3884j, -0.3478+0.6724j]],
device='cuda:0')
tensor([[-0.1322+1.1461j, -0.1098+0.8015j, 0.0617+1.0408j, -0.1988+0.4788j],
[ 0.1147+0.2296j, 0.0686+0.1161j, 0.0647+0.4044j, 0.0795+0.6407j],
[-0.2396+0.6326j, -0.3587+0.5878j, -0.1563+0.4028j, -0.2933+0.3294j],
[-0.1214+0.3678j, 0.0440+0.9951j, 0.3342+1.1360j, 0.6796+0.6590j]],
device='cuda:0')
As you can see only the 2 first values of the top row are accurate.
Any ideas on how I can make this element-wise dot product work?
Many thanks!
