tensor_convolution.bend
stelliferous@0.0.3.0 · 182 lines · 0 definitions
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.
4import Base
5import stelliferous@0.0.3.0/tensor.bend as Tensors
6import stelliferous@0.0.3.0/storage_buffer.bend as Storage
7import stelliferous@0.0.3.0/kernel_shape.bend as KernelShape
8import stelliferous@0.0.3.0/kernel_gemm_buffer.bend as GemmBuffer
9import stelliferous@0.0.3.0/kernel_tile.bend as Tile
10import stelliferous@0.0.3.0/kernel_planning.bend as KernelPlanning
11import stelliferous@0.0.3.0/default_kernel_cost.bend as DefaultCost
12import stelliferous@0.0.3.0/default_machine_limits.bend as Default
13import stelliferous@0.0.3.0/machine_limits.bend as Machine
14import stelliferous@0.0.3.0/proofs/storage_certificate.bend as Certificate
15import stelliferous@0.0.3.0/tensor_band.bend as Bands
16import stelliferous@0.0.3.0/tensor_band_window.bend as BandWindow
17import stelliferous@0.0.3.0/kernel_band_pipeline.bend as BandPipeline
18import stelliferous@0.0.3.0/kernel_band_planning.bend as BandPlanning
19import stelliferous@0.0.3.0/tensor_band_schedule.bend as BandSchedule
20import stelliferous@0.0.3.0/math_activation.bend as Activations
21
22type InputShape is Data:
23 InputShape{batch: U32,batched: Bool,channels: U32,height: U32,width: U32}
24
25type Geometry is Data:
26 Geometry{batch: U32,batched: Bool,dimensions: KernelShape.Dimensions}
27
28def 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
34def matched(matches: Bool,geometry: Geometry) -> Maybe<&2,Geometry>:
35 match matches:
36 case True{}: Some{geometry}
37 case False{}: None{}
38
39def 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
46def 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
51def 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
56def 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
61def context_window(context: BandPipeline.Context) -> BandWindow.Geometry:
62 BandPipeline.Context{window,rows,inner,parallel,activation} = context
63 window
64
65def band_output(~custom: F32 -> F32,+activation: Activations.Activation,+geometry: Geometry,parallel: Bool,input: Bands.Bands(~F32),weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,workers: U32) -> Tensors.Tensor<F32>:
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.
73def reordered(dimensions: KernelShape.Dimensions,weights: Storage.Buffer<F32>) -> Storage.Buffer<Tile.Tile8>:
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
77def packed(geometry: Geometry,weights: Storage.Buffer<F32>) -> Storage.Buffer<Tile.Tile8>:
78 Geometry{count,batched,dimensions} = geometry
79 reordered(dimensions,weights)
80
81def 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
87def 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.
92def 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
97def arrived_covered(~custom: F32 -> F32,+activation: Activations.Activation,covered: Bool,+geometry: Geometry,parallel: Bool,input: Bands.Bands(~F32),
98 weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,+shape: List<&2,U32>,fill: F32,workers: U32) -> Tensors.Tensor<F32>:
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
103def arrived(~custom: F32 -> F32,+activation: Activations.Activation,+geometry: Geometry,parallel: Bool,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,
104 +shape: List<&2,U32>,fill: F32,workers: U32,inspected: Bands.Bands(~F32) & Bands.Inspection) -> Tensors.Tensor<F32>:
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
109def selected(~custom: F32 -> F32,+activation: Activations.Activation,+geometry: Geometry,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,
110 +shape: List<&2,U32>,fill: F32,workers: U32,tensor: Tensors.Tensor<F32>,selection: BandPlanning.LayoutPlan) -> Tensors.Tensor<F32>:
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.
117def offered(~custom: F32 -> F32,+activation: Activations.Activation,+geometry: Geometry,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,
118 +shape: List<&2,U32>,fill: F32,+workers: U32,owners: Tensors.Tensor<F32> & BandSchedule.Target) -> Tensors.Tensor<F32>:
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.
128def unsupported(tensor: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
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
132def band_machine(~custom: F32 -> F32,+activation: Activations.Activation,supported: Bool,+geometry: Geometry,tensor: Tensors.Tensor<F32>,
133 weights: Storage.Buffer<F32>,bias: Storage.Buffer<F32>,+shape: List<&2,U32>,fill: F32,workers: U32) -> Tensors.Tensor<F32>:
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
138def band_supported(~custom: F32 -> F32,+activation: Activations.Activation,valid: Bool,+geometry: Geometry,tensor: Tensors.Tensor<F32>,
139 weights: Storage.Buffer<F32>,bias: Storage.Buffer<F32>,+shape: List<&2,U32>,fill: F32,workers: U32) -> Tensors.Tensor<F32>:
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.
147def 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
152def band_checked(~custom: F32 -> F32,+activation: Activations.Activation,geometry: Maybe<&2,Geometry>,tensor: Tensors.Tensor<F32>,
153 weights: Storage.Buffer<F32>,bias: Storage.Buffer<F32>,+shape: List<&2,U32>,fill: F32,workers: U32) -> Tensors.Tensor<F32>:
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
164def 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<F32>,weight_shape: List<&2,U32>,bias: Storage.Buffer<F32>,bias_shape: List<&2,U32>,stride: U32,padding: U32,+workers: U32) -> Tensors.Tensor<F32>:
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
172def input_layout(~custom: F32 -> F32,+activation: Activations.Activation,input: Tensors.Tensor<F32>,weights: Storage.Buffer<F32>,weight_shape: List<&2,U32>,bias: Storage.Buffer<F32>,bias_shape: List<&2,U32>,
173 parameter_status: Tensors.Status,parameter_workers: U32,stride: U32,padding: U32) -> Tensors.Tensor<F32>:
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.
179def conv2d(~custom: F32 -> F32,+activation: Activations.Activation,input: Tensors.Tensor<F32>,weights: Tensors.Tensor<F32>,bias: Tensors.Tensor<F32>,stride: U32,padding: U32) -> Tensors.Tensor<F32>:
180 Tensors.with_whole(~F32,~Tensors.Tensor<F32>,weights,weights => weight_shape => weight_fill => weight_status => weight_workers =>
181 Tensors.with_whole(~F32,~Tensors.Tensor<F32>,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)))