Taichi: [Lang] [Bug] templates with ti.func

[KLozes]

Describe the bug
Template parameters don't seem to working as expected for ti.func. It seems that paramters are getting turned in to an Expr when being passed into ti.funcs, instead of remaining an original type. Example below.

Log/Screenshots

[Taichi] mode=development
[Taichi] preparing sandbox at /tmp/taichi-01b5av4i
[Taichi] sandbox prepared
[Taichi] <dev mode>, supported archs: [cpu, cuda], commit 979ec630, python 3.6.9
Traceback (most recent call last):
  File "test_template.py", line 32, in <module>
    aTob1(l)
  File "/home/klozes/Documents/software/taichi/python/taichi/lang/kernel.py", line 500, in wrapped
    return primal(*args, **kwargs)
  File "/home/klozes/Documents/software/taichi/python/taichi/lang/kernel.py", line 430, in __call__
    self.materialize(key=key, args=args, arg_features=arg_features)
  File "/home/klozes/Documents/software/taichi/python/taichi/lang/kernel.py", line 310, in materialize
    taichi_kernel = taichi_kernel.define(taichi_ast_generator)
  File "/home/klozes/Documents/software/taichi/python/taichi/lang/kernel.py", line 307, in taichi_ast_generator
    compiled()
  File "test_template.py", line 23, in aTob1
    b[l][I] = sample_a(l, I)
  File "/home/klozes/Documents/software/taichi/python/taichi/lang/kernel.py", line 52, in __call__
    ret = self.compiled(*args)
  File "test_template.py", line 13, in sample_a
    return a[l][I]
  File "/home/klozes/Documents/software/taichi/python/taichi/lang/impl.py", line 82, in subscript
    return value[indices[0]]
TypeError: list indices must be integers or slices, not Expr

To Reproduce
Accessing this multilevel data structure with sample_a() inst working. Even if l is a template parameter in the ti.func and ti.kernel..

import taichi as ti

ti.init()

a = [ti.var(dt=ti.f32) for _ in range(2)]
b = [ti.var(dt=ti.f32) for _ in range(2)]

for l in range(2):
    ti.root.dense(ti.ij, 16).place(a[l],b[l])

@ti.func
def sample_a(l: ti.template(), I):
    return a[l][I]

@ti.kernel
def fill(l: ti.template()):
    for I in ti.grouped(a[l]):
        a[l][I] = l

@ti.kernel
def aTob1(l: ti.template()): # doesnt compile
    for I in ti.grouped(b[l]):
        b[l][I] = sample_a(l, I)

@ti.kernel
def aTob2(l: ti.template()): # works fine
    for I in ti.grouped(b[l]):
        b[l][I] = a[l][I]

for l in range(2):
    fill(l)
    aTob1(l)

for l in range(2):
    for i in range(16):
        for j in range(16):
            assert b[l][i,j] == l

[KLozes]

The preprocessing helps show the issue
The line in after preprocessing sample_a, l = ti.expr_init_func(l_by_value__), should not be happening I think

python3 test_template.py 
[Taichi] mode=development
[Taichi] preparing sandbox at /tmp/taichi-zwcx7ibj
[Taichi] sandbox prepared
[Taichi] <dev mode>, supported archs: [cpu, cuda], commit 2f6fbc51, python 3.6.9
Before preprocessing:
@ti.kernel
def fill(l: ti.template()):
    for I in ti.grouped(a[l]):
        a[l][I] = l

After preprocessing:
def fill():
  if 1:
    ___loop_var = ti.grouped(ti.subscript(a, l))
    I = ti.make_var_vector(size=___loop_var.loop_range().dim())
    ___expr_group = ti.make_expr_group(I)
    ti.core.begin_frontend_struct_for(___expr_group, ___loop_var.loop_range
        ().ptr)
    ti.subscript(ti.subscript(a, l), I).assign(l)
    ti.core.end_frontend_range_for()
    del I

Before preprocessing:
@ti.kernel
def aTob1(l: ti.template()):
    for I in ti.grouped(b[l]):
        b[l][I] = sample_a(l, I)

After preprocessing:
def aTob1():
  if 1:
    ___loop_var = ti.grouped(ti.subscript(b, l))
    I = ti.make_var_vector(size=___loop_var.loop_range().dim())
    ___expr_group = ti.make_expr_group(I)
    ti.core.begin_frontend_struct_for(___expr_group, ___loop_var.loop_range
        ().ptr)
    ti.subscript(ti.subscript(b, l), I).assign(sample_a(l, I))
    ti.core.end_frontend_range_for()
    del I

Before preprocessing:
def sample_a(l: ti.template(), I):
  return a[l][I]

After preprocessing:
def sample_a(l_by_value__: ti.template(), I_by_value__):
  l = ti.expr_init_func(l_by_value__)
  I = ti.expr_init_func(I_by_value__)
  return ti.subscript(ti.subscript(a, l), I)

Traceback (most recent call last):
  File "test_template.py", line 32, in <module>
    aTob1(l)
  File "/home/klozes/Documents/software/taichi/python/taichi/lang/kernel.py", line 500, in wrapped
    return primal(*args, **kwargs)
  File "/home/klozes/Documents/software/taichi/python/taichi/lang/kernel.py", line 430, in __call__
    self.materialize(key=key, args=args, arg_features=arg_features)
  File "/home/klozes/Documents/software/taichi/python/taichi/lang/kernel.py", line 310, in materialize
    taichi_kernel = taichi_kernel.define(taichi_ast_generator)
  File "/home/klozes/Documents/software/taichi/python/taichi/lang/kernel.py", line 307, in taichi_ast_generator
    compiled()
  File "test_template.py", line 23, in aTob1
    b[l][I] = sample_a(l, I)
  File "/home/klozes/Documents/software/taichi/python/taichi/lang/kernel.py", line 52, in __call__
    ret = self.compiled(*args)
  File "test_template.py", line 13, in sample_a
    return a[l][I]
  File "/home/klozes/Documents/software/taichi/python/taichi/lang/impl.py", line 82, in subscript
    return value[indices[0]]
TypeError: list indices must be integers or slices, not Expr

[KLozes]

near line 647 of transformer.py. no if statements for handling template parameters

[archibate]

Thank for reporting this! But IMO @ti.func doesn't support type-hints?

[KLozes]

Yeah. Looking at class Kernel compared to class Func in kernel.py. There is no logic for handling type hints at all in for Func. The metaprogramming doc shows an example of a ti.func with a ti.template argument. So this _should_ be supported syntax. We should definitely add this if we can.

[yuanming-hu]

Thanks for raising this!

Currently ti.func is simply force-inlined so type-hints are not really needed in most cases. In the future, after we have real-function support (https://github.com/taichi-dev/taichi/issues/602), type hints will be an optional constraint on the arguments.

For now, all ti.func arguments are passed by value, so we need to make a copy at the beginning of the function, to prevent modifications of the arguments affect the outer scope. The reason why this fails is that l = ti.expr_init_func(l_by_value__) doesn't work when l is a template...