Having a callable as an argument for jitted function

[mfkasim1]

I need to implement a jitted function with one of its arguments is a callable (jitted as well). Here is an example:

import torch
import triton
import triton.language as tl


@triton.jit
def jitfn(a, b):
    return a + b

@triton.jit
def custom_func(binfn: tl.constexpr, a_ptr, b_ptr, res_ptr, block_size: tl.constexpr):
    block_idxs = tl.arange(0, block_size)
    a = tl.load(a_ptr + block_idxs)
    b = tl.load(b_ptr + block_idxs)
    res = binfn(a, b)
    tl.store(res_ptr + block_idxs, res)

def main():
    device = torch.device("cuda")
    block_size = 16
    a = torch.randn(block_size, device=device)
    b = torch.randn(block_size, device=device)
    res = torch.empty(block_size, device=device)
    custom_func[(1,)](jitfn, a, b, res, block_size=block_size)

if __name__ == "__main__":
    main()

When I ran the code above, it complains:

Traceback (most recent call last):
  File "/home/azureuser/experiments/01-triton/triton_deer.py", line 121, in <module>
    main()
  File "/home/azureuser/experiments/01-triton/triton_deer.py", line 26, in main
    custom_func[(1,)](jitfn, a, b, res, block_size=block_size)
  File "/home/azureuser/base/lib/python3.11/site-packages/triton/runtime/jit.py", line 209, in <lambda>
    return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/azureuser/base/lib/python3.11/site-packages/triton/runtime/jit.py", line 520, in run
    raise TypeError(f"Callable constexpr at index {i} is not supported")
TypeError: Callable constexpr at index binfn is not supported

Is it just not possible to have a jitted function with a jitted function as an argument for now? Or is there any workaround for this case?

[thoangtrvn]

tl.constexpr expects a compile-time constant value, and you're using the Triton function object which is not know at compile-time. You can call it explicitly, though.

@triton.jit
def custom_func(a_ptr, b_ptr, res_ptr, block_size: tl.constexpr):
    block_idxs = tl.arange(0, block_size)
    a = tl.load(a_ptr + block_idxs)
    b = tl.load(b_ptr + block_idxs)
    res = jitfn(a, b)
    tl.store(res_ptr + block_idxs, res)