WIP: GPU Support

Project.toml

@@ -4,6 +4,7 @@ authors = ["EarthSciML Authors and Contributors"]
 version = "0.24.0"
 
 [deps]
+AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1"
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
@@ -13,17 +14,25 @@ DomainSets = "5b8099bc-c8ec-5219-889f-1d9e522a28bf"
 DynamicQuantities = "06fc5a27-2a28-4c7c-a15d-362465fb6821"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
+OrdinaryDiffEqLowOrderRK = "1344f307-1e59-4825-a18e-ace9aa3fa4c6"
+OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf"
+OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 ThreadsX = "ac1d9e8a-700a-412c-b207-f0111f4b6c0d"
 
 [compat]
+AcceleratedKernels = "0.4.3"
 Accessors = "0.1.42"
 ArrayInterface = "7.18.0"
 DiffEqCallbacks = "4.3"
@@ -36,13 +45,16 @@ JLArrays = "0.2.0"
 LinearSolve = "3.7.2"
 MacroTools = "0.5.15"
 MetaGraphsNext = "0.5, 0.6, 0.7"
+Metal = "1.7.0"
 MethodOfLines = "0.11.7"
 ModelingToolkit = "10.4"
 OrdinaryDiffEqLowOrderRK = "1.2.0"
 OrdinaryDiffEqSDIRK = "1.3.0"
 OrdinaryDiffEqTsit5 = "1.1.0"
+Reactant = "0.2.161"
 RuntimeGeneratedFunctions = "0.5.13"
 SciMLBase = "2.82"
+SciMLOperators = "1.5.0"
 SciMLSensitivity = "7.76.0"
 Statistics = "1.10"
 SymbolicIndexingInterface = "0.3.38"
@@ -58,6 +70,7 @@ DomainSets = "5b8099bc-c8ec-5219-889f-1d9e522a28bf"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 MethodOfLines = "94925ecb-adb7-4558-8ed8-f975c56a0bf4"
 OrdinaryDiffEqLowOrderRK = "1344f307-1e59-4825-a18e-ace9aa3fa4c6"
 OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf"
@@ -69,4 +82,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Documenter", "Test", "SafeTestsets", "ForwardDiff", "JLArrays", "LinearSolve", "MethodOfLines", "Dates", "DomainSets", "OrdinaryDiffEqLowOrderRK", "OrdinaryDiffEqSDIRK", "OrdinaryDiffEqTsit5", "SciMLBase", "SciMLSensitivity", "Zygote"]
+test = ["Documenter", "Test", "SafeTestsets", "ForwardDiff", "Metal", "LinearSolve", "MethodOfLines", "Dates", "DomainSets", "OrdinaryDiffEqLowOrderRK", "OrdinaryDiffEqSDIRK", "OrdinaryDiffEqTsit5", "SciMLBase", "SciMLSensitivity", "Zygote"]

src/EarthSciMLBase.jl

@@ -13,7 +13,11 @@ using LinearAlgebra
 using SymbolicIndexingInterface: setp
 using Accessors
 using Graphs
-using MacroTools, RuntimeGeneratedFunctions, ThreadsX
+using MacroTools, RuntimeGeneratedFunctions
+import SciMLOperators as SMO
+import ThreadsX
+import AcceleratedKernels as AK
+import Reactant
 
 include("add_dims.jl")
 include("domaininfo.jl")
@@ -25,9 +29,11 @@ include("advection.jl")
 include("coord_trans.jl")
 include("param_to_var.jl")
 include("graph.jl")
+include("map_algorithm.jl")
+include("blockdiagonal.jl")
+include("jacobian.jl")
 include("mtk_grid_func.jl")
 include("solver_strategies.jl")
 include("solver_strategy_strang.jl")
-include("blockdiagonal.jl")
 
 end