[Performance] Improve segment_matmul by reducing launching overheads

[yaox12]

CUTLASS grouped gemm requires copying matrix pointers and layouts to the device memory, which brings significant "launch" overheads, more concretely, 7 pageable H2D copies. This PR sets up the arguments for grouped gemm in a CPU pinned buffer manually and copy it to the device memory at once to reduce such overheads.

Other changes include setting CUDA stream for the grouped gemm, adding proper C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK.

Performance

Benchmarking with the following script, this PR reduces the op time from 0.29 ms to 0.05 ms on my desktop (RTX 3090).

import torch
import time
import pyg_lib

def bench_pyg(a, b, seg):

    seg = torch.tensor(seg)
    pyg_lib.ops.segment_matmul(a, seg, b)
    torch.cuda.synchronize()

    tic = time.time()
    for _ in range(10):
        pyg_lib.ops.segment_matmul(a, seg, b)
    torch.cuda.synchronize()
    print(f"{(time.time() - tic) * 100:.2f} ms")

if __name__ == "__main__":
    num_seg = 20
    hid_dim = 64
    num_ele = 10000

    torch.manual_seed(42)
    device = torch.device("cuda:0")

    a = torch.rand((num_ele, hid_dim), device=device)
    b = torch.rand((num_seg, hid_dim, hid_dim), device=device)

    seg = torch.randint(num_ele, (num_seg - 1,)).sort()[0].tolist()
    seg = [0] + seg + [num_ele]

    bench_pyg(a, b, seg)

cc @rusty1s @puririshi98

[codecov-commenter]

Codecov Report

Merging #213 (2a76c9e) into master (c04fb60) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##           master     #213   +/-   ##
=======================================
  Coverage   83.49%   83.49%           
=======================================
  Files          26       26           
  Lines         848      848           
=======================================
  Hits          708      708           
  Misses        140      140           

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[rusty1s]

Thank you! @puririshi98 do you mind to take a look?

[yaox12]

@puririshi98 Can I get your review?

[puririshi98]

LGTM!

[puririshi98]

thanks for this! @yaox12