proofs/kernel_tile_refinement.bend
stelliferous@0.0.3.0 · 211 lines · 0 definitions
1# A shared native/model reader contract; pair composition owns the generic step.
2import Base
3import stelliferous@0.0.3.0/proofs/paired_read_refinement.bend as ReadLaws
4import stelliferous@0.0.3.0/kernel_block.bend as Block
5import stelliferous@0.0.3.0/math_sequence.bend as Sequence
6import stelliferous@0.0.3.0/kernel_tile.bend as Tile
7
8type SampleContext<-Model: Data,-Context: Data> is Data:
9 SampleContext{model: Model,context: Context}
10
11def sample_context(~Model: Data,~Context: Data,~sample: Model -> Context -> Nat -> F32,pair: SampleContext<Model,Context>,index: Nat) -> F32:
12 SampleContext{model,context} = pair
13 sample(model,context,index)
14
15def head_sample(~Model: Data,~Context: Data,~sample: Model -> Context -> Nat -> F32,model: Model,combined: Tile.HeadContext<Context>,
16 index: Nat) -> F32:
17 Tile.HeadContext{head,context} = combined
18 match index:
19 case 0n: head
20 case 1n+rest: sample(model,context,1n+rest)
21
22def head_read_refines(
23 ~Native: Type,~Model: Data,~Context: Data,
24 ~pack: Model -> Native,~read: Native -> Context -> Nat -> Native & F32,
25 ~sample: Model -> Context -> Nat -> F32,
26 ~read_rule: @+model: Model -> @+context: Context -> @+index: Nat -> {read(pack(model),context,index) == (pack(model),sample(model,context,
27 index)) : Native & F32},
28 +model: Model,+combined: Tile.HeadContext<Context>,+index: Nat
29) -> {Tile.read_with_head(~Native,~Context,~read,pack(model),combined,index) == (pack(model),head_sample(~Model,~Context,~sample,model,combined,
30 index)) : Native & F32}:
31 Tile.HeadContext{head,context} = combined
32 match index:
33 case 0n: {==}
34 case 1n+rest: read_rule(model,context,1n+rest)
35
36def read_pair_with_head_refines(
37 ~Native: Type,~Model: Data,~Context: Data,
38 ~pack: Model -> Native,~read: Native -> Context -> Nat -> Native & F32,
39 ~sample: Model -> Context -> Nat -> F32,
40 ~read_rule: @+model: Model -> @+context: Context -> @+index: Nat -> {read(pack(model),context,index) == (pack(model),sample(model,context,
41 index)) : Native & F32},
42 +model: Model,+context: Tile.HeadContext<Context>,+index: Nat
43) -> {Block.read_pair(~Native,~Tile.HeadContext<Context>,~F32,~Tile.read_with_head(~Native,~Context,~read),~1n,pack(model),context,
44 index) == (pack(model),ReadLaws.pair_sample(~Model,~Tile.HeadContext<Context>,~F32,~head_sample(~Model,~Context,~sample),~1n,model,context,
45 index)) : Native & Block.Twin<F32>}:
46 ReadLaws.read_pair_refines(~Native,~Model,~Tile.HeadContext<Context>,~F32,~pack,~Tile.read_with_head(~Native,~Context,~read),
47 ~head_sample(~Model,~Context,~sample),~1n,model,context,index,
48 head_read_refines(~Native,~Model,~Context,~pack,~read,~sample,~read_rule,model,context,index),
49 head_read_refines(~Native,~Model,~Context,~pack,~read,~sample,~read_rule,model,context,Nat.add(index,1n)))
50
51def read_pair_of_pairs_refines(
52 ~Native: Type,~Model: Data,~Context: Data,
53 ~pack: Model -> Native,~read: Native -> Context -> Nat -> Native & F32,
54 ~sample: Model -> Context -> Nat -> F32,
55 ~read_rule: @+model: Model -> @+context: Context -> @+index: Nat -> {read(pack(model),context,index) == (pack(model),sample(model,context,
56 index)) : Native & F32},
57 +model: Model,+context: Tile.HeadContext<Context>,+index: Nat
58) -> {Block.read_pair(~Native,~Tile.HeadContext<Context>,~Block.Twin<F32>,~Block.read_pair(~Native,
59 ~Tile.HeadContext<Context>,~F32,~Tile.read_with_head(~Native,~Context,~read),~1n),~2n,pack(model),context,index) == (pack(model),
60 ReadLaws.pair_sample(~Model,~Tile.HeadContext<Context>,~Block.Twin<F32>,~ReadLaws.pair_sample(~Model,~Tile.HeadContext<Context>,
61 ~F32,~head_sample(~Model,~Context,~sample),~1n),~2n,model,context,index)) : Native & Block.Twin<Block.Twin<F32>>}:
62 ReadLaws.read_pair_refines(~Native,~Model,~Tile.HeadContext<Context>,~Block.Twin<F32>,~pack,~Block.read_pair(~Native,
63 ~Tile.HeadContext<Context>,~F32,~Tile.read_with_head(~Native,~Context,~read),~1n),~ReadLaws.pair_sample(~Model,~Tile.HeadContext<Context>,
64 ~F32,~head_sample(~Model,~Context,~sample),~1n),~2n,model,context,index,
65 read_pair_with_head_refines(~Native,~Model,~Context,~pack,~read,~sample,~read_rule,model,context,index),
66 read_pair_with_head_refines(~Native,~Model,~Context,~pack,~read,~sample,~read_rule,model,context,Nat.add(index,2n)))
67
68def read_tile_pairs_refines(
69 ~Native: Type,~Model: Data,~Context: Data,
70 ~pack: Model -> Native,~read: Native -> Context -> Nat -> Native & F32,
71 ~sample: Model -> Context -> Nat -> F32,
72 ~read_rule: @+model: Model -> @+context: Context -> @+index: Nat -> {read(pack(model),context,index) == (pack(model),sample(model,context,
73 index)) : Native & F32},
74 +model: Model,+context: Tile.HeadContext<Context>,+index: Nat
75) -> {Block.read_pair(~Native,~Tile.HeadContext<Context>,~Block.Twin<Block.Twin<F32>>,
76 ~Block.read_pair(~Native,~Tile.HeadContext<Context>,~Block.Twin<F32>,~Block.read_pair(~Native,
77 ~Tile.HeadContext<Context>,~F32,~Tile.read_with_head(~Native,~Context,~read),~1n),~2n),~4n,pack(model),context,index) == (pack(model),
78 ReadLaws.pair_sample(~Model,~Tile.HeadContext<Context>,~Block.Twin<Block.Twin<F32>>,~ReadLaws.pair_sample(~Model,
79 ~Tile.HeadContext<Context>,~Block.Twin<F32>,~ReadLaws.pair_sample(~Model,~Tile.HeadContext<Context>,~F32,~head_sample(~Model,
80 ~Context,~sample),~1n),~2n),~4n,model,context,index)) : Native & Block.Twin<Block.Twin<Block.Twin<F32>>>}:
81 ReadLaws.read_pair_refines(~Native,~Model,~Tile.HeadContext<Context>,~Block.Twin<Block.Twin<F32>>,~pack,
82 ~Block.read_pair(~Native,~Tile.HeadContext<Context>,~Block.Twin<F32>,~Block.read_pair(~Native,
83 ~Tile.HeadContext<Context>,~F32,~Tile.read_with_head(~Native,~Context,~read),~1n),~2n),~ReadLaws.pair_sample(~Model,
84 ~Tile.HeadContext<Context>,~Block.Twin<F32>,~ReadLaws.pair_sample(~Model,~Tile.HeadContext<Context>,~F32,~head_sample(~Model,
85 ~Context,~sample),~1n),~2n),~4n,model,context,index,
86 read_pair_of_pairs_refines(~Native,~Model,~Context,~pack,~read,~sample,~read_rule,model,context,index),
87 read_pair_of_pairs_refines(~Native,~Model,~Context,~pack,~read,~sample,~read_rule,model,context,Nat.add(index,4n)))
88
89def tail_block_refines(
90 ~Native: Type,~Model: Data,~Context: Data,
91 ~pack: Model -> Native,
92 ~read: Native -> Context -> Nat -> Native & F32,
93 ~sample: Model -> Context -> Nat -> F32,
94 ~read_rule: @+model: Model -> @+context: Context -> @+index: Nat -> {read(pack(model),context,index) == (pack(model),sample(model,context,
95 index)) : Native & F32},
96 +model: Model,+context: Context,+head: F32
97) -> {Tile.read_tail(~Native,~Context,~read,(pack(model),head),context) == (pack(model),Tile.from_sequence(Sequence.Elements{head,
98 Sequence.tabulate(~F32,~SampleContext<Model,Context>,~sample_context(~Model,~Context,~sample),7n,1n,SampleContext{model,
99 context})})) : Native & Tile.Tile8}:
100 %read_tile_pairs_refines(~Native,~Model,~Context,~pack,~read,~sample,~read_rule,model,Tile.HeadContext{head,context},
101 0n) : {Tile.finish_pairs(~Native,Block.read_pair(~Native,~Tile.HeadContext<Context>,~Block.Twin<Block.Twin<F32>>,
102 ~Block.read_pair(~Native,~Tile.HeadContext<Context>,~Block.Twin<F32>,~Block.read_pair(~Native,
103 ~Tile.HeadContext<Context>,~F32,~Tile.read_with_head(~Native,~Context,~read),~1n),~2n),~4n,pack(model),Tile.HeadContext{head,context},
104 0n)) == Tile.finish_pairs(~Native,_) : Native & Tile.Tile8}
105 {==}
106
107# Tile projections and extensional equality used by GEMM callers.
108def lane_of(lane: Nat) -> Sequence.Index(8n):
109 match lane:
110 case 0n: Sequence.First{}
111 case 1n+a:
112 match a:
113 case 0n: Sequence.Next{Sequence.First{}}
114 case 1n+b:
115 match b:
116 case 0n: Sequence.Next{Sequence.Next{Sequence.First{}}}
117 case 1n+c:
118 match c:
119 case 0n: Sequence.Next{Sequence.Next{Sequence.Next{Sequence.First{}}}}
120 case 1n+d:
121 match d:
122 case 0n: Sequence.Next{Sequence.Next{Sequence.Next{Sequence.Next{Sequence.First{}}}}}
123 case 1n+e:
124 match e:
125 case 0n: Sequence.Next{Sequence.Next{Sequence.Next{Sequence.Next{Sequence.Next{Sequence.First{}}}}}}
126 case 1n+f:
127 match f:
128 case 0n: Sequence.Next{Sequence.Next{Sequence.Next{Sequence.Next{Sequence.Next{Sequence.Next{Sequence.First{}}}}}}}
129 case 1n+g: Sequence.Next{Sequence.Next{Sequence.Next{Sequence.Next{Sequence.Next{Sequence.Next{Sequence.Next{Sequence.First{}}}}}}}}
130
131def select_lane(lane: Nat,values: Tile.Tile8) -> F32:
132 match lane:
133 case 0n: Tile.lane0(values)
134 case 1n+a:
135 match a:
136 case 0n: Tile.lane1(values)
137 case 1n+b:
138 match b:
139 case 0n: Tile.lane2(values)
140 case 1n+c:
141 match c:
142 case 0n: Tile.lane3(values)
143 case 1n+d:
144 match d:
145 case 0n: Tile.lane4(values)
146 case 1n+e:
147 match e:
148 case 0n: Tile.lane5(values)
149 case 1n+f:
150 match f:
151 case 0n: Tile.lane6(values)
152 case 1n+g: Tile.lane7(values)
153
154def selected_lane(+lane: Nat,+values: Tile.Tile8) ->
155 {select_lane(lane,values) == Tile.element_at(lane_of(lane),values) : F32}:
156 match lane:
157 case 0n:
158 Tile.Tile{a,b,c,d,e,f,g,h} = values
159 {==}
160 case 1n+ +a:
161 match a:
162 case 0n:
163 Tile.Tile{a,b,c,d,e,f,g,h} = values
164 {==}
165 case 1n+ +b:
166 match b:
167 case 0n:
168 Tile.Tile{a,b,c,d,e,f,g,h} = values
169 {==}
170 case 1n+ +c:
171 match c:
172 case 0n:
173 Tile.Tile{a,b,c,d,e,f,g,h} = values
174 {==}
175 case 1n+ +d:
176 match d:
177 case 0n:
178 Tile.Tile{a,b,c,d,e,f,g,h} = values
179 {==}
180 case 1n+ +e:
181 match e:
182 case 0n:
183 Tile.Tile{a,b,c,d,e,f,g,h} = values
184 {==}
185 case 1n+ +f:
186 match f:
187 case 0n:
188 Tile.Tile{a,b,c,d,e,f,g,h} = values
189 {==}
190 case 1n+ +g:
191 Tile.Tile{a,b,c,d,e,f,g,h} = values
192 {==}
193
194def tile_eta(+values: Tile.Tile8) -> {values == Tile.Tile{Tile.lane0(values),Tile.lane1(values),Tile.lane2(values),Tile.lane3(values),Tile.lane4(values),Tile.lane5(values),Tile.lane6(values),Tile.lane7(values)} : Tile.Tile8}:
195 Tile.Tile{a,b,c,d,e,f,g,h} = values
196 {==}
197
198def tiles_equal(+left: Tile.Tile8,+right: Tile.Tile8,
199 e0: {Tile.lane0(left) == Tile.lane0(right) : F32},e1: {Tile.lane1(left) == Tile.lane1(right) : F32},e2: {Tile.lane2(left) == Tile.lane2(right) : F32},e3: {Tile.lane3(left) == Tile.lane3(right) : F32},e4: {Tile.lane4(left) == Tile.lane4(right) : F32},e5: {Tile.lane5(left) == Tile.lane5(right) : F32},e6: {Tile.lane6(left) == Tile.lane6(right) : F32},e7: {Tile.lane7(left) == Tile.lane7(right) : F32}) -> {left == right : Tile.Tile8}:
200 %Equal.sym(Tile.Tile8,left,Tile.Tile{Tile.lane0(left),Tile.lane1(left),Tile.lane2(left),Tile.lane3(left),Tile.lane4(left),Tile.lane5(left),Tile.lane6(left),Tile.lane7(left)},tile_eta(left)) : {_ == right : Tile.Tile8}
201 %Equal.sym(Tile.Tile8,right,Tile.Tile{Tile.lane0(right),Tile.lane1(right),Tile.lane2(right),Tile.lane3(right),Tile.lane4(right),Tile.lane5(right),Tile.lane6(right),Tile.lane7(right)},tile_eta(right)) :
202 {Tile.Tile{Tile.lane0(left),Tile.lane1(left),Tile.lane2(left),Tile.lane3(left),Tile.lane4(left),Tile.lane5(left),Tile.lane6(left),Tile.lane7(left)} == _ : Tile.Tile8}
203 %e0 : {Tile.Tile{Tile.lane0(left),Tile.lane1(left),Tile.lane2(left),Tile.lane3(left),Tile.lane4(left),Tile.lane5(left),Tile.lane6(left),Tile.lane7(left)} == Tile.Tile{_,Tile.lane1(right),Tile.lane2(right),Tile.lane3(right),Tile.lane4(right),Tile.lane5(right),Tile.lane6(right),Tile.lane7(right)} : Tile.Tile8}
204 %e1 : {Tile.Tile{Tile.lane0(left),Tile.lane1(left),Tile.lane2(left),Tile.lane3(left),Tile.lane4(left),Tile.lane5(left),Tile.lane6(left),Tile.lane7(left)} == Tile.Tile{Tile.lane0(left),_,Tile.lane2(right),Tile.lane3(right),Tile.lane4(right),Tile.lane5(right),Tile.lane6(right),Tile.lane7(right)} : Tile.Tile8}
205 %e2 : {Tile.Tile{Tile.lane0(left),Tile.lane1(left),Tile.lane2(left),Tile.lane3(left),Tile.lane4(left),Tile.lane5(left),Tile.lane6(left),Tile.lane7(left)} == Tile.Tile{Tile.lane0(left),Tile.lane1(left),_,Tile.lane3(right),Tile.lane4(right),Tile.lane5(right),Tile.lane6(right),Tile.lane7(right)} : Tile.Tile8}
206 %e3 : {Tile.Tile{Tile.lane0(left),Tile.lane1(left),Tile.lane2(left),Tile.lane3(left),Tile.lane4(left),Tile.lane5(left),Tile.lane6(left),Tile.lane7(left)} == Tile.Tile{Tile.lane0(left),Tile.lane1(left),Tile.lane2(left),_,Tile.lane4(right),Tile.lane5(right),Tile.lane6(right),Tile.lane7(right)} : Tile.Tile8}
207 %e4 : {Tile.Tile{Tile.lane0(left),Tile.lane1(left),Tile.lane2(left),Tile.lane3(left),Tile.lane4(left),Tile.lane5(left),Tile.lane6(left),Tile.lane7(left)} == Tile.Tile{Tile.lane0(left),Tile.lane1(left),Tile.lane2(left),Tile.lane3(left),_,Tile.lane5(right),Tile.lane6(right),Tile.lane7(right)} : Tile.Tile8}
208 %e5 : {Tile.Tile{Tile.lane0(left),Tile.lane1(left),Tile.lane2(left),Tile.lane3(left),Tile.lane4(left),Tile.lane5(left),Tile.lane6(left),Tile.lane7(left)} == Tile.Tile{Tile.lane0(left),Tile.lane1(left),Tile.lane2(left),Tile.lane3(left),Tile.lane4(left),_,Tile.lane6(right),Tile.lane7(right)} : Tile.Tile8}
209 %e6 : {Tile.Tile{Tile.lane0(left),Tile.lane1(left),Tile.lane2(left),Tile.lane3(left),Tile.lane4(left),Tile.lane5(left),Tile.lane6(left),Tile.lane7(left)} == Tile.Tile{Tile.lane0(left),Tile.lane1(left),Tile.lane2(left),Tile.lane3(left),Tile.lane4(left),Tile.lane5(left),_,Tile.lane7(right)} : Tile.Tile8}
210 %e7 : {Tile.Tile{Tile.lane0(left),Tile.lane1(left),Tile.lane2(left),Tile.lane3(left),Tile.lane4(left),Tile.lane5(left),Tile.lane6(left),Tile.lane7(left)} == Tile.Tile{Tile.lane0(left),Tile.lane1(left),Tile.lane2(left),Tile.lane3(left),Tile.lane4(left),Tile.lane5(left),Tile.lane6(left),_} : Tile.Tile8}
211 {==}