Add support for index_select OP (backend implementation)

[azazhu]

This PR is backend implementation of index_select(lookup_tv, dim, index_tv) fusion:

• support arbitrarily rank
• support arbitrarily dim value
• index_tv can become consumer;
• lookup_tv cannot become consumer.

OP def is here: https://pytorch.org/docs/stable/generated/torch.index_select.html

TODO: move index calculation from codegen.cpp to lower_index.cpp

[naoyam]

@zasdfgbnm added support of select in #2179. I think it'd be relatively straightforward to build index_select on top of his work. Please take a look at the PR, especially around here: https://github.com/csarofeen/pytorch/pull/2179/files#diff-3c08060d5a1993c9599d5cbb2d7a805258ae1b4e166682b497c4f5fd376ba5d3R199-R205

[azazhu]

@zasdfgbnm added support of select in #2179. I think it'd be relatively straightforward to build index_select on top of his work. Please take a look at the PR, especially around here: https://github.com/csarofeen/pytorch/pull/2179/files#diff-3c08060d5a1993c9599d5cbb2d7a805258ae1b4e166682b497c4f5fd376ba5d3R199-R205

Referred to Xiang's impl and moved lower_index from codegen.cpp to lower_index.cpp

[naoyam]

@azazhu Can you please also add a validation check to validateIr so that all of the lookup TVs are ensured to be fusion inputs? Something like:

for (auto expr: fusion->exprs()) {
   if (expr->isA<SelectOp>() || expr->isA<IndexSelectOp>()) {
     TORCH_CHECK(expr->input(0)->isFusionInput(), "Lookup input must be a fusion input: ", expr->toString());
   }
}

[azazhu]

@azazhu Can you please also add a validation check to validateIr so that all of the lookup TVs are ensured to be fusion inputs? Something like:

for (auto expr: fusion->exprs()) {
   if (expr->isA<SelectOp>() || expr->isA<IndexSelectOp>()) {
     TORCH_CHECK(expr->input(0)->isFusionInput(), "Lookup input must be a fusion input: ", expr->toString());
   }
}

Done. see https://github.com/csarofeen/pytorch/pull/2183/files#diff-e9719a5d2b00c30404cd4d11ebaf37dea60b9447ee893aaddb9a9ef6fd9aaebbR1313-R1323

[naoyam]

LGTM. Just left one more minor comment.

Thanks for adding this new operator!