1 # A shared native/model reader contract; pair composition owns the generic step. 2 import Base 3 import ./paired_read_refinement.bend as ReadLaws 4 import ../kernel_block.bend as Block 5 import ../math_sequence.bend as Sequence 6 import ../kernel_tile.bend as Tile 7 8 type SampleContext<-Model: Data,-Context: Data> is Data: 9 SampleContext{model: Model,context: Context} 10 11 def sample_context(~Model: Data,~Context: Data,~sample: Model -> Context -> Nat -> F32,pair: SampleContext,index: Nat) -> F32: 12 SampleContext{model,context} = pair 13 sample(model,context,index) 14 15 def head_sample(~Model: Data,~Context: Data,~sample: Model -> Context -> Nat -> F32,model: Model,combined: Tile.HeadContext, 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 22 def 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,+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 36 def 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,+index: Nat 43 ) -> {Block.read_pair(~Native,~Tile.HeadContext,~F32,~Tile.read_with_head(~Native,~Context,~read),~1n,pack(model),context, 44 index) == (pack(model),ReadLaws.pair_sample(~Model,~Tile.HeadContext,~F32,~head_sample(~Model,~Context,~sample),~1n,model,context, 45 index)) : Native & Block.Twin}: 46 ReadLaws.read_pair_refines(~Native,~Model,~Tile.HeadContext,~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 51 def 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,+index: Nat 58 ) -> {Block.read_pair(~Native,~Tile.HeadContext,~Block.Twin,~Block.read_pair(~Native, 59 ~Tile.HeadContext,~F32,~Tile.read_with_head(~Native,~Context,~read),~1n),~2n,pack(model),context,index) == (pack(model), 60 ReadLaws.pair_sample(~Model,~Tile.HeadContext,~Block.Twin,~ReadLaws.pair_sample(~Model,~Tile.HeadContext, 61 ~F32,~head_sample(~Model,~Context,~sample),~1n),~2n,model,context,index)) : Native & Block.Twin>}: 62 ReadLaws.read_pair_refines(~Native,~Model,~Tile.HeadContext,~Block.Twin,~pack,~Block.read_pair(~Native, 63 ~Tile.HeadContext,~F32,~Tile.read_with_head(~Native,~Context,~read),~1n),~ReadLaws.pair_sample(~Model,~Tile.HeadContext, 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 68 def 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,+index: Nat 75 ) -> {Block.read_pair(~Native,~Tile.HeadContext,~Block.Twin>, 76 ~Block.read_pair(~Native,~Tile.HeadContext,~Block.Twin,~Block.read_pair(~Native, 77 ~Tile.HeadContext,~F32,~Tile.read_with_head(~Native,~Context,~read),~1n),~2n),~4n,pack(model),context,index) == (pack(model), 78 ReadLaws.pair_sample(~Model,~Tile.HeadContext,~Block.Twin>,~ReadLaws.pair_sample(~Model, 79 ~Tile.HeadContext,~Block.Twin,~ReadLaws.pair_sample(~Model,~Tile.HeadContext,~F32,~head_sample(~Model, 80 ~Context,~sample),~1n),~2n),~4n,model,context,index)) : Native & Block.Twin>>}: 81 ReadLaws.read_pair_refines(~Native,~Model,~Tile.HeadContext,~Block.Twin>,~pack, 82 ~Block.read_pair(~Native,~Tile.HeadContext,~Block.Twin,~Block.read_pair(~Native, 83 ~Tile.HeadContext,~F32,~Tile.read_with_head(~Native,~Context,~read),~1n),~2n),~ReadLaws.pair_sample(~Model, 84 ~Tile.HeadContext,~Block.Twin,~ReadLaws.pair_sample(~Model,~Tile.HeadContext,~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 89 def 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,~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,~Block.Twin>, 102 ~Block.read_pair(~Native,~Tile.HeadContext,~Block.Twin,~Block.read_pair(~Native, 103 ~Tile.HeadContext,~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. 108 def 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 131 def 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 154 def 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 194 def 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 198 def 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 {==}