Taichi: For Loops with Grouped NDRange

[yuanming-hu]

Concisely describe the proposed feature

for I in ti.grouped(ti.ndrange((x0, y0), (x1, y1))):
   a[I] = I[0] + I[1] * 2

# equivalent to 
for i, j in ti.ndrange((x0, y0), (x1, y1)):
   a[i, j] = i + j * 2

# a fancier application
for I in ti.grouped(ti.ndrange(*(((0, n),) * dim))):
   a[I] = 0

Describe the solution you'd like (if any)
Modify https://github.com/taichi-dev/taichi/blob/39b2f49ad2d8c88f0fa752b74c17004cdc23ad6d/python/taichi/lang/transformer.py#L274

Please make use of print_preprocessed=True to learn about the treatment of ti.grouped and ti.ndrange.

Additional comments
Will greatly simplify
https://github.com/taichi-dev/taichi/blob/39b2f49ad2d8c88f0fa752b74c17004cdc23ad6d/examples/mgpcg_advanced.py#L46

[xumingkuan]

What do loop_range at https://github.com/taichi-dev/taichi/blob/39b2f49ad2d8c88f0fa752b74c17004cdc23ad6d/python/taichi/lang/transformer.py#L388
and global at https://github.com/taichi-dev/taichi/blob/1dbd05a0b17d39193f51220be648673fd65be6d7/taichi/python/export_lang.cpp#L266 mean?

[yuanming-hu]

global is the global tensor that you are loop over. This is for struct-fors.

[xumingkuan]

I'm trying to reuse the code in if is_ndrange_for but I'm not sure how to modify
https://github.com/taichi-dev/taichi/blob/39b2f49ad2d8c88f0fa752b74c17004cdc23ad6d/python/taichi/lang/transformer.py#L310

Shall I and how do I change I in for I in ti.grouped(ti.ndrange((x0, y0), (x1, y1))): to (I[0], I[1]) at this line?

[yuanming-hu]

I'm not sure if I understand your question, but this is probably an indication that the more efficient way is to figure it out on your own :-)

[xumingkuan]

May I transform

for I in ti.grouped(ti.ndrange((x0, y0), (x1, y1))):
   a[I] = I[0] + I[1] * 2

to

if 1:
    I = [0] * 2
    for __i, __j in ti.ndrange((x0, y0), (x1, y1)):
        I[0] = __i
        I[1] = __j
        a[I] = I[0] + I[1] * 2

(Should I be a list?)

[yuanming-hu]

I see. No, I should be a ti.Vector(x, dt=ti.i32).

[yuanming-hu]

Done. Lot of thanks to @xumingkuan.