[TIR] Add additional termination condition to For node to enable While loop like feature

[masahi]

This is my proposed solution to add While loop like feature to TIR, in the simplest, the least invasive way. It generalizes the For node termination condition from

loop_var < extent

to

loop_var < extent && test

Using this, we can write binary search as follows (see the complete test case, which implements numpy searchsorted function, here).

lo[0] = 0
hi[0] = n
v = Bptr[i]
num_loop = int(np.log2(n)) + 1

with ib.for_range(0, num_loop, test=(lo[0] < hi[0])) as _:
    mid = lo[0] + tvm.tir.floordiv(hi[0] - lo[0], 2).astype("int32")
    with ib.if_scope(Aptr[mid] < v):
        lo[0] = mid + 1
    with ib.else_scope():
        hi[0] = mid

Cptr[i] = lo[0]

My motivation was to improve GPU NMS performance using while loop, and it indeed did:

NMS workload from PyTorch MaskRCNN:

Without while loop (current main): 4.11 milli sec
With while loop (my branch): 3.66 milli sec

And a crazy 120000 box + 100 max_out_size NMS workload from TF MaskRCNN. The difference is huge because the # of iterations changed from 120000 to 100 (roughly)

Without while loop (current main): 51.31 milli sec
With while loop (my branch): 17.63 milli sec

please review @tqchen @mbrookhart @kevinthesun @zhiics @Laurawly @anijain2305 @trevor-m

[anijain2305]

LGTM! This is awesome stuff.

[mbrookhart]

LGTM. Longer term we probably want to add this to hybrid script and/or tvm script

[tqchen]

It would be helpful to have a RFC discussion, since different strategies changes to the IR can have different implications

[tqchen]

Thanks @masahi . I think it is great to enable support for some form of While loop.

It would be great to have an RFC thread discussing the alternatives in the IR node design. Since the IR node design can impact the general ability to do analysis and will impact how would we engineer future transformation passes.

For example, I can see two possible variants:

• V0: Put the condition into the for loop as it is (the current approach)
• V1: Introduce a separate While node for while loops

V0 means the for loop is somewhat overloaded for both while and For. On one hand it brings the benefit of richer semantics and the minimum set of changes to enable such feature.

This does mean that the visitors to For would need to handle the semantics of break. Given that the for loop is used for regular interval analysis, it could be beneficial to distinguish between "regular structured loop" vs "un-structured loop", unless the condition testing also offers some analysis benefits.

I think this PoC is a great start and would be great to have a design discussion over the RFC

[masahi]

@tqchen Yes, what you said totally makes sense to me. As you said, this is a minimal-change solution, but I think ideally we want a separate While node. Since while loop only makes sense for sequential loop (I think), I think it is better to decouple a simpler While node that doesn't need any analysis from heavy-duty For node.

I'll send a RFC, sure.

[masahi]

The rfc posted https://discuss.tvm.apache.org/t/rfc-add-while-loop-node-to-tir/9028

[masahi]

TIR While node added in #7425