1 # The FP32 cells of Cells.map: eight consecutive values as one Tile8. 2 import Base 3 import ./kernel_tile.bend as Tile 4 import ./traversal_cells.bend as Cells 5 import ./traversal_array.bend as Traversal 6 7 def gathered(~Context: Data,~map: Context -> F32 -> F32,+context: Context,+a: F32,+b: F32,+c: F32,+d: F32,+e: F32,+f: F32,+g: F32, 8 read: Array & F32) -> Array & Tile.Tile8: 9 (values,h) = read 10 (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)}) 11 12 def gather7(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32,+b: F32,+c: F32,+d: F32,+e: F32,+f: F32, 13 read: Array & F32) -> Array & Tile.Tile8: 14 (values,g) = read 15 gathered(~Context,~map,context,a,b,c,d,e,f,g,Array.get(F32,values,(position + 1 : U32))) 16 17 def gather6(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32,+b: F32,+c: F32,+d: F32,+e: F32, 18 read: Array & F32) -> Array & Tile.Tile8: 19 (values,f) = read 20 gather7(~Context,~map,(position + 1 : U32),context,a,b,c,d,e,f,Array.get(F32,values,(position + 1 : U32))) 21 22 def gather5(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32,+b: F32,+c: F32,+d: F32, 23 read: Array & F32) -> Array & Tile.Tile8: 24 (values,e) = read 25 gather6(~Context,~map,(position + 1 : U32),context,a,b,c,d,e,Array.get(F32,values,(position + 1 : U32))) 26 27 def gather4(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32,+b: F32,+c: F32, 28 read: Array & F32) -> Array & Tile.Tile8: 29 (values,d) = read 30 gather5(~Context,~map,(position + 1 : U32),context,a,b,c,d,Array.get(F32,values,(position + 1 : U32))) 31 32 def gather3(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32,+b: F32, 33 read: Array & F32) -> Array & Tile.Tile8: 34 (values,c) = read 35 gather4(~Context,~map,(position + 1 : U32),context,a,b,c,Array.get(F32,values,(position + 1 : U32))) 36 37 def gather2(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context,+a: F32, 38 read: Array & F32) -> Array & Tile.Tile8: 39 (values,b) = read 40 gather3(~Context,~map,(position + 1 : U32),context,a,b,Array.get(F32,values,(position + 1 : U32))) 41 42 def gather1(~Context: Data,~map: Context -> F32 -> F32,+position: U32,+context: Context, 43 read: Array & F32) -> Array & Tile.Tile8: 44 (values,a) = read 45 gather2(~Context,~map,(position + 1 : U32),context,a,Array.get(F32,values,(position + 1 : U32))) 46 47 # The mapped values at position, position + 1, ... as one tile; the array is only read. 48 def gather(~Context: Data,~map: Context -> F32 -> F32,values: Array,position: U32,context: Context) -> Array & Tile.Tile8: 49 +position = position 50 gather1(~Context,~map,position,context,Array.get(F32,values,position)) 51 52 # The lanes of a tile to position, position + 1, ... 53 def scatter(values: Array,position: U32,tile: Tile.Tile8) -> Array: 54 +tile = tile 55 +position = position 56 +p1 = (position + 1 : U32) 57 +p2 = (p1 + 1 : U32) 58 +p3 = (p2 + 1 : U32) 59 +p4 = (p3 + 1 : U32) 60 +p5 = (p4 + 1 : U32) 61 +p6 = (p5 + 1 : U32) 62 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)), 63 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)) 64 65 def scratch() -> Array: 66 Array.new(Tile.Tile8,6n,Tile.filled(0.0)) 67 68 # The values the rounds leave: fewer than eight, mapped one by one. 69 def finish(~Context: Data,~map: Context -> F32 -> F32,rest: Nat,state: Cells.State,Tile.Tile8,Context>) -> Array: 70 Cells.State{values,scratch,position,slot,context} = state 71 Traversal.map_in_place(~F32,~Context,~map,rest,position,context,values) 72 73 # One eight-value loop for any owned input bundle; gather and scatter decide 74 # which owners are read or updated. The unary path is its Array specialization. 75 def rounds_with(~Owner: Type,~Context: Data,~gather: Owner -> U32 -> Context -> Owner & Tile.Tile8, 76 ~scatter: Owner -> U32 -> Tile.Tile8 -> Owner,~finish: Nat -> Cells.State -> Owner, 77 count: Nat,state: Cells.State) -> Owner: 78 match count: 79 case 1n+1n+1n+1n+1n+1n+1n+1n+rest: 80 rounds_with(~Owner,~Context,~gather,~scatter,~finish,rest,Cells.step(~Owner,~Tile.Tile8,~Context,~gather,~scatter,~8,state)) 81 case few: finish(few,state) 82 83 def rounds(~Context: Data,~map: Context -> F32 -> F32,count: Nat,state: Cells.State,Tile.Tile8,Context>) -> Array: 84 rounds_with(~Array,~Context,~gather(~Context,~map),~scatter,~finish(~Context,~map),count,state) 85 86 # map over the count values from position, in order: Traversal.map_in_place 87 # with the values of each round computed together. 88 def map(~Context: Data,~map: Context -> F32 -> F32,count: Nat,position: U32,context: Context,values: Array) -> Array: 89 rounds(~Context,~map,count,Cells.State{values,scratch(),position,0,context}) 90 91 # An array of 2^depth values holds the count positions from position, within 92 # the depths the refinement proof covers. 93 def covered(+depth: Nat,+position: U32,+count: Nat) -> Bool: 94 Nat.is_le(depth,30n) && Nat.is_le(Nat.add(U32.to_nat(position),count),U32.to_nat(U32.shln(1,depth))) 95 96 # map over count values from position when the array is known to hold them, 97 # and the scalar map otherwise: the same array either way. 98 def map_when(~Context: Data,~operation: Context -> F32 -> F32,covered: Bool,count: Nat,position: U32,context: Context,values: Array) -> Array: 99 match covered: 100 case True{}: map(~Context,~operation,count,position,context,values) 101 case False{}: Traversal.map_in_place(~F32,~Context,~operation,count,position,context,values)