[Lang] Support continue on all backends

[k-ye]

Related issue = #708

[Click here for the format server]

[k-ye]

@yuanming-hu I just realized that this will require us to encapsulate the actual kernel body as a function (LLVM does this, but Metal doesn't). Otherwise when we use one thread to handle more than one elements, this could prematurely terminate the entire kernel thread.

Alternatively, do you folks think it even makes sense to support continue at the top level? I mean Taichi users are supposed to well understand that the top-level loop as parallelized, yet continue makes sense only for serial loops.

[yuanming-hu]

Otherwise when we use one thread to handle more than one element, this could prematurely terminate the entire kernel thread.

Right, sounds like we have to goto in the source-to-source backends...The other option is to directly translate this into a continue in the source backends. But we have to be very careful when doing that...

Alternatively, do you folks think it even makes sense to support continue at the top level? I mean Taichi users are supposed to well understand that the top-level loop as parallelized, yet continue makes sense only for serial loops.

I think it still makes sense. For example,

for i in x:
    if x[i] > 0:
       # statement
       # statement
       if y[i] > 0:
          # statement
          # statement
          # statement

might be more verbose than

for i in x:
    if x[i] <= 0:
       continue
    # statement
    # statement
    if y[i] <= 0:
       continue
    # statement
    # statement
    # statement

(especially when we have multiple nested ifs)

We should probably document this in the documentation: continue at the top level is only for skipping the rest of the current iteration. There's no guarantee about iteration order in a parallel setting.

[archibate]

Using continue as a shortcut of return in top-level? Nice idea!

Alternatively, do you folks think it even makes sense to support continue at the top level? I mean Taichi users are supposed to well understand that the top-level loop as parallelized, yet continue makes sense only for serial loops.

We want to keep this easy-to-use shortcut since they acts same as the expected in python.
In comparison, we don't want break in top-level is because it can't act same as the expected in python.

[k-ye]

The other option is to directly translate this into a continue in the source backends. But we have to be very careful when doing that...

I guess this means we assume the kernel body is embedded inside a "meta" for-loop implemented by the Taichi runtime:

taichi/taichi/runtime/llvm/runtime.cpp

Lines 1074 to 1077 in 58c64ba

	while (idx < end) {
	func(context, idx);
	idx += block_dim() * grid_dim();
	}

Then I think wrapping the kernel into a function is a much more robust and cleaner approach...

[k-ye]

OK, sounds like I should keep this support. Thanks guys :)

[yuanming-hu]

Then I think wrapping the kernel into a function is a much more robust and cleaner approach...

Right, then I guess the source-to-source backends have to support something like non-linear code emission since now parallel loop bodies are outlined into functions. Pragmatically maybe using a continue is still the easiest way.

[k-ye]

Right, then I guess the source-to-source backends have to support something like non-linear code emission since now parallel loop bodies are outlined into functions.

+1

Pragmatically maybe using a continue is still the easiest way.

Right. I just realized a problem on Metal. Its range_for doesn't use grid-stride loop, yet its struct_for does, gahhh... 😦

[yuanming-hu]

Otherwise when we use one thread to handle more than one element, this could prematurely terminate the entire kernel thread.

Right, sounds like we have to goto in the source-to-source backends...The other option is to directly translate this into a continue in the source backends. But we have to be very careful when doing that...

Alternatively, do you folks think it even makes sense to support continue at the top level? I mean Taichi users are supposed to well understand that the top-level loop as parallelized, yet continue makes sense only for serial loops.

I think it still makes sense. For example,

for i in x:
    if x[i] > 0:
       # statement
       # statement
       if y[i] > 0:
          # statement
          # statement
          # statement

might be more verbose than

for i in x:
    if x[i] <= 0:
       continue
    # statement
    # statement
    if y[i] <= 0:
       continue
    # statement
    # statement
    # statement

(especially when we have multiple nested ifs)

We should probably document this in the documentation: continue at the top level is only for skipping the rest of the current iteration. There's no guarantee about iteration order in a parallel setting.

[k-ye]

(May be related to #689 ) This came up when I was writing this PR, would it be good if we split FrontendStmt and Expression into its own, higher layer of IR? I think the only pass that cares about both sets of IRs is lower_ast?

[yuanming-hu]

Right, ideally FrontendIR and Expressions are unique to the Taichi language, which should be translated into the CHI IR via lower_ast.

This separation will improve modularity since now the frontend and middle-end IR (CHI) are separate into two modules. In fact, lower_ast should not even use IRVisitor (which is for the middle-end IR). We should ideally have a FrontendASTVisitor...