Merge branch 'task-affinity' of https://github.com/AkhilAkkapelli/Dagger.jl into task-affinity

Author: AkhilAkkapelli

.buildkite/pipeline.yml

@@ -7,6 +7,9 @@
     arch: x86_64
   command: "julia --project -e 'using Pkg; Pkg.develop(;path=\"lib/TimespanLogging\")'"
 
+.gputest: &gputest
+  if: build.message !~ /\[skip tests\]/
+
 .bench: &bench
   if: build.message =~ /\[run benchmarks\]/
   agents:
@@ -28,6 +31,28 @@ steps:
       - JuliaCI/julia-coverage#v1:
           codecov: true
 
+  - label: Julia 1.10
+    timeout_in_minutes: 90
+    <<: *test
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1.10"
+      - JuliaCI/julia-test#v1:
+          julia_args: "--threads=1"
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+
+  - label: Julia 1.11
+    timeout_in_minutes: 90
+    <<: *test
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1.11"
+      - JuliaCI/julia-test#v1:
+          julia_args: "--threads=1"
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+
   - label: Julia 1
     timeout_in_minutes: 90
     <<: *test
@@ -65,6 +90,83 @@ steps:
       - JuliaCI/julia-coverage#v1:
           codecov: true
 
+  - label: Julia 1.11 (CUDA)
+    timeout_in_minutes: 20
+    <<: *gputest
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1.11"
+      - JuliaCI/julia-test#v1: ~
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    env:
+      CI_USE_CUDA: "1"
+
+  - label: Julia 1.11 (ROCm)
+    timeout_in_minutes: 20
+    <<: *gputest
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1.11"
+      - JuliaCI/julia-test#v1: ~
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+    agents:
+      queue: "juliagpu"
+      rocm: "*"
+    env:
+      CI_USE_ROCM: "1"
+
+  - label: Julia 1.11 (oneAPI)
+    timeout_in_minutes: 20
+    <<: *gputest
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1.11"
+      - JuliaCI/julia-test#v1: ~
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+    agents:
+      queue: "juliagpu"
+      intel: "*"
+    env:
+      CI_USE_ONEAPI: "1"
+
+  - label: Julia 1.11 (Metal)
+    timeout_in_minutes: 20
+    <<: *gputest
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1.11"
+      - JuliaCI/julia-test#v1: ~
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+    agents:
+      queue: "juliaecosystem"
+      os: "macos"
+      arch: "aarch64"
+    env:
+      CI_USE_METAL: "1"
+
+  - label: Julia 1.11 (OpenCL)
+    timeout_in_minutes: 20
+    <<: *gputest
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1.11"
+      - JuliaCI/julia-test#v1:
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+    agents:
+      queue: "juliaecosystem"
+      os: macos
+      arch: aarch64
+    env:
+      CI_USE_OPENCL: "1"
+
   - label: Julia 1 - TimespanLogging
     timeout_in_minutes: 20
     <<: *test

Project.toml

@@ -1,13 +1,15 @@
 name = "Dagger"
 uuid = "d58978e5-989f-55fb-8d15-ea34adc7bf54"
-version = "0.18.16"
+version = "0.18.17"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 DistributedNext = "fab6aee4-877b-4bac-a744-3eca44acbb6f"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94"
@@ -28,35 +30,52 @@ TimespanLogging = "a526e669-04d3-4846-9525-c66122c55f63"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [weakdeps]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
+OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2"
+oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 
 [extensions]
+CUDAExt = "CUDA"
 DistributionsExt = "Distributions"
 GraphVizExt = "GraphViz"
 GraphVizSimpleExt = "Colors"
+IntelExt = "oneAPI"
 JSON3Ext = "JSON3"
+MetalExt = "Metal"
+OpenCLExt = "OpenCL"
 PlotsExt = ["DataFrames", "Plots"]
 PythonExt = "PythonCall"
+ROCExt = "AMDGPU"
 
 [compat]
-Adapt = "4.0.4"
+AMDGPU = "1"
+Adapt = "4"
+CUDA = "3, 4, 5"
 Colors = "0.12, 0.13"
 DataFrames = "1"
 DataStructures = "0.18"
 DistributedNext = "1.0.0"
 Distributions = "0.25"
+FillArrays = "1.13.0"
 GraphViz = "0.2"
 Graphs = "1"
 JSON3 = "1"
+KernelAbstractions = "0.9"
 MacroTools = "0.5"
 MemPool = "0.4.12"
+Metal = "1.1"
 OnlineStats = "1"
+OpenCL = "0.10"
+oneAPI = "1, 2"
 Plots = "1"
 PrecompileTools = "1.2"
 Preferences = "1.4.3"

docs/make.jl

