traversal_lanes.bend source
traversal_lanes.bend on the hub · documented module
# The FP32 cells of Cells.map: eight consecutive values as one Tile8.import Baseimport ./kernel_tile.bend as Tileimport ./traversal_cells.bend as Cellsimport ./traversal_array.bend as Traversaldef gathered(~Context: Data,~map: Context -> F32 -> F32,+context: Context,+a: F32,+b: F32,+c: F32,+d: F32,+e: F32,+f: F32,+g: F32, read: Array<F32> & F32) -> Array<F32> & Tile.Tile8: (values,h) = read (values,Tile.Tile{map(context,a),map(context,b),map(context,c),map(context,d),map(context,e),map(context,f),map(context,g),map(context,h)})def gather7(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32,+b: F32,+c: F32,+d: F32,+e: F32,+f: F32, read: Array<F32> & F32) -> Array<F32> & Tile.Tile8: (values,g) = read gathered(~Context,~map,context,a,b,c,d,e,f,g,Array.get(F32,values,(position + 1 : U32)))def gather6(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32,+b: F32,+c: F32,+d: F32,+e: F32, read: Array<F32> & F32) -> Array<F32> & Tile.Tile8: (values,f) = read gather7(~Context,~map,(position + 1 : U32),context,a,b,c,d,e,f,Array.get(F32,values,(position + 1 : U32)))def gather5(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32,+b: F32,+c: F32,+d: F32, read: Array<F32> & F32) -> Array<F32> & Tile.Tile8: (values,e) = read gather6(~Context,~map,(position + 1 : U32),context,a,b,c,d,e,Array.get(F32,values,(position + 1 : U32)))def gather4(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32,+b: F32,+c: F32, read: Array<F32> & F32) -> Array<F32> & Tile.Tile8: (values,d) = read gather5(~Context,~map,(position + 1 : U32),context,a,b,c,d,Array.get(F32,values,(position + 1 : U32)))def gather3(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32,+b: F32, read: Array<F32> & F32) -> Array<F32> & Tile.Tile8: (values,c) = read gather4(~Context,~map,(position + 1 : U32),context,a,b,c,Array.get(F32,values,(position + 1 : U32)))def gather2(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32, read: Array<F32> & F32) -> Array<F32> & Tile.Tile8: (values,b) = read gather3(~Context,~map,(position + 1 : U32),context,a,b,Array.get(F32,values,(position + 1 : U32)))def gather1(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context, read: Array<F32> & F32) -> Array<F32> & Tile.Tile8: (values,a) = read gather2(~Context,~map,(position + 1 : U32),context,a,Array.get(F32,values,(position + 1 : U32)))# The mapped values at position, position + 1, ... as one tile; the array is only read.def gather(~Context: Data,~map: Context -> F32 -> F32,values: Array<F32>,position: U32,context: Context) -> Array<F32> & Tile.Tile8: +position = position gather1(~Context,~map,position,context,Array.get(F32,values,position))# The lanes of a tile to position, position + 1, ...def scatter(values: Array<F32>,position: U32,tile: Tile.Tile8) -> Array<F32>: +tile = tile +position = position +p1 = (position + 1 : U32) +p2 = (p1 + 1 : U32) +p3 = (p2 + 1 : U32) +p4 = (p3 + 1 : U32) +p5 = (p4 + 1 : U32) +p6 = (p5 + 1 : U32) Array.set(F32,Array.set(F32,Array.set(F32,Array.set(F32,Array.set(F32,Array.set(F32,Array.set(F32,Array.set(F32,values,position,Tile.lane0(tile)),p1,Tile.lane1(tile)), p2,Tile.lane2(tile)),p3,Tile.lane3(tile)),p4,Tile.lane4(tile)),p5,Tile.lane5(tile)),p6,Tile.lane6(tile)),(p6 + 1 : U32),Tile.lane7(tile))def scratch() -> Array<Tile.Tile8>: Array.new(Tile.Tile8,6n,Tile.filled(0.0))# The values the rounds leave: fewer than eight, mapped one by one.def finish(~Context: Data,~map: Context -> F32 -> F32,rest: Nat,state: Cells.State<Array<F32>,Tile.Tile8,Context>) -> Array<F32>: Cells.State{values,scratch,position,slot,context} = state Traversal.map_in_place(~F32,~Context,~map,rest,position,context,values)# One eight-value loop for any owned input bundle; gather and scatter decide# which owners are read or updated. The unary path is its Array specialization.def rounds_with(~Owner: Type,~Context: Data,~gather: Owner -> U32 -> Context -> Owner & Tile.Tile8, ~scatter: Owner -> U32 -> Tile.Tile8 -> Owner,~finish: Nat -> Cells.State<Owner,Tile.Tile8,Context> -> Owner, count: Nat,state: Cells.State<Owner,Tile.Tile8,Context>) -> Owner: match count: case 1n+1n+1n+1n+1n+1n+1n+1n+rest: rounds_with(~Owner,~Context,~gather,~scatter,~finish,rest,Cells.step(~Owner,~Tile.Tile8,~Context,~gather,~scatter,~8,state)) case few: finish(few,state)def rounds(~Context: Data,~map: Context -> F32 -> F32,count: Nat,state: Cells.State<Array<F32>,Tile.Tile8,Context>) -> Array<F32>: rounds_with(~Array<F32>,~Context,~gather(~Context,~map),~scatter,~finish(~Context,~map),count,state)# map over the count values from position, in order: Traversal.map_in_place# with the values of each round computed together.def map(~Context: Data,~map: Context -> F32 -> F32,count: Nat,position: U32,context: Context,values: Array<F32>) -> Array<F32>: rounds(~Context,~map,count,Cells.State{values,scratch(),position,0,context})# An array of 2^depth values holds the count positions from position, within# the depths the refinement proof covers.def covered(+depth: Nat,+position: U32,+count: Nat) -> Bool: Nat.is_le(depth,30n) && Nat.is_le(Nat.add(U32.to_nat(position),count),U32.to_nat(U32.shln(1,depth)))# map over count values from position when the array is known to hold them,# and the scalar map otherwise: the same array either way.def map_when(~Context: Data,~operation: Context -> F32 -> F32,covered: Bool,count: Nat,position: U32,context: Context,values: Array<F32>) -> Array<F32>: match covered: case True{}: map(~Context,~operation,count,position,context,values) case False{}: Traversal.map_in_place(~F32,~Context,~operation,count,position,context,values)