Let real function/kernel support returning a struct/multiple values

[lin-hitonami]

Concisely describe the proposed feature
Currently, the Taichi function supports returning a arbitrary number of return values, and the return values can be scalars, matrices and structs, and the Taichi function is inlined into the Taichi kernel at compile time.

However, the Taichi kernel only supports returning a single return value, and the return value can only be a scalar or a matrix.

The WIP real function (#602) is a Taichi function that does not inline into the kernel. The kernel calls it instead of inlining it. In LLVM-based backends, the real function receives an argument with type RuntimeContext where the arguments and the address of return value buffer are stored, same as the Taichi kernel. Currently, the real function only supports a single scalar return value.

I would like to let the real function (and maybe the Taichi kernel) support returning a struct, and then let them support multiple return values (we can construct a struct containing the return values).

Describe the solution you'd like (if any)
A struct can be seen as a tree. A scalar element is a leaf, and a struct element is a subtree. A matrix can be seen as a struct when real_matrix=False, and maybe a special scalar when real_matrix=True. For simplicity, I will not consider matrices (will consider it in the implementation) and I only consider the return value of the real function in the explanation below. The code below are simplified pseudo-code.

For this struct:

from ti.types import struct
s0 = struct(c=ti.f32, d=ti.i64)
s1 = struct(a=ti.i32, b=s0)

There are 3 leaf elements for s1: a, b.c, and b.d.

Basic Implementation

1. We allocate a u64 array which size equals to the number of the leaf elements in the struct as the return buffer in the caller.

u64 ret_buf[3];

2. We pass the pointer to the callee

RuntimeContext *rc;
rc->result_buffer=ret_buf
callee(rc);

3. We cast the elements to the respective types.

i32 a = bitcast_from_u64(ret_buf[0], i32)
f32 c = bitcast_from_u64(ret_buf[1], f32)
i64 d = bitcast_from_u64(ret_buf[2], i64)

4. We compose a struct type in Python

• In C++:

Expr return_struct_element(int index){
    return element[index];
}

• In Python:

ret = s1(a=return_struct_element(0) b=s0(c=return_struct_element(1), d=return_struct_emement(2)))

Possible improvements

We can construct the return values by sequence to a llvm::StructType, and get the return value via the GEP instructions. By doing this, we can save the memory if not all the elements are 64-bit wide.

Where should we store the return value?

There are two places we can store it: the stack or the heap. I suggest we store it on the stack if the size of the return value is very small, otherwise we should store it on the heap, because the stack of CUDA is small.

If we store it on the heap, we may want an allocator using preallocated memory because malloc on CUDA is slow. We may use the NodeManager to allocate and recycle small items which have the same size.