tensor_lanes.bend
stelliferous@0.0.2.0 · 133 lines · 26 definitions
1
2import Base
3import stelliferous@0.0.2.0/math_compare.bend as Compare
4import stelliferous@0.0.2.0/storage_buffer.bend as Storage
5import stelliferous@0.0.2.0/traversal_lanes.bend as Lanes
6import stelliferous@0.0.2.0/traversal_cells.bend as Cells
7import stelliferous@0.0.2.0/kernel_tile.bend as Tile
8import stelliferous@0.0.2.0/tensor_operations.bend as Operations
9import stelliferous@0.0.2.0/tensor_view.bend as Views
10import stelliferous@0.0.2.0/tensor_access.bend as Access
11import stelliferous@0.0.2.0/proofs/storage_certificate.bend as Certificate
12import stelliferous@0.0.2.0/proofs/traversal_lanes_certificate.bend as LanesCertificate
13import stelliferous@0.0.2.0/proofs/traversal_witness.bend as Witness
14
15def identity(context: Unit,value: F32) -> F32:
16 value
17
18def read_result(length: Nat,result: Array<F32> & Tile.Tile8,
19 certificate: Certificate.Certificate<F32,Witness.first(~Array<F32>,~Tile.Tile8,result)>) -> Storage.Buffer<F32> & Tile.Tile8:
20 (values,tile) = result
21 (Storage.Buffer{values,length,certificate},tile)
22
23def read(buffer: Storage.Buffer<F32>,+position: U32) -> Storage.Buffer<F32> & 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
27def write(buffer: Storage.Buffer<F32>,+position: U32,+tile: Tile.Tile8) -> Storage.Buffer<F32>:
28 Storage.Buffer{values,length,certificate} = buffer
29 Storage.Buffer{Lanes.scatter(values,position,tile),length,LanesCertificate.scatter(values,position,tile,certificate)}
30
31def 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
36type Binary is Type:
37 Binary{left: Storage.Buffer<F32>,right: Storage.Buffer<F32>}
38
39def pair_reads(~operation: F32 -> F32 -> F32,left: Storage.Buffer<F32> & Tile.Tile8,right: Storage.Buffer<F32> & Tile.Tile8) -> Binary & Tile.Tile8:
40 (left,a) = left
41 (right,b) = right
42 (Binary{left,right},combine_tile(~operation,a,b))
43
44def 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
49def pair_scatter(owners: Binary,position: U32,tile: Tile.Tile8) -> Binary:
50 Binary{left,right} = owners
51 Binary{write(left,position,tile),right}
52
53def pair_result(result: Storage.Buffer<F32> & Storage.Buffer<F32>) -> Binary:
54 (left,right) = result
55 Binary{left,right}
56
57
58def pair_finish(~operation: F32 -> F32 -> F32,count: Nat,state: Cells.State<Binary,Tile.Tile8,Unit>) -> 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
63def pair_map(~operation: F32 -> F32 -> F32,count: Nat,position: U32,left: Storage.Buffer<F32>,right: Storage.Buffer<F32>) -> 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
67def pair_buffers(owners: Binary) -> Storage.Buffer<F32> & Storage.Buffer<F32>:
68 Binary{left,right} = owners
69 (left,right)
70
71def pair_when(~operation: F32 -> F32 -> F32,fast: Bool,count: Nat,right_view: Views.View,offset: U32,
72 left: Storage.Buffer<F32>,right: Storage.Buffer<F32>) -> Storage.Buffer<F32> & Storage.Buffer<F32>:
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
77def combine(~operation: F32 -> F32 -> F32,+count: Nat,+right_view: Views.View,+offset: U32,
78 left: Storage.Buffer<F32>,right: Storage.Buffer<F32>) -> Storage.Buffer<F32> & Storage.Buffer<F32>:
79 pair_when(~operation,Access.zero_contiguous(right_view) && U32.is_zero(offset),count,right_view,offset,left,right)
80
81type Selected is Type:
82 Selected{mask: Storage.Buffer<F32>,yes: Storage.Buffer<F32>,no: Storage.Buffer<F32>}
83
84def 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
91def select_reads(mask: Storage.Buffer<F32> & Tile.Tile8,yes: Storage.Buffer<F32> & Tile.Tile8,no: Storage.Buffer<F32> & 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
97def 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
102def 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
106def select_joined(mask: Storage.Buffer<F32>,joined: Storage.Buffer<F32> & Storage.Buffer<F32>) -> Selected:
107 (yes,no) = joined
108 Selected{mask,yes,no}
109
110def select_rejected(+count: Nat,+position: U32,yes: Storage.Buffer<F32>,result: Storage.Buffer<F32> & Storage.Buffer<F32>) -> 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
114def select_chosen(+count: Nat,+position: U32,no: Storage.Buffer<F32>,result: Storage.Buffer<F32> & Storage.Buffer<F32>) -> 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
120
121def select_finish(count: Nat,state: Cells.State<Selected,Tile.Tile8,Unit>) -> 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
127def select_result(owners: Selected) -> Storage.Buffer<F32>:
128 Selected{mask,yes,no} = owners
129 yes
130
131def select(count: Nat,mask: Storage.Buffer<F32>,yes: Storage.Buffer<F32>,no: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
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{}}))