Cannot call a function with two arguments from a triton kernel

[alexander-zinoviev]

Hi,

I want to write a triton kernel that in turn calls another function that takes multiple arguments. It is something similar to the leaky_relu in the matrix multiplication example https://triton-lang.org/master/getting-started/tutorials/03-matrix-multiplication.html.
The only difference is that I want "my_leaky_relu" to take 2 or more arguments not 1.

@triton.jit
def my_kernel(a_ptr, b_ptr):
my_function(a_ptr, b_ptr)

@triton.jit
def my_function(a_ptr, b_ptr):
return

What I get when I declare a function with two parameters and call it with two arguments is:

Function expects 1 arg(s), got 2 [wrong-arg-count]
Expected: ()
Actually passed: (, _)

and it looks that whatever number of parameters I declare on the function declaration the python compiler still thinks it has only one parameter.

Is there any limitation of parameters a callable can take?

[wujingyue]

The following code works for me at 5686c51. This issue may have been fixed. Can you try to run your program again?

import torch
import triton
import triton.language as tl

@triton.jit
def add(x, y):
  return x + y

@triton.jit
def multiply_by_two(x_ptr, length, y_ptr, 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 < length
  x = tl.load(x_ptr + offsets, mask=mask)
  y = add(x, x)
  tl.store(y_ptr + offsets, y, mask=mask)

x = torch.tensor([-2, -1, 0, 1, 2], device='cuda')
y = torch.empty_like(x)
length = x.numel()
block_size = 8
multiply_by_two[(triton.cdiv(length, block_size),)](x, length, y, block_size=block_size)
print(x)
print(y)