Skip to content

A safe way to give a fork tree's leaves disjoint parts of one Array to write (the counterpart of #1376) #1777

Open

Description

@nuxyel

What Bend should do

Base already has the right shape in its model: an Array<T> is ALeaf{value} or ANode{xs, ys}, two disjoint halves. In the compiled program, though, matching an array copies each half into its own block (BendRT, "Arrays"), so a fork tree cannot use it. Something like

# split a into its two halves, both views of a's block: O(1), no copy
Array.halves(-T: Type, a: Array<T>) -> Array<T> & Array<T>      # model: match a { ANode{xs, ys}: (xs, ys) }
# put two halves of one block back together: O(1) when they came from Array.halves
Array.node(-T: Type, xs: Array<T>, ys: Array<T>) -> Array<T>     # model: ANode{xs, ys}

would have a model in which each leaf owns its half (writes included), so laws about the joined result could be stated and proved, and the runtime would do what the unsafe pattern does today. A leaf at depth d owns a contiguous range of 2^(capacity - d) cells. That fits tensor-fork (blocks of rows of C), stencil3d (slabs) and the histogram's rows. It does not fit a 2-D tile directly (tensor-gpu's leaves write 4 x 4 tiles of a row-major C): such a kernel would store C tile-major and permute it afterwards, or need a split by stride, which is a further question.

Questions:

  1. Is a safe split for writes on the roadmap, alongside the read-only borrow of A read-only borrow of an Array for a parallel let, without @unsafe #1376 (or the borrow inference Taelin mentioned)?
  2. Would it belong in Base (as above), or is a library the intended place, with @unsafe contained in it as bend-ml does now?
  3. Is the halves-of-a-node shape the right one for the runtime (a view into one block, freed once), or would ranges (Array.split_at(a, i)) fit it better?

I can share the packages' tests (every product against the sequential code, bit for bit, on the CPU threads and on the GPU) as a test bed for whatever shape it takes.

Why

The one-block pattern of tests/run/stencil3d.bend (and of bend-parallel's histogram, PR #1278): one Array, a fork tree whose every node holds a handle to it (Array.fork), each leaf writing its own part, the halves joined back (Array.join). It is how a GPU kernel is written in Bend today, and the fast one: in bend-ml I use it for matrix products on the CPU threads (bend-ml-tensor-fork: each leaf writes its rows of C) and on the GPU (bend-ml-tensor-gpu: each leaf writes its tile of C), and both give tensor-array's bits.

Both packages are @unsafe, and so is every program that uses them. That part is expected: Base models Array.fork(a) as (a, a) and Array.join(a, b) as a, so a write through the second handle is not in the model at all, and no law about the result can be stated. The alternatives cost what the pattern saves:

way to write disjoint parts in parallel what it costs
each leaf returns its part as a list, written into the array after the join 2-3.6x slower than not splitting, for products with a short inner dimension (bend-ml-tensor-array's NOTES); on CUDA, the lists' appends hit the lane stack (#1393)
each leaf gets its own array, concatenated after the join a copy of the whole output per call
Array.fork / Array.join O(1), the speed of the pattern, but @unsafe

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions