[Runtime][Contrib] Support cudnn softmax

[icemelon]

Using cudnn can improve the softmax performance on Nvidia GPU.

@yzhliu @Laurawly @ZihengJiang

[Hzfengsy]

I'm confused why we need int64_t here but later cast to int

[icemelon]

It's because DLTensor defines the shape in int64_t type. There'll be a cast anyway.

[tqchen]

as part of the principle, would be great if we can lookinto making the native op as fast

[icemelon]

@tqchen Yes, I understand that. But the latency difference could be 10x between tvm schedule and cudnn for the input shape like [100, 1024] on V100. I guess to achieve such performance it requires fusion across multiple stage of reduction, which it seems not easy to be implemented in tir.

[tqchen]

k, I am not trying to blocking the PR, merely trying to say it would be great to have such investigation :)

[tqchen]

@wpan11nv @yongfeng-nv can you suggest a bit about possible optimizations that can be done?

[yongfeng-nv]

@wpan11nv @yongfeng-nv can you suggest a bit about possible optimizations that can be done?

We don't know the details, but will look into it.

[wpan11nv]

@wpan11nv @yongfeng-nv can you suggest a bit about possible optimizations that can be done?

The cuda schedule emits 4 kernels, which cause lots of IO overhead. Ideally, we may emit a single kernel for small reduction sizes (e.g. reduction dim n <= 1024)

[tqchen]

#5600 for improving softmax with warp shuffle.