Merge pull request #374 from ChrisRackauckas-Claude/add-jlarrays-support

Add JLArrays.jl support for CPU-based array computations

Author: ChrisRackauckas

.github/workflows/CI.yml

@@ -33,6 +33,7 @@ jobs:
           - "pre"
         group:
           - CPU
+          - JLArrays
     uses: "SciML/.github/.github/workflows/tests.yml@master"
     with:
       julia-version: "${{ matrix.version }}"

Project.toml

@@ -28,12 +28,14 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b"
 
 [extensions]
 AMDGPUExt = ["AMDGPU"]
 CUDAExt = ["CUDA"]
+JLArraysExt = ["JLArrays"]
 MetalExt = ["Metal"]
 oneAPIExt = ["oneAPI"]
 
@@ -46,6 +48,7 @@ DiffEqBase = "6.122"
 DocStringExtensions = "0.9"
 ForwardDiff = "0.10.38, 1"
 GPUArraysCore = "0.2"
+JLArrays = "0.1, 0.2, 0.3"
 KernelAbstractions = "0.9"
 LinearSolve = "3"
 Metal = "1"
@@ -65,6 +68,7 @@ oneAPI = "2"
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

ext/JLArraysExt.jl

@@ -0,0 +1,7 @@
+module JLArraysExt
+using JLArrays: JLBackend
+import DiffEqGPU
+
+DiffEqGPU.maybe_prefer_blocks(::JLBackend) = JLBackend()
+
+end

test/Project.toml

@@ -6,6 +6,7 @@ DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"

test/utils.jl

@@ -11,6 +11,9 @@ elseif GROUP == "oneAPI"
 elseif GROUP == "Metal"
     using Metal
     Metal.MetalBackend()
+elseif GROUP == "JLArrays"
+    using JLArrays
+    JLArrays.JLBackend()
 elseif GROUP == "CPU"
     using KernelAbstractions
     KernelAbstractions.CPU()