Define alias detection for every GPU array through a location hook

NEWS.md

@@ -1,10 +1,21 @@
 # GPUArrays.jl release history
 
-This document lists the noteworthy changes in every breaking release of GPUArrays.jl,
+This document lists the noteworthy changes in GPUArrays.jl releases that back-ends act on,
 newest first, aimed at the authors of back-end packages. For the complete list of merged
 pull requests, see the [GitHub releases](https://github.com/JuliaGPU/GPUArrays.jl/releases).
 
 
+## v12.1
+
+GPUArrays defines `Base.dataids` and `Base.mightalias` for every `AbstractGPUArray`. A
+back-end that implements the new `GPUArrays.memory_location` hook (the coordinate of an
+array's allocation, normally its address, and its byte offset into it) gets precise alias
+detection: disjoint views of one buffer don't alias, empty arrays alias nothing on every
+Julia version, and `SubArray`s of two such arrays are compared without converting their
+parents to pointers. Arrays wrapped in a `SubArray` or other wrapper are still compared by
+allocation, which is conservative. Back-ends should implement the hook and delete their own
+`dataids` and `mightalias` methods, which would keep winning dispatch.
+
 ## v12.0
 
 Base's reductions, sorting, scans, `findall`, logical indexing and `reverse` are now

Project.toml

@@ -1,6 +1,6 @@
 name = "GPUArrays"
 uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
-version = "12.0.3"
+version = "12.1.0"
 
 [workspace]
 projects = ["lib/GPUArraysCore", "lib/JLArrays", "test", "docs"]

docs/src/interface.md

@@ -31,6 +31,17 @@ KernelAbstractions.get_backend(a::CA) where CA <: CustomArray = CustomBackend()
 
 There are numerous examples of potential interfaces for GPUArrays, such as with [JLArrays](https://github.com/JuliaGPU/GPUArrays.jl/blob/main/lib/JLArrays/src/JLArrays.jl), [CuArrays](https://github.com/JuliaGPU/CUDA.jl/blob/main/CUDACore/src/array.jl), and [ROCArrays](https://github.com/JuliaGPU/AMDGPU.jl/blob/main/src/array.jl).
 
+## Alias detection
+
+Without the hook below, all views of one buffer are treated as aliasing each other, so Base
+and GPUArrays copy one of them before operations that write to the other (e.g. `copyto!`,
+broadcasting, `accumulate!`). Report where an array's elements are stored to make that check
+precise:
+
+```@docs
+GPUArrays.memory_location
+```
+
 ## Caching Allocator
 
 ```@docs

lib/JLArrays/Project.toml

@@ -1,6 +1,6 @@
 name = "JLArrays"
 uuid = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
-version = "0.4.1"
+version = "0.4.2"
 authors = ["Tim Besard <tim.besard@gmail.com>"]
 
 [deps]
@@ -13,7 +13,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
 Adapt = "2.0, 3.0, 4.0"
-GPUArrays = "12"
+GPUArrays = "12.1"
 KernelAbstractions = "0.9"
 LinearAlgebra = "1"
 Random = "1"

lib/JLArrays/src/array.jl

@@ -45,19 +45,7 @@ end
 
 GPUArrays.storage(a::JLArray) = a.data
 
-## alias detection
-
-# Byte-range overlap follows CUDA's Base.mightalias(::CuArray, ::CuArray):
-# https://github.com/JuliaGPU/CUDA.jl/blob/main/CUDACore/src/array.jl#L113-L116
-# Keep shared-allocation dataids for conservative alias detection involving wrappers.
-function Base.mightalias(A::JLArray, B::JLArray)
-    nA, nB = sizeof(A), sizeof(B)
-    (iszero(nA) || iszero(nB)) && return false
-    GC.@preserve A B begin
-        pA, pB = UInt(pointer(A)), UInt(pointer(B))
-        return pA <= pB < pA + nA || pB <= pA < pB + nB
-    end
-end
+GPUArrays.memory_location(a::JLArray) = (UInt(pointer(a.data[])), a.offset)
 
 GPUArrays.dense_array_type(a::JLArray{T, N}) where {T, N} = JLArray{T, N}
 GPUArrays.dense_array_type(::Type{JLArray{T, N}}) where {T, N} = JLArray{T, N}

src/host/abstractarray.jl

@@ -101,14 +101,96 @@ function unsafe_free!(ref::DataRef)
     return
 end
 
+
+## alias detection
+
+"""
+    GPUArrays.memory_location(A::AbstractGPUArray) -> Union{Nothing, Tuple{UInt, Int}}
+
+Where the elements of `A` are stored: a coordinate `base` that identifies the allocation
+holding them (normally its address), and the byte `offset` of the first element from it.
+GPUArrays defines `Base.dataids`, `Base.mightalias` and the check Base uses to compare two
+`SubArray`s of the same parent for every `AbstractGPUArray`; with this hook, they compare the
+byte ranges arrays cover instead of the allocations they share.
+
+Two arrays are taken to share memory when their byte ranges, `sizeof(A)` bytes from
+`base + offset`, overlap. For that to be sound:
+
+- Every array derived from an allocation (a view, reshape or reinterpretation, which are
+  built with `derive`) must report the same `base`, and a nonnegative `offset` into it.
+  Memory wrapped from a pointer into the middle of an allocation (`unsafe_wrap`) reports a
+  `base` of its own: it is still compared by byte range against arrays that implement this
+  hook, but against wrapped arrays (e.g. a `SubArray`) and arrays of other types (e.g. an
+  `Array`), only when both start at the same address.
+- Arrays that share memory must report it in the same coordinates. That holds for
+  addresses in a single address space, but not for memory that is reachable at two
+  addresses (such as a host buffer and its device mapping) or that is wrapped twice with
+  separate identities (such as two handles for one buffer): the hook cannot express such
+  aliases, so a back-end that creates them has to detect them itself.
+- Memory without an address may be identified by any other value that is unique to its
+  allocation, such as a handle. Overlap with the coordinates of other memory only causes
+  false positives.
+- For valid, live arrays it must not have side effects (taking stream ownership,
+  synchronizing) and must not throw, as Base checks for aliasing in many operations,
+  including `copyto!` and broadcasting.
+
+The default, `nothing`, keeps alias detection at the level of allocations: arrays alias
+when they share the `DataRef` returned by `GPUArrays.storage`.
+"""
+memory_location(::AbstractGPUArray) = nothing
+
 # `copy(::DataRef)` shares `rc`. Views and reshapes are built that way, so
 # alias detection has to key off the shared record rather than the wrapper.
 Base.dataids(ref::DataRef) = (objectid(ref.rc),)
 
+# Base compares `dataids` whenever one of the arrays is wrapped (e.g., a `SubArray`), so they
+# identify the allocation rather than the part of it the array covers: the shared `DataRef`,
+# which arrays of a type without `memory_location` are compared by, and the allocation's
+# coordinate. The start of the array is included too, to match an array wrapped around
+# memory starting at the same address (whose `dataids` are that address). Empty arrays alias
+# nothing, which Julia 1.10's `mightalias` does not check.
 function Base.dataids(A::AbstractGPUArray)
+    isempty(A) && return ()
     ids = Base.dataids(storage(A))
-    isempty(ids) ? (objectid(A),) : ids
+    loc = memory_location(A)
+    if loc === nothing
+        return isempty(ids) ? (objectid(A),) : ids
+    end
+    base, offset = loc
+    return (ids..., base, base + offset)
+end
+
+function Base.mightalias(A::AbstractGPUArray, B::AbstractGPUArray)
+    (isempty(A) || isempty(B)) && return false
+    a, b = memory_location(A), memory_location(B)
+    # the selector bytes of isbits-union arrays are stored outside the elements' byte range
+    if a === nothing || b === nothing ||
+       Base.isbitsunion(eltype(A)) || Base.isbitsunion(eltype(B))
+        return !Base._isdisjoint(Base.dataids(A), Base.dataids(B))
+    end
+    # nonempty arrays without bytes (of a zero-size element type) can't be written to
+    nA, nB = sizeof(A), sizeof(B)
+    (iszero(nA) || iszero(nB)) && return false
+    # widened, so that coordinates near `typemax(UInt)` (e.g. handles) don't wrap around
+    startA, startB = widen(a[1]) + a[2], widen(b[1]) + b[2]
+    return startA <= startB ? startB - startA < nA : startA - startB < nB
+end
+
+# Base only compares the indices of two `SubArray`s when their parents are the same memory,
+# which it checks for dense arrays by comparing pointers. Converting a GPU array to a pointer
+# can have side effects or throw, so compare their locations instead, or only their identity
+# for arrays without one. Locations don't establish where the selector bytes of isbits-union
+# arrays are, so those only match themselves too. When parents don't match, Base compares
+# their `dataids`, so returning false is always sound.
+function Base._parentsmatch(A::AbstractGPUArray, B::AbstractGPUArray)
+    A === B && return true
+    a, b = memory_location(A), memory_location(B)
+    (a === nothing || b === nothing) && return false
+    (Base.isbitsunion(eltype(A)) || Base.isbitsunion(eltype(B))) && return false
+    return a == b && Base.elsize(A) == Base.elsize(B) && size(A) == size(B)
 end
+Base._parentsmatch(::AbstractGPUArray, ::DenseArray) = false
+Base._parentsmatch(::DenseArray, ::AbstractGPUArray) = false
 
 # array methods
 

test/testsuite/aliasing.jl

@@ -1,3 +1,4 @@
+
 @testsuite "aliasing" (AT, eltypes)->begin
     T = Float32 in eltypes ? Float32 : first(eltypes)
     buffer = AT{T}(undef, 32)
@@ -15,4 +16,17 @@
     empty = view(buffer, 9:8)
     @test !Base.mightalias(empty, left)
     @test !Base.mightalias(left, empty)
+
+    # Wrapped arrays are compared by allocation, which conservatively covers any overlap.
+    @test Base.mightalias(view(overlap, 1:2:16), left)
+    @test Base.mightalias(left, view(overlap, 1:2:16))
+    @test Base.mightalias(view(left, 1:2:16), view(overlap, 1:2:16))
+    @test !Base.mightalias(view(left, 1:2:16), view(copy(left), 1:2:16))
+
+    # Strided views of the same memory are compared by index, even if their parents are
+    # different array objects.
+    other = view(buffer, 1:16)
+    @test !Base.mightalias(view(left, 1:2:16), view(other, 2:2:16))
+    @test Base.mightalias(view(left, 1:2:16), view(other, 3:2:16))
+    @test Base.mightalias(view(left, 1:2:16), view(reshape(other, 4, 4), 1, :))
 end