Merge branch 'main' into fe/return

Author: vchuravy

.github/workflows/Formatting.yml

@@ -0,0 +1,55 @@
+name: 'Format'
+
+on:
+  pull_request_target:
+    paths: ['**/*.jl']
+    types: [opened, synchronize, reopened, ready_for_review]
+
+permissions:
+  contents: read
+
+  # needed for julia-actions/cache to delete old caches
+  actions: write
+
+  # needed for googleapis/code-suggester
+  pull-requests: write
+
+jobs:
+  runic:
+    runs-on: ubuntu-latest
+    if: github.event.pull_request.draft == false
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{github.event.pull_request.head.ref}}
+          repository: ${{github.event.pull_request.head.repo.full_name}}
+          fetch-depth: 0
+
+      - name: Setup Julia
+        uses: julia-actions/setup-julia@v2
+        with:
+          version: '1'
+          arch: 'x64'
+      - uses: julia-actions/cache@v2
+
+      - name: Install Runic
+        run: |
+          julia --project=@runic -e 'using Pkg; Pkg.add("Runic")'
+          curl -o git-runic https://raw.githubusercontent.com/fredrikekre/Runic.jl/master/bin/git-runic
+          chmod +x git-runic
+          sudo mv git-runic /usr/local/bin
+
+      - name: Run Runic
+        run: |
+          set +e
+          git runic origin/main
+          [ $? -eq 2 ] && exit 1 || exit 0
+
+      - name: Suggest changes
+        uses: googleapis/code-suggester@v2
+        env:
+          ACCESS_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        with:
+          command: review
+          pull_number: ${{ github.event.pull_request.number }}
+          git_dir: '.'

Project.toml

@@ -1,7 +1,7 @@
 name = "KernelAbstractions"
 uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 authors = ["Valentin Churavy <[email protected]> and contributors"]
-version = "0.9.29"
+version = "0.9.32"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -15,12 +15,10 @@ Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
-UnsafeAtomics = "013be700-e6cd-48c3-b4a1-df204f14c38f"
-UnsafeAtomicsLLVM = "d80eeb9a-aca5-4d75-85e5-170c8b632249"
 
 [compat]
 Adapt = "0.4, 1.0, 2.0, 3.0, 4"
-Atomix = "0.1"
+Atomix = "0.1, 1"
 EnzymeCore = "0.7, 0.8.1"
 InteractiveUtils = "1.6"
 LinearAlgebra = "1.6"
@@ -30,8 +28,6 @@ Requires = "1.3"
 SparseArrays = "<0.0.1, 1.6"
 StaticArrays = "0.12, 1.0"
 UUIDs = "<0.0.1, 1.6"
-UnsafeAtomics = "0.2.1"
-UnsafeAtomicsLLVM = "0.1, 0.2"
 julia = "1.6"
 
 [extensions]

docs/src/quickstart.md

@@ -45,7 +45,7 @@ The [`synchronize`](@ref) blocks the *host* until the kernel has completed on th
 
 ## Launching kernel on the backend
 
-To launch the kernel on a backend-supported backend `isa(backend, KA.GPU)` (e.g., `CUDABackend()`, `ROCBackend()`, `oneBackend()`), we generate the kernel
+To launch the kernel on a backend-supported backend `isa(backend, KA.GPU)` (e.g., `CUDABackend()`, `ROCBackend()`, `oneAPIBackend()`), we generate the kernel
 for this backend.
 
 First, we initialize the array using the Array constructor of the chosen backend with

src/KernelAbstractions.jl

@@ -10,7 +10,6 @@ export synchronize, get_backend, allocate
 import PrecompileTools
 
 import Atomix: @atomic, @atomicswap, @atomicreplace
-import UnsafeAtomics
 
 using MacroTools
 using StaticArrays
@@ -111,13 +110,53 @@ macro Const end
 """
     copyto!(::Backend, dest::AbstractArray, src::AbstractArray)
 
-Perform a `copyto!` operation that execution ordered with respect to the backend.
+Perform an asynchronous `copyto!` operation that is execution ordered with respect to the back-end.
+
+For most users, `Base.copyto!` should suffice, performance a simple, synchronous copy.
+Only when you know you need asynchronicity w.r.t. the host, you should consider using
+this asynchronous version, which requires additional lifetime guarantees as documented below.
+
+!!! warning
+
+    Because of the asynchronous nature of this operation, the user is required to guarantee that the lifetime
+    of the source extends past the *completion* of the copy operation as to avoid a use-after-free. It is not
+    sufficient to simply use `GC.@preserve` around the call to `copyto!`, because that only extends the
+    lifetime past the operation getting queued. Instead, it may be required to `synchronize()`,
+    or otherwise guarantee that the source will still be around when the copy is executed:
+
+    ```julia
+    arr = zeros(64)
+    GC.@preserve arr begin
+        copyto!(backend, arr, ...)
+        # other operations
+        synchronize(backend)
+    end
+    ```
 
 !!! note
-    Backend implementations **must** implement this function.
+
+    On some back-ends it may be necessary to first call [`pagelock!`](@ref) on host memory
+    to enable fully asynchronous behavior w.r.t to the host.
+
+!!! note
+    Backends **must** implement this function.
 """
 function copyto! end
 
