[Bug][Unity] Legalization of cumsum op is broken

[masahi]

R.cumsum is legalized to topi.cumsum according to

https://github.com/apache/tvm/blob/unity/python/tvm/relax/transform/legalize_ops/statistical.py#L90-L92

But this is broken on GPU. topi.cumsum is one of topi ops that cannot be defined by te.compute and thus they are written using IR builder by hand, for CPU and GPU targets separately. Since topi.cumsum is meant to be used only for CPU targets, after legalization we get TIR that has a for loop with parallel kind, for example, even on GPU. Other common ops that would have the same issue include sorting etc.

So in this case, it would be more correct to use topi.cuda.cumsum for legalizing relax.cumsum, meaning the LegalizeOp pass needs to be target dependent. But I think that would break the original semantics of LegalizeOp. DefaultGPUSchedule pass also needs to ignore such legalized ops carefully since they are already scheduled for GPU by hand.

I think we need to discuss what the Relax compilation story is going to be for such ops that were previously implemented using te.extern.

cc @Hzfengsy @MasterJH5574 @quic-sanirudh @tqchen

[tqchen]

In this particular case, we should build up a tensor ir of cumsum, and then bring dlight schedule for them

[tqchen]

opens up a discussion thread in https://discuss.tvm.apache.org/t/discuss-default-compilation-flow-for-scan-and-sort/15949