commit f1d67b8773d4330416342e708b096bf0969060b9
author Kojo Acquah <KoolJBlack@users.noreply.github.com> Mon Jun 13 23:40:37 2022 -0400
committer GitHub <noreply@github.com> Mon Jun 13 20:40:37 2022 -0700
tree 6b93c5d14ae2abb8a00fcdb3f7df1dbd24880bc0
parent 75a02aba9fdd9a46a258642fdb6991fa49da6886

Topk Parallel Reduction Initial Implementation (#9364)

Stand alone first version of the TopK parallel reduction implementation. This includes:
* New `TopkOpSplitReduction` OpRewritePattern 
* `SplitReduction.cpp` with the implementation details of the split reduction for topk
* `Passes.td` updated to integrate the topk split reduction pass. Pass ratio is passed by argument.  

There is `split-reduction.mlir` test demonstrating functionality.

## First Level Topk Reductions in Parallel 

Following the "map - reduce" style of parallelization with a single level of parallel reductions. The pass transforms an applicable TopkOp with a single reduction into a `TopkOp` with multiple reductions in parallel followed by another `TopkOp` combining the reductions into the final output. 

The first computation in parallel expands the input tensor shape by breaking the reduction dimension into multiple parallel reductions (upping the rank of the input). `TopkOp` already processes multiple reductions in parallel.  The parallel results are then collapsed into a single input for the final `TopkOp`. Currently, only static shapes are supported. Shapes must be aligned with the reduction factor. Ops with input indices aren't supported.

Part of #9383