causal_conv1d_fn fails when number of channels is 2^16 or more

[KeAWang]

Once we use 2^16 or more channels, we get the following error:

RuntimeError: CUDA error: invalid configuration argument
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Minimal reproducible code:

import torch
from causal_conv1d import causal_conv1d_fn

device = "cuda"
b,  l = 1, 2**2
k = 2


# this works
c = 2**16 - 1
x = torch.randn(b, c, l, dtype=torch.float32, device=device)  # batch, channel, seq_len
weight = torch.randn(c, k, dtype=torch.float32, device=device)
y = causal_conv1d_fn(x, weight)

# this fails (c >= 2**16)
c = 2**16
x = torch.randn(b, c, l, dtype=torch.float32, device=device)  # batch, channel, seq_len
weight = torch.randn(c, k, dtype=torch.float32, device=device)
y = causal_conv1d_fn(x, weight)

This might be because the implementation parallelizes the blocks across the batch and channel axes, but CUDA blocks only go up to 2^16-1 along the y axis:

causal-conv1d/csrc/causal_conv1d_fwd.cu

Line 65 in f8c2467

const int channel_id = blockIdx.y;

Flipping the .x and .y should fix the issue (at least for the non-channellast implementation).