Fix sort kernel launch bug when nrows exceed gridDim.y limit (65535) (#3050)

Author: guoqingbao

candle-core/src/sort.rs

@@ -94,7 +94,7 @@ mod cuda {
             let nrows = elem_count / ncols;
             let ncols_pad = next_power_of_2(ncols);
             let cfg = LaunchConfig {
-                grid_dim: (1, nrows as u32, 1),
+                grid_dim: (nrows as u32, 1, 1),
                 block_dim: (ncols_pad as u32, 1, 1),
                 shared_mem_bytes: (ncols_pad * std::mem::size_of::<u32>()) as u32,
             };

candle-kernels/src/sort.cu

@@ -15,7 +15,7 @@ template<int order, typename T>
 static __device__ void k_argsort(const T * x, uint32_t * dst, const int ncols, int ncols_pad) {
     // bitonic sort
     int col = threadIdx.x;
-    int row = blockIdx.y;
+    int row = blockIdx.x;
 
     if (col >= ncols_pad) {
         return;