futhark
Safe HaskellNone
LanguageGHC2021

Futhark.Pass.Flatten.SOAC

Description

Flattening rules for SOACs.

Synopsis

Documentation

flattenFlatMap :: FlattenOps -> Pat Type -> SubExp -> [VName] -> ExtLambda SOACS -> FlattenM () Source #

Flattening rule for a top-level FlatMap, which is a very thin wrapper over just flattening the lambda. The result produced by FlatMap corresponds exactly to the internal irregular representation, so we simply distribute the lambda over the w segments and obtain an IrregularRep for each result (they share the same segment structure, as required by the type of FlatMap).

flattenFlatMapNested :: FlattenOps -> Segments -> DistEnv -> DistInputs -> [DistResult] -> StmAux () -> SubExp -> [VName] -> ExtLambda SOACS -> FlattenM DistEnv Source #

Flattening rule for a FlatMap nested inside an enclosing map-nest. A FlatMap is just a nonuniform map with implicit concatenation, so we run it through the ordinary inner-map machinery, which already produces - for a variably-sized result - an IrregularRep whose segment sizes are the per-enclosing-segment concatenated lengths and whose data is the concatenation. That directly gives the data arrays (results 4..) and, as their per-enclosing segment sizes, the total sizes m (result 0).

The metadata is derived from the shape array (result 1) - the per-iteration output sizes - which the inner-map does not itself surface. The lambda returns exactly that size as its leading result, so mapping it over the nest yields the irregular array of per-iteration sizes, segmented by the FlatMap width. From it we then compute the offset array (result 3), as the per-enclosing exclusive prefix sum, and the flag array (result 2).

See also Note [FlatMap element counting].