[FEA] gather/scatter on other dims

[xiaonans]

I notice cutlass has examples to show how to fuse gather/scatter with gemm, such as https://github.com/NVIDIA/cutlass/tree/main/examples/36_gather_scatter_fusion. But it is assumed that “when the matrix is row major, we select rows. When the matrix is column major, we select columns". This assumption limits the application to gather/scatter on other dims, e.g. selecting columns from a row major output matrix.

Would you please provide examples to show how to fuse gather/scatter on other dims?

[thakkarV]

you cannot gather/scatter within the contiguous dimension, since by definition that would make it not contiguous anymore. that would kill performance since you would have no way to coaleacse and vectorize those. if you want want to do this pattern however, you are welcome to implement a custom kernel for it similar to example 59