1 # Owned FP32 input bundles use the same Cells/Lanes rounds as unary maps. 2 import Base 3 import ./math_compare.bend as Compare 4 import ./storage_buffer.bend as Storage 5 import ./traversal_lanes.bend as Lanes 6 import ./traversal_cells.bend as Cells 7 import ./kernel_tile.bend as Tile 8 import ./tensor_operations.bend as Operations 9 import ./tensor_view.bend as Views 10 import ./tensor_access.bend as Access 11 import ./proofs/storage_certificate.bend as Certificate 12 import ./proofs/traversal_lanes_certificate.bend as LanesCertificate 13 import ./proofs/traversal_witness.bend as Witness 14 15 def identity(context: Unit,value: F32) -> F32: 16 value 17 18 def read_result(length: Nat,result: Array & Tile.Tile8, 19 certificate: Certificate.Certificate,~Tile.Tile8,result)>) -> Storage.Buffer & Tile.Tile8: 20 (values,tile) = result 21 (Storage.Buffer{values,length,certificate},tile) 22 23 def read(buffer: Storage.Buffer,+position: U32) -> Storage.Buffer & Tile.Tile8: 24 Storage.Buffer{values,length,certificate} = buffer 25 read_result(length,Lanes.gather(~Unit,~identity,values,position,Unit{}),LanesCertificate.gather(~Unit,~identity,values,position,Unit{},certificate)) 26 27 def write(buffer: Storage.Buffer,+position: U32,+tile: Tile.Tile8) -> Storage.Buffer: 28 Storage.Buffer{values,length,certificate} = buffer 29 Storage.Buffer{Lanes.scatter(values,position,tile),length,LanesCertificate.scatter(values,position,tile,certificate)} 30 31 def combine_tile(~operation: F32 -> F32 -> F32,left: Tile.Tile8,right: Tile.Tile8) -> Tile.Tile8: 32 Tile.Tile{a,b,c,d,e,f,g,h} = left 33 Tile.Tile{i,j,k,l,m,n,o,p} = right 34 Tile.Tile{operation(a,i),operation(b,j),operation(c,k),operation(d,l),operation(e,m),operation(f,n),operation(g,o),operation(h,p)} 35 36 type Binary is Type: 37 Binary{left: Storage.Buffer,right: Storage.Buffer} 38 39 def pair_reads(~operation: F32 -> F32 -> F32,left: Storage.Buffer & Tile.Tile8,right: Storage.Buffer & Tile.Tile8) -> Binary & Tile.Tile8: 40 (left,a) = left 41 (right,b) = right 42 (Binary{left,right},combine_tile(~operation,a,b)) 43 44 def pair_gather(~operation: F32 -> F32 -> F32,owners: Binary,position: U32,context: Unit) -> Binary & Tile.Tile8: 45 Binary{left,right} = owners 46 +position = position 47 pair_reads(~operation,read(left,position),read(right,position)) 48 49 def pair_scatter(owners: Binary,position: U32,tile: Tile.Tile8) -> Binary: 50 Binary{left,right} = owners 51 Binary{write(left,position,tile),right} 52 53 def pair_result(result: Storage.Buffer & Storage.Buffer) -> Binary: 54 (left,right) = result 55 Binary{left,right} 56 57 # The checked caller has already validated both views and owners. 58 def pair_finish(~operation: F32 -> F32 -> F32,count: Nat,state: Cells.State) -> Binary: 59 Cells.State{Binary{left,right},scratch,+position,slot,context} = state 60 +count = count 61 pair_result(Operations.combine_kernel(~F32,~operation,count,Views.View{[Views.Axis{U32.from_nat(count),1}],position},position,left,right)) 62 63 def pair_map(~operation: F32 -> F32 -> F32,count: Nat,position: U32,left: Storage.Buffer,right: Storage.Buffer) -> Binary: 64 Lanes.rounds_with(~Binary,~Unit,~pair_gather(~operation),~pair_scatter,~pair_finish(~operation),count, 65 Cells.State{Binary{left,right},Lanes.scratch(),position,0,Unit{}}) 66 67 def pair_buffers(owners: Binary) -> Storage.Buffer & Storage.Buffer: 68 Binary{left,right} = owners 69 (left,right) 70 71 def pair_when(~operation: F32 -> F32 -> F32,fast: Bool,count: Nat,right_view: Views.View,offset: U32, 72 left: Storage.Buffer,right: Storage.Buffer) -> Storage.Buffer & Storage.Buffer: 73 match fast: 74 case True{}: pair_buffers(pair_map(~operation,count,0,left,right)) 75 case False{}: Operations.combine_kernel(~F32,~operation,count,right_view,offset,left,right) 76 77 def combine(~operation: F32 -> F32 -> F32,+count: Nat,+right_view: Views.View,+offset: U32, 78 left: Storage.Buffer,right: Storage.Buffer) -> Storage.Buffer & Storage.Buffer: 79 pair_when(~operation,Access.zero_contiguous(right_view) && U32.is_zero(offset),count,right_view,offset,left,right) 80 81 type Selected is Type: 82 Selected{mask: Storage.Buffer,yes: Storage.Buffer,no: Storage.Buffer} 83 84 def select_tile(mask: Tile.Tile8,yes: Tile.Tile8,no: Tile.Tile8) -> Tile.Tile8: 85 Tile.Tile{a,b,c,d,e,f,g,h} = mask 86 Tile.Tile{i,j,k,l,m,n,o,p} = yes 87 Tile.Tile{q,r,s,t,u,v,w,x} = no 88 Tile.Tile{Compare.select(a,i,q),Compare.select(b,j,r),Compare.select(c,k,s),Compare.select(d,l,t), 89 Compare.select(e,m,u),Compare.select(f,n,v),Compare.select(g,o,w),Compare.select(h,p,x)} 90 91 def select_reads(mask: Storage.Buffer & Tile.Tile8,yes: Storage.Buffer & Tile.Tile8,no: Storage.Buffer & Tile.Tile8) -> Selected & Tile.Tile8: 92 (mask,a) = mask 93 (yes,b) = yes 94 (no,c) = no 95 (Selected{mask,yes,no},select_tile(a,b,c)) 96 97 def select_gather(owners: Selected,position: U32,context: Unit) -> Selected & Tile.Tile8: 98 Selected{mask,yes,no} = owners 99 +position = position 100 select_reads(read(mask,position),read(yes,position),read(no,position)) 101 102 def select_scatter(owners: Selected,position: U32,tile: Tile.Tile8) -> Selected: 103 Selected{mask,yes,no} = owners 104 Selected{mask,write(yes,position,tile),no} 105 106 def select_joined(mask: Storage.Buffer,joined: Storage.Buffer & Storage.Buffer) -> Selected: 107 (yes,no) = joined 108 Selected{mask,yes,no} 109 110 def select_rejected(+count: Nat,+position: U32,yes: Storage.Buffer,result: Storage.Buffer & Storage.Buffer) -> Selected: 111 (no,mask) = result 112 select_joined(mask,Operations.combine_kernel(~F32,~Compare.union,count,Views.View{[Views.Axis{U32.from_nat(count),1}],position},position,yes,no)) 113 114 def select_chosen(+count: Nat,+position: U32,no: Storage.Buffer,result: Storage.Buffer & Storage.Buffer) -> Selected: 115 (yes,mask) = result 116 select_rejected(count,position,yes, 117 Operations.combine_kernel(~F32,~Compare.rejected,count,Views.View{[Views.Axis{U32.from_nat(count),1}],position},position,no,mask)) 118 119 # The final 0..7 values use the existing certified binary kernels. Only the 120 # tail makes three passes; full cells read all three inputs in one round. 121 def select_finish(count: Nat,state: Cells.State) -> Selected: 122 Cells.State{Selected{mask,yes,no},scratch,+position,slot,context} = state 123 +count = count 124 select_chosen(count,position,no, 125 Operations.combine_kernel(~F32,~Compare.selected,count,Views.View{[Views.Axis{U32.from_nat(count),1}],position},position,yes,mask)) 126 127 def select_result(owners: Selected) -> Storage.Buffer: 128 Selected{mask,yes,no} = owners 129 yes 130 131 def select(count: Nat,mask: Storage.Buffer,yes: Storage.Buffer,no: Storage.Buffer) -> Storage.Buffer: 132 select_result(Lanes.rounds_with(~Selected,~Unit,~select_gather,~select_scatter,~select_finish,count, 133 Cells.State{Selected{mask,yes,no},Lanes.scratch(),0,0,Unit{}}))