1 import Base 2 import ../kernel_spatial_packing.bend as Packing 3 import ../kernel_packing_samples.bend as PackingSamples 4 import ./storage_witness.bend as Witness 5 import ./traversal_witness.bend as TraversalWitness 6 import ../kernel_tile.bend as Tile 7 import ./kernel_tile_witness.bend as TileWitness 8 9 def family(-input: Array) -> Data: 10 Witness.ArrayWitness 11 12 def masked(-input: Array,+inside: Bool,-result: Array & F32,observed: Witness.Read) -> 13 TraversalWitness.Sample,F32,input,PackingSamples.mask_pixel(inside,result)>: 14 Witness.Read{value,equation} = observed 15 %Equal.sym(Array & F32,result,(input,value),equation) : TraversalWitness.Sample,F32,input,PackingSamples.mask_pixel(inside,_)> 16 TraversalWitness.Sample{PackingSamples.masked(value,inside),{==}} 17 18 def masked_read(-input: Array,+inside: Bool,-result: Array & F32,observed: Witness.Read) -> 19 TraversalWitness.Sample,F32,input,Packing.masked(inside,result)>: 20 Witness.Read{value,equation} = observed 21 %Equal.sym(Array & F32,result,(input,value),equation) : TraversalWitness.Sample,F32,input,Packing.masked(inside,_)> 22 TraversalWitness.Sample{PackingSamples.masked(value,inside),{==}} 23 24 def pixel(-input: Array,+row: U32,+column: U32,+channel: U32,+height: U32,+width: U32,evidence: family(input)) -> 25 TraversalWitness.Sample,F32,input,Packing.pixel(input,row,column,channel,height,width)>: 26 +index = (channel + row * width + column : U32) 27 masked_read(input,U32.is_lt(row,height) && U32.is_lt(column,width),Array.get(F32,input,index),Witness.observed(F32,input,index,evidence)) 28 29 def row_read(-input: Array,context: Packing.PixelBlock,index: Nat,evidence: family(input)) -> 30 TraversalWitness.Sample,F32,input,Packing.read_block_pixel(input,context,index)>: 31 Packing.PixelBlock{row,column,channel,height,width,stride} = context 32 pixel(input,row,(column + U32.from_nat(index) * stride : U32),channel,height,width,evidence) 33 34 def row_block(-input: Array,+row: U32,+column: U32,+channel: U32,+height: U32,+width: U32,+stride: U32,+next_row: U32,+offset: U32,+output_width: U32,+kernel_column: U32,+padding: U32,+evidence: family(input)) -> 35 TraversalWitness.Sample,Tile.Tile8,input,Packing.block(input,row,column,channel,height,width,stride,next_row,offset,output_width,kernel_column,padding)>: 36 TileWitness.tail(~Array,~Packing.PixelBlock,~family,~Packing.read_block_pixel,~row_read,input,Packing.PixelBlock{row,column,channel,height,width,stride}, 37 Packing.pixel(input,row,column,channel,height,width),evidence,pixel(input,row,column,channel,height,width,evidence)) 38 39 def wrap_pixel(same: Bool,-input: Array,+row: U32,+column: U32,+channel: U32,+height: U32,+width: U32,+stride: U32,+next_row: U32,+offset: U32,+output_width: U32,+kernel_column: U32,+padding: U32,+lane: U32,evidence: family(input)) -> 40 TraversalWitness.Sample,F32,input,Packing.wrap_pixel(same,input,row,column,channel,height,width,stride,next_row,offset,output_width,kernel_column,padding,lane)>: 41 match same: 42 case True{}: pixel(input,row,(column + lane * stride : U32),channel,height,width,evidence) 43 case False{}: pixel(input,next_row,((offset + lane - output_width) * stride + kernel_column - padding : U32),channel,height,width,evidence) 44 45 def wrap_lane(-input: Array,+row: U32,+column: U32,+channel: U32,+height: U32,+width: U32,+stride: U32,+next_row: U32,+offset: U32,+output_width: U32,+kernel_column: U32,+padding: U32,+lane: U32,evidence: family(input)) -> 46 TraversalWitness.Sample,F32,input,Packing.wrap_lane(input,row,column,channel,height,width,stride,next_row,offset,output_width,kernel_column,padding,lane)>: 47 wrap_pixel(U32.is_lt((offset + lane : U32),output_width),input,row,column,channel,height,width,stride,next_row,offset,output_width,kernel_column,padding,lane,evidence) 48 49 def wrap_read(-input: Array,context: Packing.WrapBlock,index: Nat,evidence: family(input)) -> 50 TraversalWitness.Sample,F32,input,Packing.read_wrap_pixel(input,context,index)>: 51 Packing.WrapBlock{+row,+column,+channel,+height,+width,+stride,+next_row,+offset,+output_width,+kernel_column,+padding} = context 52 wrap_lane(input,row,column,channel,height,width,stride,next_row,offset,output_width,kernel_column,padding,U32.from_nat(index),evidence) 53 54 def wrap_block(-input: Array,+row: U32,+column: U32,+channel: U32,+height: U32,+width: U32,+stride: U32,+next_row: U32,+offset: U32,+output_width: U32,+kernel_column: U32,+padding: U32,+evidence: family(input)) -> 55 TraversalWitness.Sample,Tile.Tile8,input,Packing.wrap_block(input,row,column,channel,height,width,stride,next_row,offset,output_width,kernel_column,padding)>: 56 TileWitness.tail(~Array,~Packing.WrapBlock,~family,~Packing.read_wrap_pixel,~wrap_read,input,Packing.WrapBlock{row,column,channel,height,width,stride,next_row,offset,output_width,kernel_column,padding}, 57 Packing.pixel(input,row,column,channel,height,width),evidence,pixel(input,row,column,channel,height,width,evidence)) 58 59 def address(-input: Array,address: PackingSamples.PixelAddress,evidence: family(input)) -> 60 TraversalWitness.Sample,F32,input,PackingSamples.read_pixel_address(input,address)>: 61 PackingSamples.PixelAddress{+index,inside} = address 62 masked(input,inside,Array.get(F32,input,index),Witness.observed(F32,input,index,evidence)) 63 64 def crossing_read(-input: Array,context: PackingSamples.PixelBlock,index: Nat,evidence: family(input)) -> 65 TraversalWitness.Sample,F32,input,PackingSamples.read_block_pixel(input,context,index)>: 66 address(input,PackingSamples.pixel_address(context,index),evidence) 67 68 def first(-input: Array,+position: U32,+reduction: U32,+height: U32,+width: U32,+kernel: U32,+stride: U32,+padding: U32,+output_width: U32,evidence: family(input)) -> 69 TraversalWitness.Sample,F32,input,PackingSamples.pixel(input,position,reduction,height,width,kernel,stride,padding,output_width)>: 70 +row = (position / output_width * stride + reduction / kernel % kernel - padding : U32) 71 +column = (position % output_width * stride + reduction % kernel - padding : U32) 72 +index = (reduction / (kernel * kernel) * height * width + row * width + column : U32) 73 masked(input,U32.is_lt(row,height) && U32.is_lt(column,width),Array.get(F32,input,index),Witness.observed(F32,input,index,evidence))