Repository navigation
Support arrays with offset axes, and fix other indexing bugs - #177
Merged
Merged
Conversation
A TaskPartitioner that ends up with one task, e.g. on a single thread or with fewer than twice `min_elems` elements, stores no tasks, so passing it to task_partition or itask_partition threw a BoundsError. Call `f` on the whole range instead.
reverse!(dst, src) offset the destination index by the source's `eachindex`, which throws for sources with Cartesian indices such as a view of a matrix.
The ConcurrentWrite kernel and the threaded algorithm read `v[i]` for `i` in `1:length(v)`, which reads out of bounds for arrays whose axes do not start at one. The threaded algorithm also set its result flag from several tasks without synchronization; make it atomic.
The threaded reductions index their source from 1, so arrays whose axes start elsewhere (an `OffsetArray`, for instance) reduced the wrong elements on the host, whole and along `dims`, also as the source of `mapreducedim!`. On GPUs, a one-element source was read out of bounds, also by `any`/`all` with `ViaReduce`. Reductions now pass the kernels and threaded loops the source and destination with axes starting at 1: a reshape of arrays of several dimensions, whose linear indexing already starts at 1 (and which `OffsetArrays` makes the parent array), and a view of vectors. Fills of such views on GPUs run as a kernel, as Base's `fill!` would index them element by element. The tests get a minimal array type with offset axes, which replaces the one of the reduction tests.
The threaded whole-array scan indexes its array from 1, so it read out of bounds for arrays whose axes start elsewhere, and the GPU scans wrote the wrong elements. `accumulate` could not scan such arrays at all, as the result it allocates has axes starting at 1, which `accumulate!` rejected for an offset source. As for reductions, the scans now run on views of the destination and source with axes starting at 1. Copies into such views from the scratch array, and fills of them, run as kernels on GPUs.
The sorts index their arrays from 1, so arrays whose axes start elsewhere threw a `BoundsError` on the host and were read and written out of bounds by the GPU kernels; along `dims` they were rejected. The sorts now run on views of their arrays with axes starting at 1, as reductions and scans do. `sortperm!` and `sortperm` shift the permutation into the source's own linear indices, as Base's do, and `sortperm` along `dims` no longer rejects its own result, whose axes start at 1. The GPU merge sorts initialize the permutation with linear indices also for Cartesian-indexed index arrays.
The mirror index of an element along `dims` was computed from 1-based coordinates, which indexed arrays whose axes start elsewhere out of bounds. Reverse 1-based views of them instead, as the other operations do.
The batched searches only require `ix` and `xs` to have the same length, but wrote `ix[i]` for `i in eachindex(xs)`, out of bounds when either has offset axes.
Before Julia 1.12, reductions of several arrays or of a `Broadcasted` object first materialize it, as it cannot be indexed linearly. That used `copyto!(similar(bc), bc)`, whose broadcast style comes from the arrays: for array types without a GPU broadcast style, such as lazy wrappers or arrays with offset axes around GPU arrays, it allocated a host array and indexed the GPU arrays element by element, which GPU backends disallow. Evaluate it with a kernel on the reduction's backend instead, through Cartesian indices.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Bugs found while auditing the package for a cleanup. The main one: most operations index arrays from 1, so arrays whose axes start elsewhere (e.g. an
OffsetArraywrapping aCuArray, as Oceananigans uses) gave wrong results or were read and written out of bounds:dims, andmapreducedim!from such a source) reduced the wrong elements, becauseAutotakes the threaded path, which reducessrc[irange]with 1-based ranges, even for small inputs;BoundsErroron the host and wrote the wrong elements on GPUs;sort!/sortperm!threw on the host and indexed out of bounds in GPU kernels; alongdimsthey were rejected;any/all,reverse!alongdims, andsearchsorted*!with offset queries or results were read or written out of bounds.Rather than teaching every kernel about offsets, the operations now hand their kernels and threaded loops 1-based versions of the arrays at the host entry: a reshape for arrays of several dimensions (whose linear indexing already starts at 1, and which
OffsetArraysturns into the parent array, so anOffsetArray{CuArray}reaches the kernels as a plainCuArray), a view for vectors, and likewise for the arrays in aBroadcastedsource. The checks (shapes, aliasing, workspaces) still see the caller's arrays.sortperm!returns indices of the source, asBase.sortpermdoes. Copies and fills into such views run as kernels on GPUs, since Base's would index element by element. Arrays whose axes start at 1 take the same paths as before.Smaller fixes:
task_partition(f, tp)anditask_partition(f, tp)threw aBoundsErrorwhen theTaskPartitionerhad a single task (e.g. on one thread), as such a partitioner stores no tasks.reverse!(dst, src)threw aMethodErrorfor Cartesian-indexed sources such asview(A, 1:2, 1:3).any/allset their result from several tasks through a plainRef; it is now atomic.Broadcastedobject materialize it first, withcopyto!(similar(bc), bc). For arrays without a GPU broadcast style (lazy wrappers, or offset arrays around GPU arrays) that indexed the GPU arrays element by element, which GPU backends disallow. A kernel materializes it now.Each fix comes with tests that fail without it, using a minimal offset-axes array type in the test setup. Tested on CPU threads, PoCL (also with AK's kernels on the host backend) and CUDA.