search

rapidsai / raft

RAFT contains fundamental widely-used algorithms and primitives for machine learning and information retrieval. The algorithms are CUDA-accelerated and form building blocks for more easily writing high performance applications.
https://docs.rapids.ai/api/raft/stable/
Apache License 2.0
781 stars 195 forks source link

[FEA] A helper for transposing an mdarray #1010

Open achirkin opened 2 years ago

achirkin commented 2 years ago

When optimizing cuda kernels, modifying the layout of the data seems to be a common approach - to facilitate coalesced data access. That often boils down to transposing a multidimensional data along a pair of indices. It would be handy to have something like this as a helper for the mdarray type to reduce the boilerplate code.

Inspired by https://github.com/rapidsai/raft/pull/926#discussion_r1019992068

cjnolet commented 2 years ago

This is great! A few of us were just talking about the need for something like this yesterday.

Though do you see any reason to do this only for mdarray and not for mdspan?

Tagging @hcho3 @wphicks @dantegd @divyegala

tfeher commented 2 years ago

The particular use case we have would take an mdspan as an input and create a new mdarray with the new layout.