Impossible to generate random numbers with constexpr seed

[GuillaumeLeclerc]

@triton.jit
def blop(
    output,
    RAND_SIZE: tl.constexpr,
    SEED: tl.constexpr
):
    o = tl.arange(0, RAND_SIZE) + tl.program_id(axis=0)
    random_values = tl.rand(SEED, o)
result = ch.zeros(10, dtype=ch.int32).cuda()
blop[(1,)](result, RAND_SIZE=16, SEED=0)

yields

AttributeError("'constexpr' object has no attribute 'to'")

[jlebar]

cc @manman-ren

[manman-ren]

I will try to triage this later this afternoon. If anyone wants to take a look before, feel free!

[manman-ren]

This "help wanted" label looks new :] @jlebar

[manman-ren]

I can reproduce with top of trunk. With a simple fix

-def philox(seed, c0, c1, c2, c3, n_rounds: tl.constexpr = N_ROUNDS_DEFAULT):
+def philox(seedin, c0, c1, c2, c3, n_rounds: tl.constexpr = N_ROUNDS_DEFAULT):
+    seed = seedin

It works, the error is because we can't apply "to(tl.uint64)" to a constexpr. The fix (which I am not sure if it is correct) is to copy the constexpr to another variable and apply "to(tl.uint64)" on the variable.

[Jokeren]

1. You can convert seed to a triton tensor in the philox function.
2. Or to workaround, you can remove the constexpr annotation from SEED.

[lijinpei]

Can be fixed by inserting to_tensor in philox, but maybe ast visiter should automatically insert a to_tensor call when we are about to get a attribute access error?

[Jokeren]

That might be difficult. You could have a string or other objects as a constexpr, which cannot be converted to a tensor.

[Jokeren]

Then in the AST visitor, you'll have to check if the constexpr's value is an integer or a float, only then you can convert to a tensor