Propagate Mosaic GPU kernel launch errors to XLA instead of calling abort().
Previously, `mosaic_gpu_launch_kernel` in `runtime.cc` unconditionally invoked `abort_on_error` when `cuLaunchKernelEx` returned a failure (such as `CUDA_ERROR_CAPTURE_ALREADY_FAILED` during CUDA Graph / Command Buffer tracing). This raised `SIGABRT` and killed the process immediately. This change: 1. Updates `mosaic_gpu_launch_kernel` in `runtime.cc` to return `CUresult` directly instead of calling `abort()`. 2. Updates `emitRuntimeDecls` and `launchPreloadedKernel` in `launch_lowering.cc` to declare and return `i32` from the preloaded kernel launcher and host function. 3. Updates `MosaicGpuExecute` and `MosaicHostFunc` in `custom_call.cc` to check the returned `CUresult` and return a descriptive `absl::InternalError` when kernel launch fails, allowing XLA FFI to handle it cleanly. PiperOrigin-RevId: 971955078
R
Rohan Jain committed
f3fcb67b5f07af401a9b4d5b822334cbef77e709
Parent: 8f607bd
Committed by jax authors <google-ml-automation@google.com>
on 8/27/2026, 2:59:16 PM