Experimental support for ti.func multi-return with scalar values #536 (3)

[archibate]

Related issue id = #536

[archibate]

Sorry for messing commits from #539 here... but this multi-return can't be done without that AST fix :) So pls merge #539 first.

[yuanming-hu]

Thanks! This is really helpful. Sorry about the forced-push. I had to maintain a clear changeset for reviewing.

[yuanming-hu]

Would it be possible to briefly test this? For example, max and abs. Note that we should not spend too much time on this - The IR should be slightly improved to systematically support functions, including those with multiple returns. I'm still thinking about how that should be done.

[archibate]

This is breaking test because ti.Matrix return is not supported. Should __retval = ti.expr_init(mat.zeros())

[archibate]

How to zero-initialize given types? eg. ti.expr_zero_init(ti.f32).
I just use ti.expr_init(None) and ignoring -> ti.Matrix for now. (21067393

[archibate]

Suggestions are super welcome. I use a new function ti.type_init(ti.mat(2, 2)) to do zero-init.

[yuanming-hu]

Sorry I was on a meeting this morning. Can we try something like ti.expr_init(Matrix.zero(2, 2, dt=ti.f32))?

[archibate]

Sorry I was on a meeting this morning. Can we try something like ti.expr_init(Matrix.zero(2, 2, dt=ti.f32))?

That's exactly what I've done. The difference is because, we can't write def func() -> ti.Matrix(2, 2), since ti.Matrix(2, 2) will be evaluated and you get segfault. So we need ti.mat(2, 2) - a partially-applicated version of ti.Matrix.

[archibate]

Check out this example:

import taichi as ti

ti.init()#print_preprocessed=True)

x = ti.var(ti.f32, ())
u = ti.Vector(2, dt=ti.f32, shape=())
z = ti.Vector(3, dt=ti.f32, shape=())

@ti.kernel
def func(t: ti.f32):
  a = ti.Vector([2, 3])
  s = fake_absolute(t)
  x[None] = s
  u[None] = normalized(a)
  z[None] = zero_by_default(a)

@ti.func
def fake_absolute(x): # will guess return type by alloca, better not
  if x > 0:
    return float(x)
  else:
    return int(-x)

@ti.func
def zero_by_default(x) -> ti.vec(3): # different from ti.Vector (the value container), ti.vec is only used as type signature
  pass

@ti.func
def normalized(v) -> ti.vec(2, dt=ti.f32):
  return v / 2

func(233.33)
print(x[None])
func(-233.33)
print(x[None])
print(u[None][0])
print(u[None][1])
print(z[None][2])

[yuanming-hu]

Considering the fact that

• We want @ti.func to be generic, so that, say, your normalize function will work not only for Vector(2) but also Vector(3).
• Very soon we will migrate to a system that builds functions using Taichi IR instead of Python inlining, so here we should consider this migration cost

I suggest

• In the ti.func transform, insert a taichi_lang_core.initialize_funcion_scope call at the very beginning and taichi_lang_core.finalize_funcion_scope at the end. This helps us track the return values.
• Then transforming return statements, insert a taichi_lang_core.create_return.
• In each function_scope, on the first create_return call, allocate a set of mutable local variables to store the return values. If returning a scalar, allocate a scalar; if returning a matrix, allocate N mutable local variables, where N is the number of entries in the matrix. Note that instead of using type annotation (which stops generic programming), the return type of a function is defined by the first return statement in the function body
• For each return, store the return value to the mutable local variables. The data type of these return values will be inferred on the first (in AST order) LocalStoreStmt during type checking.
• At the end of the function, simply return the mutable local variables. Maybe reconstruct Python classes (e.g. Matrix) so that users get a Matrix instead of a tuple of ti.Expr

The switch to inline-based functions to IR-based functions will not be too hard if we follow this strategy.

[archibate]

We want @ti.func to be generic, so that, say, your normalize function will work not only for Vector(2) but also Vector(3).

This is worked only with inline ones, right? Like ti.template() does.

[archibate]

Sounds cool! Will work on this trr.

[archibate]

Got Basic Block in function 'func_c4_0__kernel_0_serial' does not have terminator!.
May want me create a new llvm::Function (3).

[yuanming-hu]

Got Basic Block in function 'func_c4_0__kernel_0_serial' does not have terminator!.

This is likely because a return statement is missing. Every llvm function should be terminated by a ret, even for void functions (use builder->CreateRetVoid).

[archibate]

Tried adding to different place, become Terminator found in the middle of a basic block! now.

[yuanming-hu]

Tried adding to different place, become Terminator found in the middle of a basic block! now.

Oh, please print out the LLVM IR to see if you have a return statement at the middle of any basic block. Note that LLVM follows a control flow graph + basic block IR strategy, and return statements can only be the last statement of each basic block.

[archibate]

==========
kernel {
  $0 = offloaded  {
    function body // new BasicBlock named function_leave
    return $5 // branch to function_leave
    return $8
    function leave // call CreateRetVoid here
    <i32x1>  $5 = const [666]
    print a, $5
    // was CreateRetVoid expected to be here?
  }
}
==========

(The function is inlined for now)

[yuanming-hu]

Oh I see. For the LLVM backend, maybe we can simply create a real LLVM function (llvm::Function) and then insert an LLVM function call (call)?

Or maybe it's easier if you start with adding function body definition/function calls to your OpenGL backend, since you will then be issuing C++ code which is much easier to debug.

(Of course we will need a function inlining pass sooner or later)

[yuanming-hu]

After digging into your implementation, I feel like there's a misunderstanding.

We want to end up with an IR system that supports function calls, for example

==========
func <i32x1> max(<i32x1> a, <i32x1> b) {
  <i32x1> cond = a > b
  if (cond) {
    return a
  } else {
    return b
  }
}

kernel {
  $0 = offloaded  {
    <i32x1> $1 = const [233]
    <i32x1> $2 = const [666]
    <i32x3> $3 = call max ($1, $2)            // this func call can also be optionally inlined. 
    print a, $3
  }
}
==========

taichi_lang_core.initialize_funcion_scope and taichi_lang_core.finalize_funcion_scope are just for creating the function AST, which will get lowered into a SSA function body like the max func above.

[yuanming-hu]

After implementing this IR extension, multi-return functions could be handled systematically. However, this may be a lot of work and probably we should use multiple PRs. I can go ahead and implement the scaffold of the IR system if you want.

(Originally the quick & hacky solution is to create a local variable in the Python scope to store the return value, without creating any Taichi functions since the Python function will get inlined. However, that's just a temporary solution.)

[archibate]

Was get rid of Terminator found in the middle of a basic block! impossible without creating a llvm::Function? I'm thinking about step-by-step strategy, and I understand what we'll get at the end.

[yuanming-hu]

Was get rid of Terminator found in the middle of a basic block! impossible without creating a llvm::Function? I'm thinking about step-by-step strategy, and I understand what we'll get at the end.

When a function is inlined, it's return value should be store in a local variable for its caller to read, instead of creating a return. In LLVM, return means exiting the current function (e.g. the outer-most level parallel for).

However, ideally, we don't want to inline everything - that's why we are creating functions in Taichi IR, instead of inlining everything in the Python scope.

[archibate]

I can go ahead and implement the scaffold of the IR system if you want.

Oh that would be great! I don't have too much time on reading LLVM docs during school days, if you could help, may helps me a lot!

[yuanming-hu]

Oh that would be great! I don't have too much time on reading LLVM docs during school days, if you could help, may helps me a lot!

Sounds good! I may be slow since I have many things to do simultaneously, but I already have a clear plan in my mind :-)

[archibate]

In LLVM, return means exiting the current function (e.g. the outer-most level parallel for).

So CreateRetVoid is actually exit() in LLVM?
Btw I found that the relation between return and exit, is just like the relation between break and return.

However, ideally, we don't want to inline everything - that's why we are creating functions in Taichi IR, instead of inlining everything in the Python scope.

Yeah, maybe translate FuncBody to inline ones is what lowering should do... just like how for i in range(3) is expanded.

[yuanming-hu]

So CreateRetVoid is actually exit() in LLVM?

It's just like return in C++ and will exit the current function scope.

Yeah, maybe translate FuncBody to inline ones is what lowering should do... just like how for i in range(3) is expanded.

The Python-side AST generation currently does the inlining job. You are right that the mechanism is the same as for i in ti.static(range(3)). (Note ti.static here)

I have to go to sleep as I have an early meeting tomorrow morning. Good night!

[archibate]

Have a good dream!

[archibate]

How to add llvm::Function other than main kernel?

[archibate]

tmp give up, since not in v0.6.0 roadmap.

[yuanming-hu]

How to add llvm::Function other than main kernel?

tmp give up, since not in v0.6.0 roadmap.

Let me spend a few more days further systematically thinking about this issue. Introducing new components to the IR system needs careful considerations. Sorry about the delay.

[archibate]

I found that llvm::Function must spec super clear return type like Int32, does templates solve return type too?
Btw, failed due to spdlog after merged with master.

[yuanming-hu]

I found that llvm::Function must spec super clear return type like Int32, does templates solve return type too?

That's exactly what I'm thinking about. We need the IR to lower the compound return values into what LLVM supports. Again, designing a new IR extension needs a lot of careful considerations and I would spend some time on it before starting implementation.

[archibate]

Replaced by #612