[BUG] cudf::binary_operation ignores cuda context when registering JIT compiled PTX

[magnatelee]

Describe the bug

cudf::binary_operation currently ignores the CUDA context of the caller thread, which makes the JIT compiled PTX loaded on a wrong device. Even worse is that cudf::binary_operation does not check the CUresult from the kernel launch, so the error is being silently ignored, and noticed only with cuda-memcheck.

[jlowe]

Any idea how the calling thread's context is being ignored here? Is this a case where a thread is being created without an explicit context and CUDA is auto-selecting a (potentially incorrect) device when it implicitly initializes the context? If so that will cause problems in Spark with the RAPIDS plugins in a multi-GPU setup where the GPU device is assigned at the application level (not through CUDA_VISIBLE_DEVICES).

[devavret]

This is happening in Jit where the compiled kernel is being registered with only one context. On a subsequent call from a different context, this fails. It would affect cases where the different threads are assigned different devices, as is the case with @magnatelee's usage.

If spark uses one libcudf process per GPU then this won't affect it. It if uses one thread per GPU then it will.

[devavret]

I'm investigating a fix such that the in-memory cache is stored per context.

[jlowe]

If spark uses one libcudf process per GPU then this won't affect it.

Ah, great to hear. The Spark RAPIDS plugin currently only uses one GPU per process.

[harrism]

So what is left in this bug, @devavret ?

[devavret]

The issue also asks for a check

Even worse is that cudf::binary_operation does not check the CUresult from the kernel launch, so the error is being silently ignored, and noticed only with cuda-memcheck.

I implemented this in NVIDIA/jitify#67 and after that the launch call in Launcher.h needs to be replaced withsafe_launch

[jrhemstad]

@devavret @magnatelee is this still an issue?

[devavret]

Not a lot is left. Just needs to replace the launch in

cudf/cpp/src/jit/launcher.h

Line 97 in 2780a8c

get_kernel().configure_1d_max_occupancy(0, 0, 0, stream.value()).launch(args...);

with safe_launch. I'll make a quick PR tomorrow.

[ttnghia]

Hi! How is this going?

[devavret]

I had a branch for it but I can't find it anymore. Must be lost in my corrupted git. Here's a new one #7510