1 # CHW / NCHW entry. Geometry and allocation checks finish before planning. 2 # Every convolution runs on row bands: the input arrives in the layout the 3 # planner chooses for it and the band pipeline runs there. 4 import Base 5 import ./tensor.bend as Tensors 6 import ./storage_buffer.bend as Storage 7 import ./kernel_shape.bend as KernelShape 8 import ./kernel_gemm_buffer.bend as GemmBuffer 9 import ./kernel_tile.bend as Tile 10 import ./kernel_planning.bend as KernelPlanning 11 import ./default_kernel_cost.bend as DefaultCost 12 import ./default_machine_limits.bend as Default 13 import ./machine_limits.bend as Machine 14 import ./proofs/storage_certificate.bend as Certificate 15 import ./tensor_band.bend as Bands 16 import ./tensor_band_window.bend as BandWindow 17 import ./kernel_band_pipeline.bend as BandPipeline 18 import ./kernel_band_planning.bend as BandPlanning 19 import ./tensor_band_schedule.bend as BandSchedule 20 import ./math_activation.bend as Activations 21 22 type InputShape is Data: 23 InputShape{batch: U32,batched: Bool,channels: U32,height: U32,width: U32} 24 25 type Geometry is Data: 26 Geometry{batch: U32,batched: Bool,dimensions: KernelShape.Dimensions} 27 28 def input_shape(shape: List<&2,U32>) -> Maybe<&2,InputShape>: 29 match shape: 30 case [channels,height,width]: Some{InputShape{1,False{},channels,height,width}} 31 case [batch,channels,height,width]: Some{InputShape{batch,True{},channels,height,width}} 32 case shape: None{} 33 34 def matched(matches: Bool,geometry: Geometry) -> Maybe<&2,Geometry>: 35 match matches: 36 case True{}: Some{geometry} 37 case False{}: None{} 38 39 def geometry(input: Maybe<&2,InputShape>,weights: List<&2,U32>,bias: List<&2,U32>,stride: U32,padding: U32) -> Maybe<&2,Geometry>: 40 match input weights bias: 41 case Some{InputShape{batch,batched,+channels,height,width}} Con{+rows,Con{weight_channels,Con{+kernel,Con{kernel_width,Nil{}}}}} Con{bias_channels,Nil{}}: 42 matched(U32.is_eq(channels,weight_channels) && U32.is_eq(kernel,kernel_width) && U32.is_eq(rows,bias_channels), 43 Geometry{batch,batched,KernelShape.Dimensions{U32.to_nat(channels),U32.to_nat(rows),U32.to_nat(height),U32.to_nat(width),U32.to_nat(kernel),U32.to_nat(stride),U32.to_nat(padding)}}) 44 case input weights bias: None{} 45 46 def output_shape(batched: Bool,batch: U32,rows: U32,height: U32,width: U32) -> List<&2,U32>: 47 match batched: 48 case False{}: [rows,height,width] 49 case True{}: [batch,rows,height,width] 50 51 def band_context(geometry: Geometry,parallel: Bool,+activation: Activations.Activation) -> BandPipeline.Context: 52 Geometry{count,batched,KernelShape.Dimensions{+ci,rows,+height,+width,+kernel,+stride,+padding}} = geometry 53 BandPipeline.Context{BandWindow.Geometry{Bands.Geometry{count,U32.from_nat(ci),U32.from_nat(height),U32.from_nat(width)},U32.from_nat(kernel),U32.from_nat(stride),U32.from_nat(padding), 54 U32.from_nat(KernelShape.axis_size(height,kernel,stride,padding)),U32.from_nat(KernelShape.axis_size(width,kernel,stride,padding))},U32.from_nat(rows),U32.from_nat(Nat.mul(ci,Nat.mul(kernel,kernel))),parallel,activation} 55 56 def band_work(context: BandPipeline.Context,bias_capacity: U32) -> KernelPlanning.Work: 57 BandPipeline.Context{BandWindow.Geometry{Bands.Geometry{batch,+channels,+height,+width},kernel,stride,padding,oh,+ow},rows,inner,parallel,activation} = context 58 KernelPlanning.Work{(oh * ow : U32),inner,rows,KernelPlanning.Owners{U32.from_nat(KernelPlanning.capacity((channels * height * width : U32))),bias_capacity}, 59 KernelPlanning.Geometry{channels,height,width,kernel,stride,padding,ow},activation} 60 61 def context_window(context: BandPipeline.Context) -> BandWindow.Geometry: 62 BandPipeline.Context{window,rows,inner,parallel,activation} = context 63 window 64 65 def band_output(~custom: F32 -> F32,+activation: Activations.Activation,+geometry: Geometry,parallel: Bool,input: Bands.Bands(~F32),weights: Storage.Buffer,bias: Storage.Buffer,workers: U32) -> Tensors.Tensor: 66 Geometry{+count,batched,KernelShape.Dimensions{ci,+rows,height,width,+kernel,+stride,+padding}} = geometry 67 +oh = U32.from_nat(KernelShape.axis_size(height,kernel,stride,padding)) 68 +ow = U32.from_nat(KernelShape.axis_size(width,kernel,stride,padding)) 69 Tensors.Tensor{BandPipeline.execute(~custom,band_context(geometry,parallel,activation),input,weights,bias), 70 Bands.Geometry{count,U32.from_nat(rows),oh,ow},output_shape(batched,count,U32.from_nat(rows),oh,ow),0.0,Tensors.Usable{},workers} 71 72 # The weights in the order the kernel reads them: once, after the checks. 73 def reordered(dimensions: KernelShape.Dimensions,weights: Storage.Buffer) -> Storage.Buffer: 74 KernelShape.Dimensions{ci,rows,h,w,+k,s,p} = dimensions 75 GemmBuffer.reorder(U32.from_nat(rows),U32.from_nat(Nat.mul(ci,Nat.mul(k,k))),weights) 76 77 def packed(geometry: Geometry,weights: Storage.Buffer) -> Storage.Buffer: 78 Geometry{count,batched,dimensions} = geometry 79 reordered(dimensions,weights) 80 81 def machine_support(geometry: Geometry,+limits: Machine.Limits) -> Bool: 82 Geometry{+count,batched,+dimensions} = geometry 83 KernelShape.Dimensions{+ci,+rows,+h,+w,+k,+s,+p} = dimensions 84 KernelShape.supports_packed(limits,dimensions,Nat.mul(ci,Nat.mul(h,w)),Nat.mul(rows,Nat.mul(ci,Nat.mul(k,k)))) && 85 Nat.is_le(Nat.mul(U32.to_nat(count),Nat.mul(rows,Nat.mul(KernelShape.axis_size(h,k,s,p),KernelShape.axis_size(w,k,s,p)))),U32.to_nat(Machine.capacity(limits))) 86 87 def input_geometry(geometry: Geometry) -> Bands.Geometry: 88 Geometry{count,batched,KernelShape.Dimensions{ci,rows,h,w,k,s,p}} = geometry 89 Bands.Geometry{count,U32.from_nat(ci),U32.from_nat(h),U32.from_nat(w)} 90 91 # Arrived bands cover every row of the images of the input. 92 def covered(+dimensions: KernelShape.Dimensions,+count: U32,inspection: Bands.Inspection) -> Bool: 93 KernelShape.Dimensions{ci,rows,+h,w,k,s,p} = dimensions 94 Bands.Inspection{valid,first,end,leaves} = inspection 95 valid && U32.is_zero(first) && Nat.is_eq(U32.to_nat(end),Nat.mul(U32.to_nat(count),h)) 96 97 def arrived_covered(~custom: F32 -> F32,+activation: Activations.Activation,covered: Bool,+geometry: Geometry,parallel: Bool,input: Bands.Bands(~F32), 98 weights: Storage.Buffer,bias: Storage.Buffer,+shape: List<&2,U32>,fill: F32,workers: U32) -> Tensors.Tensor: 99 match covered: 100 case True{}: band_output(~custom,activation,geometry,parallel,input,weights,bias,workers) 101 case False{}: Tensors.Tensor{input,input_geometry(geometry),shape,fill,Tensors.Failed{Tensors.Problem{"conv2d",Tensors.InvalidArgument{shape}}},workers} 102 103 def arrived(~custom: F32 -> F32,+activation: Activations.Activation,+geometry: Geometry,parallel: Bool,weights: Storage.Buffer,bias: Storage.Buffer, 104 +shape: List<&2,U32>,fill: F32,workers: U32,inspected: Bands.Bands(~F32) & Bands.Inspection) -> Tensors.Tensor: 105 Geometry{+count,batched,+dimensions} = geometry 106 (input,inspection) = inspected 107 arrived_covered(~custom,activation,covered(dimensions,count,inspection),Geometry{count,batched,dimensions},parallel,input,weights,bias,shape,fill,workers) 108 109 def selected(~custom: F32 -> F32,+activation: Activations.Activation,+geometry: Geometry,weights: Storage.Buffer,bias: Storage.Buffer, 110 +shape: List<&2,U32>,fill: F32,workers: U32,tensor: Tensors.Tensor,selection: BandPlanning.LayoutPlan) -> Tensors.Tensor: 111 BandPlanning.LayoutPlan{layout,parallel,time,repartition} = selection 112 +input = input_geometry(geometry) 113 arrived(~custom,activation,geometry,parallel,weights,bias,shape,fill,workers,Bands.inspect(~F32,input,Tensors.arrive(~F32,input,layout,tensor))) 114 115 # The planner chooses the layout again from the input's plan when the input 116 # holds the input geometry, otherwise it chooses the first split. 117 def offered(~custom: F32 -> F32,+activation: Activations.Activation,+geometry: Geometry,weights: Storage.Buffer,bias: Storage.Buffer, 118 +shape: List<&2,U32>,fill: F32,+workers: U32,owners: Tensors.Tensor & BandSchedule.Target) -> Tensors.Tensor: 119 Storage.Buffer{bias,length,Certificate.Certificate{+depth,equation}} = bias 120 (tensor,target) = owners 121 BandSchedule.Target{other,+previous,split} = target 122 +context = band_context(geometry,True{},activation) 123 selected(~custom,activation,geometry,weights,Storage.Buffer{bias,length,Certificate.Certificate{depth,equation}},shape,fill,workers,tensor, 124 BandPlanning.layout_for(Bands.same_geometry(input_geometry(geometry),other),previous,band_work(context,U32.shln(1,depth)),context_window(context),workers, 125 DefaultCost.configuration())) 126 127 # A shape the machine cannot hold. 128 def unsupported(tensor: Tensors.Tensor) -> Tensors.Tensor: 129 Tensors.Tensor{bands,geometry,+shape,fill,status,workers} = tensor 130 Tensors.Tensor{bands,geometry,shape,fill,Tensors.Failed{Tensors.Problem{"conv2d",Tensors.InvalidShape{shape}}},workers} 131 132 def band_machine(~custom: F32 -> F32,+activation: Activations.Activation,supported: Bool,+geometry: Geometry,tensor: Tensors.Tensor, 133 weights: Storage.Buffer,bias: Storage.Buffer,+shape: List<&2,U32>,fill: F32,workers: U32) -> Tensors.Tensor: 134 match supported: 135 case True{}: offered(~custom,activation,geometry,packed(geometry,weights),bias,shape,fill,workers,Tensors.offer(~F32,tensor)) 136 case False{}: unsupported(tensor) 137 138 def band_supported(~custom: F32 -> F32,+activation: Activations.Activation,valid: Bool,+geometry: Geometry,tensor: Tensors.Tensor, 139 weights: Storage.Buffer,bias: Storage.Buffer,+shape: List<&2,U32>,fill: F32,workers: U32) -> Tensors.Tensor: 140 match valid: 141 case False{}: Tensors.rejected(~F32,"conv2d",tensor) 142 case True{}: 143 band_machine(~custom,activation,machine_support(geometry,Default.limits()),geometry,tensor,weights,bias,shape,fill,workers) 144 145 # What a call needs of its parameters: lengths that fit the shape and the 146 # capacities that hold them. 147 def parameters_valid(+dimensions: KernelShape.Dimensions,+weight_length: Nat,+weight_depth: Nat,+bias_length: Nat,+bias_depth: Nat) -> Bool: 148 KernelShape.Dimensions{+ci,rows,+h,+w,k,s,p} = dimensions 149 KernelShape.valid(dimensions,Nat.mul(ci,Nat.mul(h,w)),weight_length,bias_length) && 150 Nat.is_le(weight_length,U32.to_nat(U32.shln(1,weight_depth))) && Nat.is_le(bias_length,U32.to_nat(U32.shln(1,bias_depth))) 151 152 def band_checked(~custom: F32 -> F32,+activation: Activations.Activation,geometry: Maybe<&2,Geometry>,tensor: Tensors.Tensor, 153 weights: Storage.Buffer,bias: Storage.Buffer,+shape: List<&2,U32>,fill: F32,workers: U32) -> Tensors.Tensor: 154 match geometry: 155 case None{}: Tensors.rejected(~F32,"conv2d",tensor) 156 case Some{+geometry}: 157 Geometry{count,batched,+dimensions} = geometry 158 Storage.Buffer{weights,+weight_length,Certificate.Certificate{+weight_depth,weight_equation}} = weights 159 Storage.Buffer{bias,+bias_length,Certificate.Certificate{+bias_depth,bias_equation}} = bias 160 band_supported(~custom,activation,parameters_valid(dimensions,weight_length,weight_depth,bias_length,bias_depth),geometry,tensor, 161 Storage.Buffer{weights,weight_length,Certificate.Certificate{weight_depth,weight_equation}}, 162 Storage.Buffer{bias,bias_length,Certificate.Certificate{bias_depth,bias_equation}},shape,fill,workers) 163 164 def band_status(~custom: F32 -> F32,+activation: Activations.Activation,status: Tensors.Status,layout: Bands.Geometry,input: Bands.Bands(~F32),+shape: List<&2,U32>,+fill: F32, 165 weights: Storage.Buffer,weight_shape: List<&2,U32>,bias: Storage.Buffer,bias_shape: List<&2,U32>,stride: U32,padding: U32,+workers: U32) -> Tensors.Tensor: 166 match status: 167 case Tensors.Failed{problem}: Tensors.Tensor{input,layout,shape,fill,Tensors.Failed{problem},workers} 168 case Tensors.Usable{}: 169 band_checked(~custom,activation,geometry(input_shape(shape),weight_shape,bias_shape,stride,padding),Tensors.Tensor{input,layout,shape,fill,Tensors.Usable{},workers}, 170 weights,bias,shape,fill,workers) 171 172 def input_layout(~custom: F32 -> F32,+activation: Activations.Activation,input: Tensors.Tensor,weights: Storage.Buffer,weight_shape: List<&2,U32>,bias: Storage.Buffer,bias_shape: List<&2,U32>, 173 parameter_status: Tensors.Status,parameter_workers: U32,stride: U32,padding: U32) -> Tensors.Tensor: 174 Tensors.Tensor{input,layout,+shape,fill,input_status,workers} = input 175 band_status(~custom,activation,Tensors.first_failure(input_status,parameter_status),layout,input,shape,fill,weights,weight_shape,bias,bias_shape,stride,padding, 176 U32.max(workers,parameter_workers)) 177 178 # Weights and bias are read whole, so band parameters are gathered first. 179 def conv2d(~custom: F32 -> F32,+activation: Activations.Activation,input: Tensors.Tensor,weights: Tensors.Tensor,bias: Tensors.Tensor,stride: U32,padding: U32) -> Tensors.Tensor: 180 Tensors.with_whole(~F32,~Tensors.Tensor,weights,weights => weight_shape => weight_fill => weight_status => weight_workers => 181 Tensors.with_whole(~F32,~Tensors.Tensor,bias,bias => bias_shape => bias_fill => bias_status => bias_workers => 182 input_layout(~custom,activation,input,weights,weight_shape,bias,bias_shape,Tensors.first_failure(weight_status,bias_status),U32.max(weight_workers,bias_workers),stride,padding)))