Cross product of two 3-dimensional vectors?

[cTatu]

Hi, I wanted to try implement the cross product using Triton (same functionality as torch.linalg.cross).

This is what I tried:

import torch
import triton
import triton.language as tl

@triton.jit
def cross_kernel(A, B, output, n_elements, BLOCK_SIZE: tl.constexpr):
  pid = tl.program_id(axis=0)
  block_start = pid * BLOCK_SIZE
  offsets = block_start + tl.arange(0, BLOCK_SIZE)
  mask = offsets < n_elements

  a_block = tl.load(A + offsets, mask=mask)
  b_block = tl.load(B + offsets, mask=mask)
  res_block = tl.zeros_like(a_block)

  res_block[:, 0] = a_block[:, 1] * b_block[:, 2] - a_block[:, 2] * b_block[:, 1]
  res_block[:, 1] = a_block[:, 2] * b_block[:, 0] - a_block[:, 0] * b_block[:, 2]
  res_block[:, 2] = a_block[:, 0] * b_block[:, 1] - a_block[:, 1] * b_block[:, 0]

  tl.store(output + offsets, res_block, mask=mask)

def cross(x: torch.Tensor, y: torch.Tensor):
    output = torch.empty_like(x)
    assert x.is_cuda and y.is_cuda and output.is_cuda
    n_elements = output.shape[0]
    
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )
    cross_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=256)

    return output

if __name__ == '__main__':
  size = 256*100
  x = torch.rand(size, 3, device='cuda')
  y = torch.rand(size, 3, device='cuda')
  output_triton = cross(x, y)

But this error came up:

CompilationError: at 11:2:def cross_kernel(A, B, output, n_elements, BLOCK_SIZE: tl.constexpr):
  pid = tl.program_id(axis=0)
  block_start = pid * BLOCK_SIZE
  offsets = block_start + tl.arange(0, BLOCK_SIZE)
  mask = offsets < n_elements

  a_block = tl.load(A + offsets, mask=mask)
  b_block = tl.load(B + offsets, mask=mask)
  res_block = tl.zeros_like(a_block)

  res_block[:, 0] = a_block[:, 1] * b_block[:, 2] - a_block[:, 2] * b_block[:, 1]
  ^
AssertionError()

Clearly I did something wrong. Can anyone nudge me in the right direction please?

Thank you

[Hprairie]

Triton doesn't support indexing, meaning that you can't do something like this res_block[:, 0].

[cTatu]

Thank you for the answer. But how could I implement the three dimensional cross product with triton then?

[Hprairie]

Idk I don't know much about cross product algorithms. I would try to find other ways to parallelize, so that the data you are loading into your kernel requires no indexing.