Repository navigation
Give the workspaces of sorts and scans along dims a concrete type - #171
Merged
Merged
Conversation
Contributor
Author
|
How can I get this merged? |
Member
|
I'll rebase this. |
maleadt
force-pushed
the
ab/infer-merge-sort-workspace
branch
from
October 26, 2026 08:25
e05dcbd to
7aae07e
Compare
maleadt
changed the title
dims a concrete type
Member
|
Rebased and slightly reworked. |
Merge sort's plans listed different buffers depending on the length (none for an input of one element, swap buffers only where tiles are merged) and, for `sortperm!`, on `lowmem`. Inference widens these unions, so the workspaces of `sort!`, `sortperm!`, `sortperm` and `sort_by_key!` had abstract types, and a caller that passed one to the operation, or used its buffers, dispatched at run time. Merge sort now lists the same buffers for every length, without elements where the call does not need them, as the whole-array reductions already do with their partials. The low-memory `sortperm!` path lists the same buffers as the default one, with empty keys, and swaps the indices through `temp_values`, so that the plan does not depend on `lowmem` either; only keys of a union type, which some backends cannot allocate even empty, are left out. Zero-dimensional arrays, which have no empty buffers of their shape, keep needing none. The device memory each call allocates is unchanged.
`select_sort` resolves `Auto()` to `MergeSort()` for a permutation or a key/value sort on a GPU, but `perm` and `pairs` did not reach it as constants: the algorithm inferred as the union of the three kernel algorithms, and so did the plans and workspaces of `sortperm!` and `sort_by_key!`. Constant-propagating `resolve_sort` aggressively lets them infer as merge sort's (or, on the host, as `CPUThreads.SampleSort`'s).
Radix sort's plan listed no buffers for slices of at most one element or slices sorted in one block, and its buffers and nested plans otherwise, so `workspace(sort!, v; alg=RadixSort())` inferred as an abstract `Workspace`. It now lists the same buffers and nested plans for every length, without elements where the call does not need them, with `no_elements` emptying a whole plan.
`SliceScan` listed its chunk totals only where a slice splits into chunks, so the plans and workspaces of `accumulate` along `dims` inferred as abstract. It now lists them for every shape, without elements for slices that fit in one chunk.
maleadt
force-pushed
the
ab/infer-merge-sort-workspace
branch
from
October 26, 2026 10:06
7aae07e to
6f38584
Compare
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.
Fixes #170.
The
WorkspacethatAK.workspacereturned for merge sort's operations (sort!,sortperm!,sortperm,sort_by_key!), forRadixSort, and foraccumulatealongdimshad an abstract type. A caller that kept its scratch in a workspace then dispatched at run time every time it used it: the call that took it (sortperm!(...; workspace=ws)), and any use of its buffers.The plans listed different buffers depending on the length:
sortperm!different buffers withlowmem;SliceScan: chunk totals only where a slice splits into chunks.Inference widened these unions.
Each of these plans now lists the same buffers (and nested plans) for every length, without elements where the call doesn't need them (
scratch(T, dims; needed),no_elements). The whole-array reductions already did this with theirpartials. The low-memorysortperm!path lists the same buffers as the default one, with empty keys, and swaps the indices throughtemp_values. The kernels touch the same memory as before, and each call allocates the same device memory.With
Auto(),sortperm!andsort_by_key!also inferred their algorithm as a union of the three kernel sorts, even thoughselect_sortalways picksMergeSortfor them.resolve_sortis now constant-propagated aggressively so thatpermandpairsreach it as constants.Every merge sort, radix sort and
accumulate(; dims)workspace now infers concretely, includingsortperm!with a run-timelowmem::Bool.@report_optreports nothing for the example in #170.Costs:
sortperm!workspace is now tied to the key type.Not covered:
sort!with the defaultAuto()on a GPU picks bitonic, merge or radix sort at run time, so its workspace stays a union. A caller that needs a concrete workspace passes an explicit algorithm.dimspick their launch shape at run time, so their workspaces stay abstract too.