@@ -27,7 +27,11 @@ makedocs(;
         "Scopes" => "scopes.md",
         "Processors" => "processors.md",
         "Task Queues" => "task-queues.md",
-        "Datadeps" => "datadeps.md",
+        "Datadeps" => [
+            "Basics" => "datadeps.md",
+            "Stencils" => "stencils.md",
+        ],
+        "GPUs" => "gpu.md",
         "Option Propagation" => "propagation.md",
         "Logging and Visualization" => [
             "Logging: Basics" => "logging.md",

docs/src/data-management.md

@@ -26,6 +26,31 @@ VRAM, etc.) the value resides, and where the value is allowed to be transferred
 and dereferenced. See [Processors](@ref) and [Scopes](@ref) for more details on
 how these properties can be used to control scheduling behavior around `Chunk`s.
 
+## Data movement rules
+
+Dagger utilizes a 3-argument function `Dagger.move(from_proc::Dagger.Processor, to_proc::Dagger.Processor, x)` to manage data movement between processors. This function is invoked by the scheduler for every argument of a task, including the task's function itself, before the task is executed. The purpose of `move` is to transfer the argument `x` from its current processor (`from_proc`) to the target processor (`to_proc`) where the task will run, and to perform any necessary data conversion or unwrapping before execution.
+
+This `move` mechanism is fundamental to how Dagger handles `Chunk` objects. When a `Chunk` is passed as an argument to a task, the `move` function is responsible for unwrapping the `Chunk` and providing its underlying value to the task.
+
+While users can define custom `move` implementations for their specific data types if needed, the default fallback implementation of `move` is designed to handle most common use cases effectively. Therefore, custom implementations are generally unnecessary.
+
+Here's an example of a custom `move` implementation:
+
+```julia
+struct MyCustomType
+    data::Vector{Float64}
+end
+
+# Custom move function for MyCustomType
+function Dagger.move(from_proc::Dagger.Processor, to_proc::Dagger.Processor, x::MyCustomType)
+    return x.data
+end
+
+A = MyCustomType(rand(100))
+s = fetch(Dagger.@spawn sum(A))
+@assert s == sum(A.data)
+```
+
 ## Mutation
 
 Normally, Dagger tasks should be functional and "pure": never mutating their

docs/src/datadeps.md

@@ -179,9 +179,44 @@ Dagger.spawn_datadeps() do
 end
 ```
 
-You can pass any number of aliasing modifiers to `Deps`. This is particularly
-useful for declaring aliasing with `Diagonal`, `Bidiagonal`, `Tridiagonal`, and
-`SymTridiagonal` access, as these "wrappers" make a copy of their parent array
-and thus can't be used to "mask" access to the parent like `UpperTriangular`
-and `UnitLowerTriangular` can (which is valuable for writing memory-efficient,
-generic algorithms in Julia).
+We call `InOut(Diagonal)` an "aliasing modifier". The purpose of `Deps` is to
+pass an argument (here, `A`) as-is, while specifying to Datadeps what portions
+of the argument will be accessed (in this case, the diagonal elements) and how
+(read/write/both). You can pass any number of aliasing modifiers to `Deps`.
+
+`Deps` is particularly useful for declaring aliasing with `Diagonal`,
+`Bidiagonal`, `Tridiagonal`, and `SymTridiagonal` access, as these "wrappers"
+make a copy of their parent array and thus can't be used to "mask" access to the
+parent like `UpperTriangular` and `UnitLowerTriangular` can (which is valuable
+for writing memory-efficient, generic algorithms in Julia).
+
+### Supported Aliasing Modifiers
+
+- Any function that returns the original object or a view of the original object
+- `UpperTriangular`/`LowerTriangular`/`UnitUpperTriangular`/`UnitLowerTriangular`
+- `Diagonal`/`Bidiagonal`/`Tridiagonal`/`SymTridiagonal` (via `Deps`, e.g. to read from the diagonal of `X`: `Dagger.@spawn sum(Deps(X, In(Diagonal)))`)
+- `Symbol` for field access (via `Deps`, e.g. to write to `X.value`: `Dagger.@spawn setindex!(Deps(X, InOut(:value)), :value, 42)`
+
+## In-place data movement rules
+
+Datadeps uses a specialized 5-argument function, `Dagger.move!(dep_mod, from_space::Dagger.MemorySpace, to_space::Dagger.MemorySpace, from, to)`, for managing in-place data movement. This function is an in-place variant of the more general `move` function (see [Data movement rules](@ref)) and is exclusively used within the Datadeps system. The `dep_mod` argument is usually just `identity`, but it can also be an access modifier function like `UpperTriangular`, which limits what portion of the data should be read from and written to.
+
+The core responsibility of `move!` is to read data from the `from` argument and write it directly into the `to` argument. This is crucial for operations that modify data in place, as often encountered in numerical computing and linear algebra.
+
+The default implementation of `move!` handles `Chunk` objects by unwrapping them and then recursively calling `move!` on the underlying values. This ensures that the in-place operation is performed on the actual data.
+
+Users have the option to define their own `move!` implementations for custom data types. However, this is typically not necessary for types that are subtypes of `AbstractArray`, provided that these types support the standard `Base.copyto!(to, from)` function. The default `move!` will leverage `copyto!` for such array types, enabling efficient in-place updates.
+
+Here's an example of a custom `move!` implementation:
+
+```julia
+struct MyCustomArrayWrapper{T,N}
+    data::Array{T,N}
+end
+
+# Custom move! function for MyCustomArrayWrapper
+function Dagger.move!(dep_mod::Any, from_space::Dagger.MemorySpace, to_space::Dagger.MemorySpace, from::MyCustomArrayWrapper, to::MyCustomArrayWrapper)
+    copyto!(dep_mod(to.data), dep_mod(from.data))
+    return
+end
+```