+"""
+    pagelock!(::Backend, dest::AbstractArray)
+
+Pagelock (pin) a host memory buffer for a backend device. This may be necessary for [`copyto!`](@ref)
+to perform asynchronously w.r.t to the host/
+
+This function should return `nothing`; or `missing` if not implemented.
+
+
+!!! note
+    Backends **may** implement this function.
+"""
+function pagelock! end
+
 """
     synchronize(::Backend)
 
@@ -547,6 +586,34 @@ function priority!(::Backend, prio::Symbol)
     return nothing
 end
 
+"""
+    device(::Backend)::Int
+
+Returns the ordinal number of the currently active device starting at one.
+"""
+function device(::Backend)
+    return 1
+end
+
+"""
+    ndevices(::Backend)::Int
+
+Returns the number of devices the backend supports.
+"""
+function ndevices(::Backend)
+    return 1
+end
+
+"""
+    device!(::Backend, id::Int)
+"""
+function device!(backend::Backend, id::Int)
+    if !(0 < id <= ndevices(backend))
+        throw(ArgumentError("Device id $id out of bounds."))
+    end
+    return nothing
+end
+
 """
     functional(::Backend)
 
@@ -563,6 +630,10 @@ function functional(::Backend)
     return missing
 end
 
+function pagelock!(::Backend, x)
+    return missing
+end
+
 include("nditeration.jl")
 using .NDIteration
 import .NDIteration: get

src/compiler.jl

@@ -30,3 +30,4 @@ end
 
 @inline groupsize(ctx) = __groupsize(ctx)
 @inline ndrange(ctx) = __ndrange(ctx)
+@inline Base.ndims(ctx) = ndims(__iterspace(ctx))

src/cpu.jl

@@ -1,5 +1,3 @@
-import UnsafeAtomicsLLVM
-
 unsafe_free!(::AbstractArray) = return
 synchronize(::CPU) = nothing
 
@@ -35,6 +33,7 @@ function copyto!(backend::CPU, A, B)
 end
 
 functional(::CPU) = true
+pagelock!(::CPU, x) = nothing
 
 function (obj::Kernel{CPU})(args...; ndrange = nothing, workgroupsize = nothing)
     ndrange, workgroupsize, iterspace, dynamic = launch_config(obj, ndrange, workgroupsize)
@@ -43,7 +42,8 @@ function (obj::Kernel{CPU})(args...; ndrange = nothing, workgroupsize = nothing)
         return nothing
     end
 
-    return __run(obj, ndrange, iterspace, args, dynamic, obj.backend.static)
+    __run(obj, ndrange, iterspace, args, dynamic, obj.backend.static)
+    return nothing
 end
 
 const CPU_GRAINSIZE = 1024 # Vectorization, 4x unrolling, minimal grain size

src/nditeration.jl

@@ -63,6 +63,7 @@ end
 @inline workitems(range::NDRange{N, B, W}) where {N, B, W <: StaticSize} = CartesianIndices(get(W))::CartesianIndices{N}
 @inline blocks(range::NDRange{N, B}) where {N, B <: DynamicSize} = range.blocks::CartesianIndices{N}
 @inline blocks(range::NDRange{N, B}) where {N, B <: StaticSize} = CartesianIndices(get(B))::CartesianIndices{N}
+@inline Base.ndims(::NDRange{N}) where {N} = N
 
 import Base.iterate
 @inline iterate(range::NDRange) = iterate(blocks(range))
@@ -80,8 +81,44 @@ Base.length(range::NDRange) = length(blocks(range))
     return CartesianIndex(nI)
 end
 
+
+"""
+    assume(cond::Bool)
+
+Assume that the condition `cond` is true. This is a hint to the compiler, possibly enabling
+it to optimize more aggressively.
+"""
+@inline assume(cond::Bool) = Base.llvmcall(
+    (
+        """
+        declare void @llvm.assume(i1)
+
+        define void @entry(i8) #0 {
+            %cond = icmp eq i8 %0, 1
+            call void @llvm.assume(i1 %cond)
+            ret void
+        }
+
+        attributes #0 = { alwaysinline }""", "entry",
+    ),
+    Nothing, Tuple{Bool}, cond
+)
+
+@inline function assume_nonzero(CI::CartesianIndices)
+    return ntuple(Val(ndims(CI))) do I
+        Base.@_inline_meta
+        indices = CI.indices[I]
+        assume(indices.stop > 0)
+    end
+end
+
 Base.@propagate_inbounds function expand(ndrange::NDRange, groupidx::Integer, idx::Integer)
