From 8c16db2c42d5c6da3ec204193fe9923db55f592a Mon Sep 17 00:00:00 2001
From: Lucas C Wilcox <lucas@swirlee.com>
Date: Wed, 3 Apr 2024 16:20:51 -0700
Subject: [PATCH] Avoid capturing `MArray`s in `BoundsError`

This quirk allows bounds checking of `MArray`s on device in more cases.
This is a workaround to address
<https://github.com/JuliaGPU/CUDA.jl/issues/2313>.

This change was proposed upstream here
   <https://github.com/JuliaGPU/CUDA.jl/pull/2314>
and is no longer needed if accepted upstream.
---
 ext/RavenCUDAExt.jl | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/ext/RavenCUDAExt.jl b/ext/RavenCUDAExt.jl
index 83c1efc..833b739 100644
--- a/ext/RavenCUDAExt.jl
+++ b/ext/RavenCUDAExt.jl
@@ -3,6 +3,7 @@ module RavenCUDAExt
 import Raven
 import Adapt
 import MPI
+import StaticArrays
 
 isdefined(Base, :get_extension) ? (using CUDA) : (using ..CUDA)
 isdefined(Base, :get_extension) ? (using CUDA.CUDAKernels) : (using ..CUDA.CUDAKernels)
@@ -58,4 +59,11 @@ else
     Adapt.adapt_storage(::CUDA.KernelAdaptor, ::MPI.Comm) = nothing
 end
 
+CUDA.@device_override function Base.checkbounds(A::StaticArrays.MArray, I...)
+    @inline
+    checkbounds(Bool, A, I...) ||
+        CUDA.@print_and_throw("BoundsError while indexing an MArray.")
+    nothing
+end
+
 end # module RavenCUDAExt