Add prototype for optimized dense2jagged kernel for 1 jagged dim

[chirayuG-nvidia]

Sharing optimized prototype kernel for 2D dense_to_jagged operations using vectorized operations.
CC: @mjanderson09

[netlify]

✅ Deploy Preview for eclectic-stroopwafel-199537 canceled.

Name	Link
🔨 Latest commit	23e692e
🔍 Latest deploy log	https://app.netlify.com/sites/eclectic-stroopwafel-199537/deploys/62df4e7ae3678f000928a1f6

[facebook-github-bot]

@mjanderson09 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[jspark1105]

This is amazing! Thanks for the PR. I added a few comments here for visibility but these can be certainly addressed on our side.

[jspark1105]

Strides are manually used vs. tensor accessor. Does it make a difference in perf?

[chirayuG-nvidia]

It shouldn't make any difference in perf as long as accessor is also adding same amount of instructions. I used explicit calculations as it makes the code self explanatory.

[jspark1105]

nit: we can use CUDA_KERNEL_LOOP(idx, total_jagged_rows) which is defined in pytorch/aten/src/ATen/cuda/detail/KernelUtils.h

[jspark1105]

Instead of applying vectorization sort of all-or-nothing, can we do best effort like the first loop applied up to columns rounded down to 128-bit, the second loop applied to up to columns rounded down to 64-bit, and so on? In this way when e_dim is say 576-bit we can use 128-bit vectorization up to 512-bit and then 64-bit vectorization. Or this will regress perf?

[chirayuG-nvidia]

Performance will depend on the overhead of math instructions that will get added to do the alignment check by each thread of the warp every time it executes. I think you can find both types of problem size where doing the incremental vector-bit decrement will help performance and other (perfectly aligned cases) where it will reduce it.
IMO ideal solution is to allocate contiguous tensors with appropriate pitch in each dimension that aligns the vectorized load/store. Not sure if pytorch does that already.

[mjanderson09]

So you're only checking alignment with respect to e_dim, not alignment of the actual pointers, right? If that's the case, can we just do the following? Or do we have to check alignment of "values" and "dense" pointers.

    VecType128::TType * values128 = (VecType128::TType *)&(values[0]);
    VecType128::TType * dense128 = (VecType128::TType *)&(dense[0]);
    VecType64::TType * values64 = (VecType64::TType *)&(values[0]);
    VecType64::TType * dense64 = (VecType64::TType *)&(dense[0]);
    VecType32::TType * values32 = (VecType32::TType *)&(values[0]);
    VecType32::TType * dense32 = (VecType32::TType *)&(dense[0]);
    for(int real_row = values_row ; real_row < nnz ; real_row += blockDim.y * gridDim.y) {
        int dense_row = rows[real_row];
        int dense_col = cols[real_row];
        for(int tid = threadIdx.x ; tid < E/8 ; tid += blockDim.x) {
            values128[tid + real_row * (E/8)] = dense128[tid + dense_col * (E/8) + dense_row * L * (E/8)];
        }
        for(int tid = threadIdx.x + (E/8)*8 ; tid < E/4 ; tid += blockDim.x) {
            values64[tid + real_row * (E/4)] = dense64[tid + dense_col * (E/4) + dense_row * L * (E/4)];
        }
        for(int tid = threadIdx.x + (E/4)*4 ; tid < E/2 ; tid += blockDim.x) {
            values32[tid + real_row * (E/2)] = dense32[tid + dense_col * (E/2) + dense_row * L * (E/2)];
        }
        for(int tid = threadIdx.x + (E/2)*2 ; tid < E ; tid += blockDim.x) {
            values[tid + real_row * (E)] = dense[tid + dense_col * (E) + dense_row * L * (E)];
        }
    }

[chirayuG-nvidia]

I am doing a alignment check with e_dim * sizeof(scalar_t) because :

1. It will ensure jagged_out tensor is aligned
2. Because of broadcast and the only effective stride in dense tensor being same as e_dim, so my check ensures overall things will align for all the consequent pointer location that kernel will try to read from.

The alternate that you suggested will run into pointer alignment error. For eg: say output jagged is 130 bit wide (65 half's), so very first row will work fine with your logic, however the 2nd row will try to write first 8 half's with address that is not aligned for 128-bit operation.

for(int tid = threadIdx.x ; tid < E/8 ; tid += blockDim.x) {

Note that in my code I use a warp to service two rows (block.x = 16) , I use tid = threadIdx.x % blockDim.x