Assigning to kernel argument leads to no error but wrong behavior

[yuanming-hu]

Describe the bug

The code below

import taichi as ti

ti.init()

@ti.kernel
def test(a: ti.i32):
    if True:
        a = 1
    print(a)

test(0)

Should either emits an error (kernel argument cannot be assigned), or print 1. But it prints 0:

[Taichi] version 1.0.4, llvm 10.0.0, commit 2827db2c, linux, python 3.8.10
[Taichi] Starting on arch=x64
0

[bobcao3]

I'd advocate for non-mutable function parameters unless it's a pointer/template type, and should make this an error. (ndarrays and fields parameters can be considered as a reference semantically)

[yuanming-hu]

I'd advocate for non-mutable function parameters unless it's a pointer/template type, and should make this an error. (ndarrays and fields parameters can be considered as a reference semantically)

Agreed. Let me make a decision here: kernel arguments (scalars/vectors/matrices) should be non-mutable.

[lin-hitonami]

I once made it immutable in #3105 but this broke some internal code because we allowed them to be mutable in the past, and the PR was reverted very soon. If we do make this change this time, we may need to notify the users in the release note.

[yuanming-hu]

I once made it immutable in #3105 but this broke some internal code because we allowed them to be mutable in the past, and the PR was reverted very soon. If we do make this change this time, we may need to notify the users in the release note.

OK... Thanks for the info! Sounds like we need a deeper discussion on this. If we take the mutable solution and make it watertight, do we need much more work? For example, allowing atomic adds to the argument sounds like a lot of work (you may need a pointer to the argument list?):

@ti.kernel
def foo(a: ti.i32):
    for i in range(10):
        a += i
    print(a)

If the amount of work to make the arguments mutable is huge, we'd rather make it immutable.

[bobcao3]

If it's internal code we are breaking we should fix them ourselves. I'm not exactly sure how this is done? In the IR we have ArgLoad but I don't think it works the same way as allocas. I think we should just make it a compiler warning at first and keep it what it is in the codegen, and then make it an error some versions later

[lin-hitonami]

I once made it immutable in #3105 but this broke some internal code because we allowed them to be mutable in the past, and the PR was reverted very soon. If we do make this change this time, we may need to notify the users in the release note.

OK... Thanks for the info! Sounds like we need a deeper discussion on this. If we take the mutable solution and make it watertight, do we need much more work? For example, allowing atomic adds to the argument sounds like a lot of work (you may need a pointer to the argument list?):

@ti.kernel
def foo(a: ti.i32):
    for i in range(10):
        a += i
    print(a)

If the amount of work to make the arguments mutable is huge, we'd rather make it immutable.

The solution turns out to be pretty simple: we can assign the argument to a new variable.

[yuanming-hu]

The solution turns out to be pretty simple: we can assign the argument to a new variable.

Yeah, that could be a solution. But I'm afraid that may lead to some confusion: if one thread modifies that variable, is that visible to other threads? Users will assume the kernel arguments are shared among threads (since it's kernel-level) instead of being thread-local. What about serial offloaded tasks? E.g.,

@ti.kernel
def foo(a: ti.i32):
    for i in range(10):
        a = i
    print(a) # What should be the result here?

The other thing I worry about is register usage on GPUs. We may need an extra optimization pass to eliminate these introduced temporary variables.

[lin-hitonami]

Yeah, that could be a solution. But I'm afraid that may lead to some confusion: if one thread modifies that variable, is that visible to other threads? Users will assume the kernel arguments are shared among threads (since it's kernel-level) instead of being thread-local. What about serial offloaded tasks? E.g.,

@ti.kernel
def foo(a: ti.i32):
    for i in range(10):
        a = i
    print(a) # What should be the result here?

If we make a a new variable, a is shared among threads. If one thread modifies that variable, it is not guaranteed to be visible to other threads, and it's the user's responsibility to handle that.

The other thing I worry about is register usage on GPUs. We may need an extra optimization pass to eliminate these introduced temporary variables.

Yeah. Writing this may take a while. It's possible that I can finish writing the pass before 1.1.0 releases (2 days left) but I can't 100% guarantee. The fastest temporary solution is to throw a warning (or error?) when user alter a scalar argument in non-outermost scope.

[lin-hitonami]

According to offline discussions, we decide to make kernel arguments immutable.