1 # Checked operations preserve all owners and their balanced-storage certificates. 2 # Element functions enter hot loops only as direct templates. 3 import Base 4 import ./storage_buffer.bend as Storage 5 import ./tensor_shape.bend as Shapes 6 import ./tensor_view.bend as Views 7 import ./traversal_array.bend as Traversal 8 import ./tensor_access.bend as Access 9 import ./proofs/storage_certificate.bend as Certificate 10 import ./proofs/traversal_certificate.bend as TraversalCertificate 11 import ./proofs/traversal_witness.bend as TraversalWitness 12 import ./proofs/tensor_access_certificate.bend as AccessCertificate 13 import ./traversal_lanes.bend as Lanes 14 import ./proofs/traversal_lanes_certificate.bend as LanesCertificate 15 16 type Outcome<-Value: Type> is Type: 17 Accepted{value: Value} 18 Rejected{value: Value} 19 20 def valid_buffer(+length: Nat,capacity: U32,view: Views.View) -> Bool: 21 Nat.is_le(length,U32.to_nat(capacity)) && Views.fits(view,U32.from_nat(length)) 22 23 def pair_buffers(~Element: Data,input_length: Nat,output_length: Nat,result: Array & Array, 24 certificates: TraversalCertificate.Certificates,Element,Certificate.family(~Element),result>) -> 25 Storage.Buffer & Storage.Buffer: 26 (input,output) = result 27 TraversalCertificate.Certificates{input_certificate,output_certificate} = certificates 28 (Storage.Buffer{input,input_length,input_certificate},Storage.Buffer{output,output_length,output_certificate}) 29 30 def copy_when(~Element: Data,valid: Bool,count: Maybe<&2,U32>,+source: Views.View,target: Views.View, 31 input: Storage.Buffer,output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 32 match valid count: 33 case True{} Some{+count}: 34 Views.View{axes,+offset} = target 35 Storage.Buffer{input,input_length,input_certificate} = input 36 Storage.Buffer{output,output_length,output_certificate} = output 37 +contiguous = Views.is_contiguous(source) 38 Accepted{pair_buffers(~Element,input_length,output_length, 39 Access.copy_run(~Element,contiguous,U32.to_nat(count),source,offset,input,output), 40 AccessCertificate.copy(~Element,contiguous,U32.to_nat(count),source,offset,input,output,input_certificate,output_certificate))} 41 case valid count: Rejected{(input,output)} 42 43 def copy(~Element: Data,input: Storage.Buffer,+source: Views.View, 44 output: Storage.Buffer,+target: Views.View) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 45 Storage.Buffer{input,+input_length,input_certificate} = input 46 Certificate.Certificate{+input_depth,input_equation} = input_certificate 47 Storage.Buffer{output,+output_length,output_certificate} = output 48 Certificate.Certificate{+output_depth,output_equation} = output_certificate 49 copy_when(~Element,Shapes.same(Views.shape(source),Views.shape(target)) && Views.is_contiguous(target) && 50 valid_buffer(input_length,U32.shln(1,input_depth),source) && 51 valid_buffer(output_length,U32.shln(1,output_depth),target),Views.numel(target),source,target, 52 Storage.Buffer{input,input_length,Certificate.Certificate{input_depth,input_equation}},Storage.Buffer{output,output_length,Certificate.Certificate{output_depth,output_equation}}) 53 54 # output[address of element i of target] = input[offset of source + i]: a 55 # contiguous source written through a strided target view, in row-major order. 56 def scatter_when(~Element: Data,valid: Bool,count: Maybe<&2,U32>,source: Views.View,+target: Views.View, 57 input: Storage.Buffer,output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 58 match valid count: 59 case True{} Some{+count}: 60 Views.View{axes,+offset} = source 61 Storage.Buffer{input,input_length,input_certificate} = input 62 Storage.Buffer{output,output_length,output_certificate} = output 63 Accepted{pair_buffers(~Element,input_length,output_length, 64 Access.scatter_run(~Element,U32.to_nat(count),target,offset,input,output), 65 AccessCertificate.scatter(~Element,U32.to_nat(count),target,offset,input,output,input_certificate,output_certificate))} 66 case valid count: Rejected{(input,output)} 67 68 def scatter(~Element: Data,input: Storage.Buffer,+source: Views.View, 69 output: Storage.Buffer,+target: Views.View) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 70 Storage.Buffer{input,+input_length,input_certificate} = input 71 Certificate.Certificate{+input_depth,input_equation} = input_certificate 72 Storage.Buffer{output,+output_length,output_certificate} = output 73 Certificate.Certificate{+output_depth,output_equation} = output_certificate 74 scatter_when(~Element,Shapes.same(Views.shape(source),Views.shape(target)) && Views.is_contiguous(source) && 75 valid_buffer(input_length,U32.shln(1,input_depth),source) && 76 valid_buffer(output_length,U32.shln(1,output_depth),target),Views.numel(target),source,target, 77 Storage.Buffer{input,input_length,Certificate.Certificate{input_depth,input_equation}},Storage.Buffer{output,output_length,Certificate.Certificate{output_depth,output_equation}}) 78 79 def combine_buffers(~Element: Data,left_length: Nat,right_length: Nat,result: Array & Array, 80 certificates: TraversalCertificate.Certificates,Element,Certificate.family(~Element),result>) -> Storage.Buffer & Storage.Buffer: 81 (right,left) = result 82 TraversalCertificate.Certificates{right_certificate,left_certificate} = certificates 83 (Storage.Buffer{left,left_length,left_certificate},Storage.Buffer{right,right_length,right_certificate}) 84 85 def combine_kernel(~Element: Data,~operation: Element -> Element -> Element,+count: Nat,+right_view: Views.View,+offset: U32, 86 left: Storage.Buffer,right: Storage.Buffer) -> Storage.Buffer & Storage.Buffer: 87 Storage.Buffer{left,left_length,left_certificate} = left 88 Storage.Buffer{right,right_length,right_certificate} = right 89 +fast = Access.zero_contiguous(right_view) && U32.is_zero(offset) 90 combine_buffers(~Element,left_length,right_length, 91 Access.combine_run(~Element,~operation,fast,count,right_view,offset,right,left), 92 AccessCertificate.combine(~Element,~operation,fast,count,right_view,offset,right,left,right_certificate,left_certificate)) 93 94 def combine_when_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer -> Storage.Buffer -> Storage.Buffer & Storage.Buffer, 95 valid: Bool,count: Maybe<&2,U32>,right_view: Views.View,target: Views.View,left: Storage.Buffer,right: Storage.Buffer) -> 96 Outcome<(Storage.Buffer & Storage.Buffer)>: 97 match valid count: 98 case True{} Some{count}: 99 Views.View{axes,offset} = target 100 Accepted{kernel(U32.to_nat(count),right_view,offset,left,right)} 101 case valid count: Rejected{(left,right)} 102 103 def combine_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer -> Storage.Buffer -> Storage.Buffer & Storage.Buffer, 104 left: Storage.Buffer,+left_view: Views.View,right: Storage.Buffer,+right_view: Views.View) -> 105 Outcome<(Storage.Buffer & Storage.Buffer)>: 106 Storage.Buffer{left,+left_length,left_certificate} = left 107 Certificate.Certificate{+left_depth,left_equation} = left_certificate 108 Storage.Buffer{right,+right_length,right_certificate} = right 109 Certificate.Certificate{+right_depth,right_equation} = right_certificate 110 combine_when_with(~Element,~kernel,Shapes.same(Views.shape(left_view),Views.shape(right_view)) && Views.is_contiguous(left_view) && 111 valid_buffer(left_length,U32.shln(1,left_depth),left_view) && 112 valid_buffer(right_length,U32.shln(1,right_depth),right_view),Views.numel(left_view),right_view,left_view, 113 Storage.Buffer{left,left_length,Certificate.Certificate{left_depth,left_equation}},Storage.Buffer{right,right_length,Certificate.Certificate{right_depth,right_equation}}) 114 115 def combine(~Element: Data,~operation: Element -> Element -> Element, 116 left: Storage.Buffer,+left_view: Views.View,right: Storage.Buffer,+right_view: Views.View) -> 117 Outcome<(Storage.Buffer & Storage.Buffer)>: 118 combine_with(~Element,~combine_kernel(~Element,~operation),left,left_view,right,right_view) 119 120 # A mapper is the ordered map of count values from position on a certified 121 # Buffer. mapped applies the operation to one value at a time. 122 def mapped(~Element: Data,~operation: Element -> Element,+count: Nat,+position: U32,buffer: Storage.Buffer) -> Storage.Buffer: 123 Storage.Buffer{storage,length,certificate} = buffer 124 Storage.Buffer{Traversal.map_in_place(~Element,~Unit,~Access.apply_element(~Element,~operation),count,position,Unit{},storage),length, 125 TraversalCertificate.map_in_place(~Element,~Unit,~Access.apply_element(~Element,~operation),count,position,Unit{},storage,certificate)} 126 127 # The same map with the values of each round of eight computed together: an 128 # operation written with FP32 arithmetic alone then compiles to vector code. 129 # The certificate's depth decides whether the array holds the positions. 130 def lanes(~Context: Data,~operation: Context -> F32 -> F32,+count: Nat,+position: U32,+context: Context,buffer: Storage.Buffer) -> Storage.Buffer: 131 Storage.Buffer{storage,length,Certificate.Certificate{+depth,equation}} = buffer 132 Storage.Buffer{Lanes.map_when(~Context,~operation,Lanes.covered(depth,position,count),count,position,context,storage),length, 133 LanesCertificate.map_when(~Context,~operation,Lanes.covered(depth,position,count),count,position,context,storage,Certificate.Certificate{depth,equation})} 134 135 def mapped_lanes(~operation: F32 -> F32,+count: Nat,+position: U32,buffer: Storage.Buffer) -> Storage.Buffer: 136 lanes(~Unit,~Access.apply_element(~F32,~operation),count,position,Unit{},buffer) 137 138 def apply_context_when(~Element: Data,~Context: Data,~mapper: @+count: Nat -> @+position: U32 -> @+context: Context -> Storage.Buffer -> Storage.Buffer,context: Context,valid: Bool,count: Maybe<&2,U32>,view: Views.View, 139 buffer: Storage.Buffer) -> Outcome>: 140 match valid count: 141 case True{} Some{+count}: 142 Views.View{axes,+offset} = view 143 Accepted{mapper(U32.to_nat(count),offset,context,buffer)} 144 case valid count: Rejected{buffer} 145 146 def apply_context(~Element: Data,~Context: Data,~mapper: @+count: Nat -> @+position: U32 -> @+context: Context -> Storage.Buffer -> Storage.Buffer,context: Context,buffer: Storage.Buffer,+view: Views.View) -> Outcome>: 147 Storage.Buffer{storage,+length,certificate} = buffer 148 Certificate.Certificate{+depth,equation} = certificate 149 apply_context_when(~Element,~Context,~mapper,context,Views.is_contiguous(view) && valid_buffer(length,U32.shln(1,depth),view),Views.numel(view),view, 150 Storage.Buffer{storage,length,Certificate.Certificate{depth,equation}}) 151 152 # A mapper without a context, as the Unit-context case. 153 def without_context(~Element: Data,~mapper: @+count: Nat -> @+position: U32 -> Storage.Buffer -> Storage.Buffer,+count: Nat,+position: U32,+context: Unit,buffer: Storage.Buffer) -> Storage.Buffer: 154 mapper(count,position,buffer) 155 156 def apply_indexed_when(~Element: Data,~operation: U32 -> Element -> Element,valid: Bool,count: Maybe<&2,U32>,+view: Views.View, 157 input: Storage.Buffer,output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 158 match valid count: 159 case True{} Some{+count}: 160 Storage.Buffer{input,input_length,input_certificate} = input 161 Storage.Buffer{output,output_length,output_certificate} = output 162 Accepted{pair_buffers(~Element,input_length,output_length, 163 Access.apply_indexed_run(~Element,~operation,U32.to_nat(count),view,input,output), 164 AccessCertificate.apply_indexed(~Element,~operation,U32.to_nat(count),view,input,output,input_certificate,output_certificate))} 165 case valid count: Rejected{(input,output)} 166 167 def holds(count: Maybe<&2,U32>,+length: Nat,capacity: U32) -> Bool: 168 match count: 169 case None{}: False{} 170 case Some{count}: Nat.is_le(U32.to_nat(count),length) && Nat.is_le(length,U32.to_nat(capacity)) 171 172 # Writes f(index, element) for a contiguous input view into output[0..count). 173 def apply_indexed(~Element: Data,~operation: U32 -> Element -> Element,input: Storage.Buffer,+view: Views.View, 174 output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 175 Storage.Buffer{input,+input_length,input_certificate} = input 176 Certificate.Certificate{+input_depth,input_equation} = input_certificate 177 Storage.Buffer{output,+output_length,output_certificate} = output 178 Certificate.Certificate{+output_depth,output_equation} = output_certificate 179 apply_indexed_when(~Element,~operation,Views.is_contiguous(view) && valid_buffer(input_length,U32.shln(1,input_depth),view) && 180 holds(Views.numel(view),output_length,U32.shln(1,output_depth)),Views.numel(view),view, 181 Storage.Buffer{input,input_length,Certificate.Certificate{input_depth,input_equation}},Storage.Buffer{output,output_length,Certificate.Certificate{output_depth,output_equation}}) 182 183 def small(value: U32) -> Bool: 184 U32.is_le(value,65535) 185 186 # The window geometry is exact and every sample address stays in U32. 187 def pool_geometry(+pool: Access.Pool) -> Bool: 188 Access.Pool{+height,+width,+kernel,+stride,+padding,+out_height,+out_width} = pool 189 small(height) && small(width) && small(kernel) && small(stride) && small(padding) && U32.is_lt(0,kernel) && U32.is_lt(0,stride) && 190 U32.is_le(kernel,(height + 2 * padding : U32)) && U32.is_le(kernel,(width + 2 * padding : U32)) && 191 U32.is_eq(out_height,((height + 2 * padding - kernel : U32) / stride + 1 : U32)) && 192 U32.is_eq(out_width,((width + 2 * padding - kernel : U32) / stride + 1 : U32)) 193 194 # The unchecked sweep of two owners; the callers establish its bounds. 195 def swept(~Element: Data,~operation: Element -> Element -> Element,+count: Nat,+context: Access.SweepContext, 196 input: Storage.Buffer,output: Storage.Buffer) -> Storage.Buffer & Storage.Buffer: 197 Storage.Buffer{input,input_length,input_certificate} = input 198 Storage.Buffer{output,output_length,output_certificate} = output 199 pair_buffers(~Element,input_length,output_length, 200 Access.sweep_run(~Element,~operation,count,context,input,output), 201 AccessCertificate.sweep(~Element,~operation,count,context,input,output,input_certificate,output_certificate)) 202 203 def sweep_when(~Element: Data,~operation: Element -> Element -> Element,valid: Bool,count: Maybe<&2,U32>,+context: Access.SweepContext, 204 input: Storage.Buffer,output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 205 match valid count: 206 case True{} Some{count}: Accepted{swept(~Element,~operation,U32.to_nat(count),context,input,output)} 207 case valid count: Rejected{(input,output)} 208 209 # The window geometry of a whole axis is exact and every sample address stays in U32. 210 def sweep_geometry(+sweep: Access.Sweep) -> Bool: 211 Access.Sweep{+extent,+step,+kernel,+stride,+padding,+out_extent,+origin} = sweep 212 small(extent) && small(kernel) && small(stride) && small(padding) && U32.is_lt(0,kernel) && U32.is_lt(0,stride) && U32.is_lt(0,step) && 213 U32.is_eq(origin,0) && U32.is_le(kernel,(extent + 2 * padding : U32)) && 214 U32.is_eq(out_extent,((extent + 2 * padding - kernel : U32) / stride + 1 : U32)) 215 216 def sweep_input(+sweep: Access.Sweep,+outer: U32) -> Maybe<&2,U32>: 217 Access.Sweep{+extent,+step,kernel,stride,padding,out_extent,origin} = sweep 218 Shapes.product([outer,extent,step],Some{1}) 219 220 def sweep_output(+sweep: Access.Sweep,+outer: U32) -> Maybe<&2,U32>: 221 Access.Sweep{extent,+step,kernel,stride,padding,+out_extent,origin} = sweep 222 Shapes.product([outer,out_extent,step],Some{1}) 223 224 # Windows along one axis of outer contiguous blocks; see Access.Sweep. 225 def sweep(~Element: Data,~operation: Element -> Element -> Element,+pad: Element,input: Storage.Buffer,+sweep: Access.Sweep,+outer: U32, 226 output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 227 Storage.Buffer{input,+input_length,input_certificate} = input 228 Certificate.Certificate{+input_depth,input_equation} = input_certificate 229 Storage.Buffer{output,+output_length,output_certificate} = output 230 Certificate.Certificate{+output_depth,output_equation} = output_certificate 231 sweep_when(~Element,~operation,sweep_geometry(sweep) && holds(sweep_input(sweep,outer),input_length,U32.shln(1,input_depth)) && 232 holds(sweep_output(sweep,outer),output_length,U32.shln(1,output_depth)),sweep_output(sweep,outer),Access.SweepContext{sweep,pad}, 233 Storage.Buffer{input,input_length,Certificate.Certificate{input_depth,input_equation}},Storage.Buffer{output,output_length,Certificate.Certificate{output_depth,output_equation}}) 234 235 # The input and the output of a pool around its column-swept planes. 236 def pool_finished(~Element: Data,input: Storage.Buffer,result: Outcome<(Storage.Buffer & Storage.Buffer)>) -> 237 Outcome<(Storage.Buffer & Storage.Buffer)>: 238 match result: 239 case Accepted{(middle,output)}: Accepted{(input,output)} 240 case Rejected{(middle,output)}: Rejected{(input,output)} 241 242 def pool_rows(~Element: Data,~operation: Element -> Element -> Element,+pad: Element,result: Outcome<(Storage.Buffer & Storage.Buffer)>, 243 +rows: Access.Sweep,+channels: U32,output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 244 match result: 245 case Accepted{(input,middle)}: pool_finished(~Element,input,sweep(~Element,~operation,pad,middle,rows,channels,output)) 246 case Rejected{(input,middle)}: Rejected{(input,output)} 247 248 def pool_lines(~Element: Data,~operation: Element -> Element -> Element,+pad: Element,lines: Maybe<&2,U32>,input: Storage.Buffer,+pool: Access.Pool, 249 +channels: U32,middle: Storage.Buffer,output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 250 match lines: 251 case None{}: Rejected{(input,output)} 252 case Some{+lines}: 253 pool_rows(~Element,~operation,pad,sweep(~Element,~operation,pad,input,Access.pool_columns(pool),lines,middle),Access.pool_rows(pool,0),channels,output) 254 255 # Windows over the last two axes of channels contiguous planes: a sweep along 256 # the columns of every row into middle, [channels,height,out_width], then a 257 # sweep along the rows into output. Each sweep checks its own geometry and 258 # counts. 259 def pool(~Element: Data,~operation: Element -> Element -> Element,+pad: Element,input: Storage.Buffer,+pool: Access.Pool,+channels: U32, 260 middle: Storage.Buffer,output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 261 Access.Pool{+height,+width,+kernel,+stride,+padding,+out_height,+out_width} = pool 262 pool_lines(~Element,~operation,pad,Shapes.product([channels,height],Some{1}),input,Access.Pool{height,width,kernel,stride,padding,out_height,out_width},channels,middle,output) 263 264 def reduce_buffer(~Element: Data,~Accumulator: Data,length: Nat,result: Array & Accumulator, 265 certificate: Certificate.Certificate,~Accumulator,result)>) -> Storage.Buffer & Accumulator: 266 (storage,value) = result 267 (Storage.Buffer{storage,length,certificate},value) 268 269 def reduce_when(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator, 270 valid: Bool,count: Maybe<&2,U32>,+view: Views.View,buffer: Storage.Buffer,initial: Accumulator) -> 271 Outcome<(Storage.Buffer & Accumulator)>: 272 match valid count: 273 case True{} Some{+count}: 274 Storage.Buffer{storage,length,certificate} = buffer 275 +fast = Views.is_contiguous(view) 276 Accepted{reduce_buffer(~Element,~Accumulator,length, 277 Access.reduce_run(~Element,~Accumulator,~operation,fast,U32.to_nat(count),view,storage,initial), 278 AccessCertificate.reduce(~Element,~Accumulator,~operation,fast,U32.to_nat(count),view,storage,initial,certificate))} 279 case valid count: Rejected{(buffer,initial)} 280 281 def reduce(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator, 282 buffer: Storage.Buffer,+view: Views.View,initial: Accumulator) -> Outcome<(Storage.Buffer & Accumulator)>: 283 Storage.Buffer{storage,+length,certificate} = buffer 284 Certificate.Certificate{+depth,equation} = certificate 285 reduce_when(~Element,~Accumulator,~operation,valid_buffer(length,U32.shln(1,depth),view),Views.numel(view),view, 286 Storage.Buffer{storage,length,Certificate.Certificate{depth,equation}},initial)