ti.real_func parameter / return types support more complicated types?

[Enigmatisms]

I am trying to use ti.real_func (which is actually ti.experimental.real_func), in order to accelerate compiling. I read from the documentations that real_func only supports scalar inputs and return, yet I wish to return exactly two scalar digits (which should form a tuple) but I can't do that since real_func requires type annotation.

I also learned that we might return a struct (from this PR #7059), so I tried the following code. The code is just a simple test for input/output, but what confuses me is: When I tried to directly return a vec3, the compiler will tell me "unsupported return type":

File "~/graphics/AdaPT/assets/ti_tests/./real_func.py", line 30, in calculate:
v3_field[i, j] = get_vec(field1[i, j], field2[i, j])
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Unsupported return type: <taichi.lang.matrix.VectorType object at 0x7f4f5c5c4b80>

But I can also just wrap the vec3 by a struct, then return it from real_func decorated function and everything works. Why? This doesn't make sense to me, it seems just so meaningless to wrap the vec3 (and other "unsupported return types") by a struct as a working-around. By the way, the experimental functionality real_func is released since 1.0.0 and until 1.6.0 it remains as experimental. I think this is a good feature, so, when will real_func get fully supported and moved from experimental?

This is my env:

[Taichi] version 1.4.1, llvm 15.0.4, commit e67c674, linux, python 3.10.9

Here is the code, get_struct throws no compilation error while get_vec does.

import taichi as ti
from taichi.math import vec3

ti.init()

field1 = ti.field(float, (6, 6))
field2 = ti.field(float, (6, 6))

@ti.dataclass
class Struct:
    vec: vec3

str_field = Struct.field()
ti.root.dense(ti.ij, (6, 6)).place(str_field)

v3_field = ti.Vector.field(3, float, (6, 6))

@ti.experimental.real_func
def get_struct(a: float, b: float) -> Struct:
    return Struct(vec = vec3([a, b, a + b]))

@ti.experimental.real_func
def get_vec(a: float, b: float) -> vec3:
    return vec3([a, b, a + b])

@ti.kernel
def calculate():
    for i, j in str_field:
        str_field[i, j] = get_struct(field1[i, j], field2[i, j])
        v3_field[i, j] = get_vec(field1[i, j], field2[i, j])


for i in range(6):
    for j in range(6):
        field1[i, j] = j
        field2[i, j] = i

calculate()

for i in range(6):
    for j in range(6):
        struct = str_field[i, j]
        print(f"[{struct.vec.to_numpy().tolist()}])", end = ', ')
    print("")

[lin-hitonami]

We plan to support more types like structs, matrices and ndarrays in the future (maybe in a few months). Please stay tuned.

[Enigmatisms]

Thanks! Hope to see that real_func becomes fully supported and removed from experimental.

I've noticed that today there is a new commit by you: Support ndarray argument for real function. Seems that you are making progress there. So is it difficult for real_functo be just likefunc` except not being inline function?

[lin-hitonami]

It is a bit difficult because we handle a kernel without considering function calls inside it when performing many optimizations. There's a lot to do to let all the optimization passes support real function.

[Enigmatisms]

Does ti.experimental.real_func support ti.template() as parameter type annotation? I tried some, no luck, even with this simple example:

import taichi as ti

@ti.dataclass
class TestStruct:
    p1: float
    p2: float
    val: float

@ti.experimental.real_func
def compute(struct1: ti.template(), struct2: ti.template()) -> float:
    p = (struct1.p1 + struct2.p2) / 2
    return p

@ti.kernel
def process_struct(field1: ti.template(), field2: ti.template()):
    for i in field1:
        field1[i].val = compute(field1[i], field2[i])

ti.init(arch = ti.cuda)

field1 = TestStruct.field()
field2 = TestStruct.field()

ti.root.dense(ti.i, 64).place(field1, field2)

for i in range(64):
    field1[i] = TestStruct(p1 = 1, p2 = 2)
    field2[i] = TestStruct(p1 = 2, p2 = 1)

process_struct(field1, field2)

Outputs:

[Taichi] version 1.6.0, llvm 15.0.4, commit f1c6fbbd, linux, python 3.10.9
[Taichi] Starting on arch=cuda
...
RuntimeError: [demote_operations.cpp:visit@160] Assertion failure: is_integral(lhs->element_type()) && is_integral(rhs->element_type())

[lin-hitonami]

No, expressions in the kernel can't be passed as template to the real function right now...

[Enigmatisms]

One more problem: it seems that real_func can return struct type but can only receive scalar type parameters? So... it's confusing again, why can't we make input parameter receive a struct? The example can be modified to be:

@ti.experimental.real_func
def get_struct(a: Struct) -> Struct:
    return Struct(vec = vec3([a.vec[0], a.vec[1], a.vec[0] + a.vec[1]]))

and

@ti.kernel
def calculate():
    for i, j in str_field:
        struct = Struct(vec = vec3([field1[i, j], field2[i, j], field1[i, j]]))
        str_field[i, j] = get_struct(struct)

Then we will have exception:

[Taichi] version 1.6.0, llvm 15.0.4, commit f1c6fbb, linux, python 3.10.9
[Taichi] Starting on arch=x64
Traceback (most recent call last):
File "./graphics/AdaPT/assets/ti_tests/./real_func.py", line 39, in
calculate()
File "./.conda/envs/adv/lib/python3.10/site-packages/taichi/lang/kernel_impl.py", line 976, in wrapped
raise type(e)("\n" + str(e)) from None
taichi.lang.exception.TaichiTypeError:
File "./graphics/AdaPT/assets/ti_tests/./real_func.py", line 30, in calculate:
str_field[i, j] = get_struct(struct)
^^^^^^^^^^^^^^^^^^
Invalid constant scalar data type: <class 'taichi.lang.struct.Struct3'>

[lin-hitonami]

IIRC it is already supported in the nightly version.

[Enigmatisms]

Ok, thanks