| 28 | |
| 29 | @triton.jit |
| 30 | def argsort_2d_kernel( |
| 31 | x_ptr, |
| 32 | out_vals_ptr, |
| 33 | out_ids_ptr, |
| 34 | stride_m, |
| 35 | N: tl.constexpr, # noqa: N803 |
| 36 | M: tl.constexpr, # noqa: N803 |
| 37 | ): |
| 38 | pid = tl.program_id(0) |
| 39 | offs_n = tl.arange(0, N) |
| 40 | x = tl.load(x_ptr + pid * stride_m + offs_n) |
| 41 | ids = offs_n.to(tl.int32) |
| 42 | sorted_x, sorted_ids = argsort(x, ids, descending=True) |
| 43 | tl.store(out_vals_ptr + pid * stride_m + offs_n, sorted_x) |
| 44 | tl.store(out_ids_ptr + pid * stride_m + offs_n, sorted_ids) |
| 45 | |
| 46 | |
| 47 | @pytest.mark.parametrize("N", [16, 64, 128, 256]) |