-    return expand(ndrange, blocks(ndrange)[groupidx], workitems(ndrange)[idx])
+    # this causes a exception branch and a div
+    B = blocks(ndrange)
+    W = workitems(ndrange)
+    assume_nonzero(B)
+    assume_nonzero(W)
+    return expand(ndrange, B[groupidx], workitems(ndrange)[idx])
 end
 
 Base.@propagate_inbounds function expand(ndrange::NDRange{N}, groupidx::CartesianIndex{N}, idx::Integer) where {N}

test/compiler.jl

@@ -41,6 +41,7 @@ function compiler_testsuite(backend, ArrayT)
     kernel = index(CPU(), DynamicSize(), DynamicSize())
     iterspace = NDRange{1, StaticSize{(128,)}, StaticSize{(8,)}}()
     ctx = KernelAbstractions.mkcontext(kernel, 1, nothing, iterspace, Val(KernelAbstractions.NoDynamicCheck()))
+    @test ndims(ctx) == 1
     @test KernelAbstractions.__index_Global_NTuple(ctx, CartesianIndex(1)) == (1,)
 
     A = ArrayT{Int}(undef, 1)

test/devices.jl

@@ -0,0 +1,14 @@
+function devices_testsuite(Backend)
+    backend = Backend()
+
+    current_device = KernelAbstractions.device(backend)
+    for i in KernelAbstractions.ndevices(backend)
+        KernelAbstractions.device!(backend, i)
+        @test KernelAbstractions.device(backend) == i
+    end
+
+    @test_throws ArgumentError KernelAbstractions.device!(backend, 0)
+    @test_throws ArgumentError KernelAbstractions.device!(backend, KernelAbstractions.ndevices(backend) + 1)
+    KernelAbstractions.device!(backend, current_device)
+    return nothing
+end

test/nditeration.jl

@@ -7,10 +7,12 @@ function nditeration_testsuite()
         let ndrange = NDRange{2, DynamicSize, DynamicSize}(CartesianIndices((256, 256)), CartesianIndices((32, 32)))
             @test length(ndrange) == 256 * 256
             @test all(p -> p[1] == p[2], zip(ndrange, CartesianIndices((256, 256))))
+            @test ndims(ndrange) == 2
         end
         let ndrange = NDRange{2, StaticSize{(256, 256)}, DynamicSize}(nothing, CartesianIndices((32, 32)))
             @test length(ndrange) == 256 * 256
             @test all(p -> p[1] == p[2], zip(ndrange, CartesianIndices((256, 256))))
+            @test ndims(ndrange) == 2
         end
     end
 
@@ -40,13 +42,15 @@ function nditeration_testsuite()
                 I = Tuple(I)
                 @test check(idx, i - 1, ntuple(i -> I[i] - 1, length(I))..., Dim_x, Dim_y)
             end
+            @test ndims(ndrange) == 2
         end
         let ndrange = NDRange{2, DynamicSize, DynamicSize}(CartesianIndices((4, 4)), CartesianIndices((Dim_x, Dim_y)))
             idx = linear_iteration(ndrange)
             for (i, I) in zip(1:length(blocks(ndrange)), blocks(ndrange))
                 I = Tuple(I)
                 @test check(idx, i - 1, ntuple(i -> I[i] - 1, length(I))..., Dim_x, Dim_y)
             end
+            @test ndims(ndrange) == 2
         end
 
         Dim_x = 32
@@ -58,13 +62,15 @@ function nditeration_testsuite()
                 I = Tuple(I)
                 @test check(idx, i - 1, ntuple(i -> I[i] - 1, length(I))..., Dim_x, Dim_y)
             end
+            @test ndims(ndrange) == 2
         end
         let ndrange = NDRange{2, DynamicSize, DynamicSize}(CartesianIndices((4, 4 * 32)), CartesianIndices((Dim_x, Dim_y)))
             idx = linear_iteration(ndrange)
             for (i, I) in zip(1:length(blocks(ndrange)), blocks(ndrange))
                 I = Tuple(I)
                 @test check(idx, i - 1, ntuple(i -> I[i] - 1, length(I))..., Dim_x, Dim_y)
             end
+            @test ndims(ndrange) == 2
         end
 
         Dim_x = 1
@@ -76,13 +82,15 @@ function nditeration_testsuite()
                 I = Tuple(I)
                 @test check(idx, i - 1, ntuple(i -> I[i] - 1, length(I))..., Dim_x, Dim_y)
             end
+            @test ndims(ndrange) == 2
         end
         let ndrange = NDRange{2, DynamicSize, DynamicSize}(CartesianIndices((4 * 32, 4)), CartesianIndices((Dim_x, Dim_y)))
             idx = linear_iteration(ndrange)
             for (i, I) in zip(1:length(blocks(ndrange)), blocks(ndrange))
                 I = Tuple(I)
                 @test check(idx, i - 1, ntuple(i -> I[i] - 1, length(I))..., Dim_x, Dim_y)
             end
+            @test ndims(ndrange) == 2
         end
     end
     return