求助一下CUDA版本的argmax实现问题

如题，想请教一下如果有一个二维数组data[M][N]，想求它在某个维度上的argmax，有没有什么好的使用CUDA加速的办法?最简单的办法是把某一个维度分配到block上然后用for循环去求，但是感觉这种做法没有充分利用GPU。更进一步的如果源数组是一个有很多维度的tensor，如data[N1][N2][N3][...][Nk]这种，有没有什么高效的办法实现这个函数？