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 ./kernel_tile.bend as Tile 15 import ./proofs/traversal_lanes_certificate.bend as LanesCertificate 16 17 type Outcome<-Value: Type> is Type: 18 Accepted{value: Value} 19 Rejected{value: Value} 20 21 def valid_buffer(+length: Nat,capacity: U32,view: Views.View) -> Bool: 22 Nat.is_le(length,U32.to_nat(capacity)) && Views.fits(view,U32.from_nat(length)) 23 24 def pair_buffers(~Element: Data,input_length: Nat,output_length: Nat,result: Array & Array, 25 certificates: TraversalCertificate.Certificates,Element,Certificate.family(~Element),result>) -> 26 Storage.Buffer & Storage.Buffer: 27 (input,output) = result 28 TraversalCertificate.Certificates{input_certificate,output_certificate} = certificates 29 (Storage.Buffer{input,input_length,input_certificate},Storage.Buffer{output,output_length,output_certificate}) 30 31 def copy_when(~Element: Data,valid: Bool,count: Maybe<&2,U32>,+source: Views.View,target: Views.View, 32 input: Storage.Buffer,output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 33 match valid count: 34 case True{} Some{+count}: 35 Views.View{axes,+offset} = target 36 Storage.Buffer{input,input_length,input_certificate} = input 37 Storage.Buffer{output,output_length,output_certificate} = output 38 +contiguous = Views.is_contiguous(source) 39 Accepted{pair_buffers(~Element,input_length,output_length, 40 Access.copy_run(~Element,contiguous,U32.to_nat(count),source,offset,input,output), 41 AccessCertificate.copy(~Element,contiguous,U32.to_nat(count),source,offset,input,output,input_certificate,output_certificate))} 42 case valid count: Rejected{(input,output)} 43 44 # output[offset of target + i] = element i of source read with the flagged 45 # axes in reverse, in row-major order of the result; an axis without a flag 46 # is read forward. The guard checks the source itself, and the copy runs 47 # through its mirrored view. 48 def copy(~Element: Data,input: Storage.Buffer,+source: Views.View,+reversed: List<&2,Bool>, 49 output: Storage.Buffer,+target: Views.View) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 50 Storage.Buffer{input,+input_length,input_certificate} = input 51 Certificate.Certificate{+input_depth,input_equation} = input_certificate 52 Storage.Buffer{output,+output_length,output_certificate} = output 53 Certificate.Certificate{+output_depth,output_equation} = output_certificate 54 copy_when(~Element,Shapes.same(Views.shape(source),Views.shape(target)) && Views.is_contiguous(target) && 55 valid_buffer(input_length,U32.shln(1,input_depth),source) && 56 valid_buffer(output_length,U32.shln(1,output_depth),target),Views.numel(target),Views.mirrored(source,reversed),target, 57 Storage.Buffer{input,input_length,Certificate.Certificate{input_depth,input_equation}},Storage.Buffer{output,output_length,Certificate.Certificate{output_depth,output_equation}}) 58 59 # output[address of element i of target] = input[offset of source + i]: a 60 # contiguous source written through a strided target view, in row-major order. 61 def scatter_counted(~Element: Data,+count: Nat,+target: Views.View,+offset: U32, 62 input: Storage.Buffer,output: Storage.Buffer) -> Storage.Buffer & Storage.Buffer: 63 Storage.Buffer{input,input_length,input_certificate} = input 64 Storage.Buffer{output,output_length,output_certificate} = output 65 pair_buffers(~Element,input_length,output_length, 66 Access.scatter_run(~Element,count,target,offset,input,output), 67 AccessCertificate.scatter(~Element,count,target,offset,input,output,input_certificate,output_certificate)) 68 69 def scatter_when(~Element: Data,valid: Bool,count: Maybe<&2,U32>,source: Views.View,+target: Views.View, 70 input: Storage.Buffer,output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 71 match valid count: 72 case True{} Some{+count}: 73 Views.View{axes,+offset} = source 74 Accepted{scatter_counted(~Element,U32.to_nat(count),target,offset,input,output)} 75 case valid count: Rejected{(input,output)} 76 77 def scatter(~Element: Data,input: Storage.Buffer,+source: Views.View, 78 output: Storage.Buffer,+target: Views.View) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 79 Storage.Buffer{input,+input_length,input_certificate} = input 80 Certificate.Certificate{+input_depth,input_equation} = input_certificate 81 Storage.Buffer{output,+output_length,output_certificate} = output 82 Certificate.Certificate{+output_depth,output_equation} = output_certificate 83 scatter_when(~Element,Shapes.same(Views.shape(source),Views.shape(target)) && Views.is_contiguous(source) && 84 valid_buffer(input_length,U32.shln(1,input_depth),source) && 85 valid_buffer(output_length,U32.shln(1,output_depth),target),Views.numel(target),source,target, 86 Storage.Buffer{input,input_length,Certificate.Certificate{input_depth,input_equation}},Storage.Buffer{output,output_length,Certificate.Certificate{output_depth,output_equation}}) 87 88 def combine_buffers(~Element: Data,left_length: Nat,right_length: Nat,result: Array & Array, 89 certificates: TraversalCertificate.Certificates,Element,Certificate.family(~Element),result>) -> Storage.Buffer & Storage.Buffer: 90 (right,left) = result 91 TraversalCertificate.Certificates{right_certificate,left_certificate} = certificates 92 (Storage.Buffer{left,left_length,left_certificate},Storage.Buffer{right,right_length,right_certificate}) 93 94 def combine_kernel(~Element: Data,~operation: Element -> Element -> Element,+count: Nat,+right_view: Views.View,+offset: U32, 95 left: Storage.Buffer,right: Storage.Buffer) -> Storage.Buffer & Storage.Buffer: 96 Storage.Buffer{left,left_length,left_certificate} = left 97 Storage.Buffer{right,right_length,right_certificate} = right 98 +fast = Access.zero_contiguous(right_view) && U32.is_zero(offset) 99 combine_buffers(~Element,left_length,right_length, 100 Access.combine_run(~Element,~operation,fast,count,right_view,offset,right,left), 101 AccessCertificate.combine(~Element,~operation,fast,count,right_view,offset,right,left,right_certificate,left_certificate)) 102 103 def combine_when_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer -> Storage.Buffer -> Storage.Buffer & Storage.Buffer, 104 valid: Bool,count: Maybe<&2,U32>,right_view: Views.View,target: Views.View,left: Storage.Buffer,right: Storage.Buffer) -> 105 Outcome<(Storage.Buffer & Storage.Buffer)>: 106 match valid count: 107 case True{} Some{count}: 108 Views.View{axes,offset} = target 109 Accepted{kernel(U32.to_nat(count),right_view,offset,left,right)} 110 case valid count: Rejected{(left,right)} 111 112 def combine_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer -> Storage.Buffer -> Storage.Buffer & Storage.Buffer, 113 left: Storage.Buffer,+left_view: Views.View,right: Storage.Buffer,+right_view: Views.View) -> 114 Outcome<(Storage.Buffer & Storage.Buffer)>: 115 Storage.Buffer{left,+left_length,left_certificate} = left 116 Certificate.Certificate{+left_depth,left_equation} = left_certificate 117 Storage.Buffer{right,+right_length,right_certificate} = right 118 Certificate.Certificate{+right_depth,right_equation} = right_certificate 119 combine_when_with(~Element,~kernel,Shapes.same(Views.shape(left_view),Views.shape(right_view)) && Views.is_contiguous(left_view) && 120 valid_buffer(left_length,U32.shln(1,left_depth),left_view) && 121 valid_buffer(right_length,U32.shln(1,right_depth),right_view),Views.numel(left_view),right_view,left_view, 122 Storage.Buffer{left,left_length,Certificate.Certificate{left_depth,left_equation}},Storage.Buffer{right,right_length,Certificate.Certificate{right_depth,right_equation}}) 123 124 def combine(~Element: Data,~operation: Element -> Element -> Element, 125 left: Storage.Buffer,+left_view: Views.View,right: Storage.Buffer,+right_view: Views.View) -> 126 Outcome<(Storage.Buffer & Storage.Buffer)>: 127 combine_with(~Element,~combine_kernel(~Element,~operation),left,left_view,right,right_view) 128 129 # A mapper is the ordered map of count values from position on a certified 130 # Buffer. mapped applies the operation to one value at a time. 131 def mapped(~Element: Data,~operation: Element -> Element,+count: Nat,+position: U32,buffer: Storage.Buffer) -> Storage.Buffer: 132 Storage.Buffer{storage,length,certificate} = buffer 133 Storage.Buffer{Traversal.map_in_place(~Element,~Unit,~Access.apply_element(~Element,~operation),count,position,Unit{},storage),length, 134 TraversalCertificate.map_in_place(~Element,~Unit,~Access.apply_element(~Element,~operation),count,position,Unit{},storage,certificate)} 135 136 # The same map with the values of each round of eight computed together: an 137 # operation written with FP32 arithmetic alone then compiles to vector code. 138 # The certificate's depth decides whether the array holds the positions. 139 def lanes(~Context: Data,~operation: Context -> F32 -> F32,+count: Nat,+position: U32,+context: Context,buffer: Storage.Buffer) -> Storage.Buffer: 140 Storage.Buffer{storage,length,Certificate.Certificate{+depth,equation}} = buffer 141 Storage.Buffer{Lanes.map_when(~Context,~operation,Lanes.covered(depth,position,count),count,position,context,storage),length, 142 LanesCertificate.map_when(~Context,~operation,Lanes.covered(depth,position,count),count,position,context,storage,Certificate.Certificate{depth,equation})} 143 144 def mapped_lanes(~operation: F32 -> F32,+count: Nat,+position: U32,buffer: Storage.Buffer) -> Storage.Buffer: 145 lanes(~Unit,~Access.apply_element(~F32,~operation),count,position,Unit{},buffer) 146 147 # A mapper without a context, as the Unit-context case. 148 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: 149 mapper(count,position,buffer) 150 151 def holds(count: Maybe<&2,U32>,+length: Nat,capacity: U32) -> Bool: 152 match count: 153 case None{}: False{} 154 case Some{count}: Nat.is_le(U32.to_nat(count),length) && Nat.is_le(length,U32.to_nat(capacity)) 155 156 def small(value: U32) -> Bool: 157 U32.is_le(value,65535) 158 159 # The window geometry is exact and every sample address stays in U32. 160 def pool_geometry(+pool: Access.Pool) -> Bool: 161 Access.Pool{+height,+width,+kernel,+stride,+padding,+out_height,+out_width} = pool 162 small(height) && small(width) && small(kernel) && small(stride) && small(padding) && U32.is_lt(0,kernel) && U32.is_lt(0,stride) && 163 U32.is_le(kernel,(height + 2 * padding : U32)) && U32.is_le(kernel,(width + 2 * padding : U32)) && 164 U32.is_eq(out_height,((height + 2 * padding - kernel : U32) / stride + 1 : U32)) && 165 U32.is_eq(out_width,((width + 2 * padding - kernel : U32) / stride + 1 : U32)) 166 167 # The unchecked sweep of two owners; the callers establish its bounds. 168 def swept(~Element: Data,~operation: Element -> Element -> Element,+count: Nat,+context: Access.SweepContext, 169 input: Storage.Buffer,output: Storage.Buffer) -> Storage.Buffer & Storage.Buffer: 170 Storage.Buffer{input,input_length,input_certificate} = input 171 Storage.Buffer{output,output_length,output_certificate} = output 172 pair_buffers(~Element,input_length,output_length, 173 Access.sweep_run(~Element,~operation,count,context,input,output), 174 AccessCertificate.sweep(~Element,~operation,count,context,input,output,input_certificate,output_certificate)) 175 176 def sweep_when(~Element: Data,~operation: Element -> Element -> Element,valid: Bool,count: Maybe<&2,U32>,+context: Access.SweepContext, 177 input: Storage.Buffer,output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 178 match valid count: 179 case True{} Some{count}: Accepted{swept(~Element,~operation,U32.to_nat(count),context,input,output)} 180 case valid count: Rejected{(input,output)} 181 182 # Every window fits the padded axis and every sample address stays in U32. 183 # The axis may be a window of rows of a taller image, whose first output 184 # window starts at origin. 185 def sweep_geometry(+sweep: Access.Sweep) -> Bool: 186 Access.Sweep{+extent,+step,+kernel,+stride,+padding,out_extent,+origin} = sweep 187 small(extent) && small(kernel) && small(stride) && small(padding) && small(origin) && U32.is_lt(0,kernel) && U32.is_lt(0,stride) && 188 U32.is_lt(0,step) && U32.is_le(kernel,(extent + 2 * padding : U32)) 189 190 def sweep_input(+sweep: Access.Sweep,+outer: U32) -> Maybe<&2,U32>: 191 Access.Sweep{+extent,+step,kernel,stride,padding,out_extent,origin} = sweep 192 Shapes.product([outer,extent,step],Some{1}) 193 194 def sweep_output(+sweep: Access.Sweep,+outer: U32) -> Maybe<&2,U32>: 195 Access.Sweep{extent,+step,kernel,stride,padding,+out_extent,origin} = sweep 196 Shapes.product([outer,out_extent,step],Some{1}) 197 198 # Windows along one axis of outer contiguous blocks; see Access.Sweep. 199 def sweep(~Element: Data,~operation: Element -> Element -> Element,+pad: Element,input: Storage.Buffer,+sweep: Access.Sweep,+outer: U32, 200 output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 201 Storage.Buffer{input,+input_length,input_certificate} = input 202 Certificate.Certificate{+input_depth,input_equation} = input_certificate 203 Storage.Buffer{output,+output_length,output_certificate} = output 204 Certificate.Certificate{+output_depth,output_equation} = output_certificate 205 sweep_when(~Element,~operation,sweep_geometry(sweep) && holds(sweep_input(sweep,outer),input_length,U32.shln(1,input_depth)) && 206 holds(sweep_output(sweep,outer),output_length,U32.shln(1,output_depth)),sweep_output(sweep,outer),Access.SweepContext{sweep,pad}, 207 Storage.Buffer{input,input_length,Certificate.Certificate{input_depth,input_equation}},Storage.Buffer{output,output_length,Certificate.Certificate{output_depth,output_equation}}) 208 209 # The input and the output of a pool around its column-swept planes. 210 def pool_finished(~Element: Data,input: Storage.Buffer,result: Outcome<(Storage.Buffer & Storage.Buffer)>) -> 211 Outcome<(Storage.Buffer & Storage.Buffer)>: 212 match result: 213 case Accepted{(middle,output)}: Accepted{(input,output)} 214 case Rejected{(middle,output)}: Rejected{(input,output)} 215 216 def pool_rows(~Element: Data,~operation: Element -> Element -> Element,+pad: Element,result: Outcome<(Storage.Buffer & Storage.Buffer)>, 217 +rows: Access.Sweep,+channels: U32,output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 218 match result: 219 case Accepted{(input,middle)}: pool_finished(~Element,input,sweep(~Element,~operation,pad,middle,rows,channels,output)) 220 case Rejected{(input,middle)}: Rejected{(input,output)} 221 222 def pool_lines(~Element: Data,~operation: Element -> Element -> Element,+pad: Element,lines: Maybe<&2,U32>,input: Storage.Buffer,+pool: Access.Pool, 223 +origin: U32,+channels: U32,middle: Storage.Buffer,output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 224 match lines: 225 case None{}: Rejected{(input,output)} 226 case Some{+lines}: 227 pool_rows(~Element,~operation,pad,sweep(~Element,~operation,pad,input,Access.pool_columns(pool),lines,middle),Access.pool_rows(pool,origin),channels,output) 228 229 # Windows over the last two axes of channels contiguous planes: a sweep along 230 # the columns of every row into middle, [channels,height,out_width], then a 231 # sweep along the rows into output. Each sweep checks its own geometry and 232 # counts. The stored rows may be a window of a taller image; origin is then 233 # the row where the window of the first output row starts (see Access.Sweep). 234 def pool(~Element: Data,~operation: Element -> Element -> Element,+pad: Element,input: Storage.Buffer,+pool: Access.Pool,+origin: U32,+channels: U32, 235 middle: Storage.Buffer,output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 236 Access.Pool{+height,+width,+kernel,+stride,+padding,+out_height,+out_width} = pool 237 pool_lines(~Element,~operation,pad,Shapes.product([channels,height],Some{1}),input,Access.Pool{height,width,kernel,stride,padding,out_height,out_width},origin,channels,middle,output) 238 239 # The axes of a line before its last axis, and the last axis; a line of no 240 # axis is one value. 241 def prepend_run(axis: Views.Axis,runs: List<&2,Views.Axis> & Views.Axis) -> List<&2,Views.Axis> & Views.Axis: 242 (outer,last) = runs 243 (axis <> outer,last) 244 245 def runs_after(rest: List<&2,Views.Axis>,axis: Views.Axis) -> List<&2,Views.Axis> & Views.Axis: 246 match rest: 247 case []: ([],axis) 248 case next <> tail: prepend_run(axis,runs_after(tail,next)) 249 250 def line_runs(line: List<&2,Views.Axis>) -> List<&2,Views.Axis> & Views.Axis: 251 match line: 252 case []: ([],Views.Axis{1,0}) 253 case axis <> rest: runs_after(rest,axis) 254 255 def elements_of(axes: List<&2,Views.Axis>) -> Nat: 256 match axes: 257 case []: 1n 258 case Views.Axis{extent,stride} <> rest: Nat.mul(U32.to_nat(extent),elements_of(rest)) 259 260 # The lines, unchecked; the callers establish the bounds. 261 def lined(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element, 262 +count: Nat,+view: Views.View,+runs: Nat,+extent: U32,+stride: U32,+initial: Accumulator,input: Storage.Buffer,output: Storage.Buffer) -> 263 Storage.Buffer & Storage.Buffer: 264 Storage.Buffer{input,input_length,input_certificate} = input 265 Storage.Buffer{output,output_length,output_certificate} = output 266 pair_buffers(~Element,input_length,output_length, 267 Access.reduce_lines(~Element,~Accumulator,~operation,~finish,count,view,runs,extent,stride,initial,input,output), 268 AccessCertificate.lines(~Element,~Accumulator,~operation,~finish,count,view,runs,extent,stride,initial,input,output,input_certificate,output_certificate)) 269 270 def runs_lined(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element, 271 +count: Nat,+kept: List<&2,Views.Axis>,+offset: U32,runs: List<&2,Views.Axis> & Views.Axis,+initial: Accumulator,input: Storage.Buffer, 272 output: Storage.Buffer) -> Storage.Buffer & Storage.Buffer: 273 (+outer,Views.Axis{+extent,+stride}) = runs 274 lined(~Element,~Accumulator,~operation,~finish,count,Views.View{List.append(&2,Views.Axis,kept,outer),offset},elements_of(outer),extent,stride,initial,input,output) 275 276 def lines_when(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element, 277 valid: Bool,count: Maybe<&2,U32>,+kept: List<&2,Views.Axis>,+line: List<&2,Views.Axis>,+offset: U32,+initial: Accumulator, 278 input: Storage.Buffer,output: Storage.Buffer) -> Outcome<(Storage.Buffer & Storage.Buffer)>: 279 match valid count: 280 case True{} Some{+count}: 281 Accepted{runs_lined(~Element,~Accumulator,~operation,~finish,U32.to_nat(count),kept,offset,line_runs(line),initial,input,output)} 282 case valid count: Rejected{(input,output)} 283 284 # The lines of a view: line i holds the values of the last reduced axes at 285 # index i of the axes before them. Output i folds line i in row-major order 286 # from initial and holds finish of the fold. The view must fit the input and 287 # the outputs the output. 288 def reduce_lines(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,~finish: Accumulator -> Element, 289 input: Storage.Buffer,+axes: List<&2,Views.Axis>,+reduced: Nat,+offset: U32,+initial: Accumulator,output: Storage.Buffer) -> 290 Outcome<(Storage.Buffer & Storage.Buffer)>: 291 Storage.Buffer{input,+input_length,input_certificate} = input 292 Certificate.Certificate{+input_depth,input_equation} = input_certificate 293 Storage.Buffer{output,+output_length,output_certificate} = output 294 Certificate.Certificate{+output_depth,output_equation} = output_certificate 295 +kept_count = Nat.sub(List.length(&2,Views.Axis,axes),reduced) 296 +kept = List.take(&2,Views.Axis,axes,kept_count) 297 source = {Views.View{axes,offset} : Views.View} 298 outputs = {Views.View{kept,0} : Views.View} 299 counted = {Views.View{kept,0} : Views.View} 300 lines_when(~Element,~Accumulator,~operation,~finish, 301 valid_buffer(input_length,U32.shln(1,input_depth),source) && holds(Views.numel(outputs),output_length,U32.shln(1,output_depth)), 302 Views.numel(counted),kept,List.drop(&2,Views.Axis,axes,kept_count),offset,initial, 303 Storage.Buffer{input,input_length,Certificate.Certificate{input_depth,input_equation}},Storage.Buffer{output,output_length,Certificate.Certificate{output_depth,output_equation}}) 304 305 # A staged tile transform uses the same certified buffer map and scalar tail. 306 def lanes_tiles(~Context: Data,~operation: Context -> F32 -> F32,~transform: Context -> Tile.Tile8 -> Tile.Tile8,+count: Nat,+position: U32,+context: Context,buffer: Storage.Buffer) -> Storage.Buffer: 307 Storage.Buffer{storage,length,Certificate.Certificate{+depth,equation}} = buffer 308 Storage.Buffer{Lanes.map_with_when(~Context,~operation,~Lanes.gather_transformed(~Context,~transform),Lanes.covered(depth,position,count),count,position,context,storage),length, 309 LanesCertificate.transformed_map_when(~Context,~operation,~transform,Lanes.covered(depth,position,count),count,position,context,storage,Certificate.Certificate{depth,equation})}