Towards an Automatic GPU Batching Interface (AKA GPU go brrr)

GeneralisedFilters/Project.toml

@@ -6,33 +6,40 @@ authors = ["THargreaves <tim.hargreaves@icloud.com>", "Charles Knipp <charleskni
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1"
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+IRTools = "7869d1d1-7146-5819-86e3-90919afe41df"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
+Magma = "a4173727-5e3e-4567-b12d-2e3cf2fa2f28"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SSMProblems = "26aad666-b158-4e64-9d35-0e672562fa48"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 
 [compat]
 AbstractMCMC = "5"
 AcceleratedKernels = "0.3, 0.4"
+Adapt = "4.5.0"
 Aqua = "0.8"
 CUDA = "5"
 DataStructures = "0.18.20, 0.19"
 Distributions = "0.25"
+IRTools = "0.4.15"
 LogExpFunctions = "0.3"
 OffsetArrays = "1.14.1"
 PDMats = "0.11.35"
 SSMProblems = "0.6"
 StaticArrays = "1.9.16"
 Statistics = "1"
 StatsBase = "0.34.3"
+StructArrays = "0.7.2"
 Test = "1"
 julia = "1.10"
 

GeneralisedFilters/src/GeneralisedFilters.jl

@@ -11,11 +11,19 @@ using StatsBase
 # TODO: heavy modules—move to extension
 using CUDA
 
+export initialise, step, predict, update, filter
+
 # Filtering utilities
 include("callbacks.jl")
 include("containers.jl")
 include("resamplers.jl")
 
+# Batching utilities
+include("batching/types.jl")
+include("batching/broadcasting.jl")
+include("batching/wrappers.jl")
+include("batching/operations.jl")
+
 ## FILTERING BASE ##########################################################################
 
 abstract type AbstractFilter <: AbstractSampler end