[TE] Support mixing normal and cross-thread reduction

[roastduck]

Currently TVM only supports pure normal (i.e. sequential) reduction or pure cross-thread reduction. Since TVM has not supported nested reduction yet, one is even unable to schedule a mixed reduction manually. I modified the function that lowers cross-thread reduction, to support mixed reduction as well.

The approach is straight forward: First perform normal reduction into local variables in each threads first, and then invoke the original cross-thread reduction intrinsic. It works like this (pseudo-code):

// Divide the loop nest into two parts
normal_red = sequantial loops nest
common = other loops nest
    
// If normal_red is empty, fallback to original code

normal_init = generate init for the temp var
normal_update = generate sequential reduction on the temp var
body genereate cross-thread reduction // original code

// Merge loop nests and add some checks
body = SeqStmt(normal_init, MergeNest(normal_red, normal_update), body)
body = MergeNest(common, body)
return body

A test case is added as a Python unit test.

This is my first PR to TVM, and I am not sure whom to invite as a reviewer. Since this is compiling pass related, @tqchen can you review my code?

[tqchen]

Thanks @roastduck for contributing. cc @vinx13 @merrymercy @Hzfengsy @wpan11nv please also help to take a look:)

[roastduck]

All resolved. @wpan11nv

[wpan11nv]

LGTM. Thanks.

[tqchen]

Thanks @roastduck @wpan11nv !