tensor_convolution.bend

stelliferous@0.0.2.0 · 354 lines · 53 definitions

1# CHW / NCHW entry. Geometry and allocation checks finish before planning.
2# Every image uses the same reordered weights; the batch is processed in order.
3import Base
4import stelliferous@0.0.2.0/tensor.bend as Tensors
5import stelliferous@0.0.2.0/storage_buffer.bend as Storage
6import stelliferous@0.0.2.0/storage_clone.bend as Clone
7import stelliferous@0.0.2.0/storage_copy.bend as Copy
8import stelliferous@0.0.2.0/kernel_shape.bend as KernelShape
9import stelliferous@0.0.2.0/kernel_pipeline.bend as Pipeline
10import stelliferous@0.0.2.0/kernel_gemm_buffer.bend as GemmBuffer
11import stelliferous@0.0.2.0/kernel_tile.bend as Tile
12import stelliferous@0.0.2.0/kernel_planning.bend as KernelPlanning
13import stelliferous@0.0.2.0/default_kernel_cost.bend as DefaultCost
14import stelliferous@0.0.2.0/kernel_cost.bend as Cost
15import stelliferous@0.0.2.0/default_machine_limits.bend as Default
16import stelliferous@0.0.2.0/machine_limits.bend as Machine
17import stelliferous@0.0.2.0/proofs/storage_certificate.bend as Certificate
18import stelliferous@0.0.2.0/tensor_band.bend as Bands
19import stelliferous@0.0.2.0/tensor_band_window.bend as BandWindow
20import stelliferous@0.0.2.0/kernel_band_pipeline.bend as BandPipeline
21import stelliferous@0.0.2.0/kernel_band_planning.bend as BandPlanning
22import stelliferous@0.0.2.0/math_activation.bend as Activations
23
24type InputShape is Data:
25  InputShape{batch: U32,batched: Bool,channels: U32,height: U32,width: U32}
26
27type Geometry is Data:
28  Geometry{batch: U32,batched: Bool,dimensions: KernelShape.Dimensions}
29
30type DenseArguments is Data:
31  DenseArguments{image_count: U32,work: Pipeline.Work,input_size: U32,output_size: U32,output_extents: List<&2,U32>,selected_depth: U32}
32
33def input_shape(shape: List<&2,U32>) -> Maybe<&2,InputShape>:
34  match shape:
35    case [channels,height,width]: Some{InputShape{1,False{},channels,height,width}}
36    case [batch,channels,height,width]: Some{InputShape{batch,True{},channels,height,width}}
37    case shape: None{}
38
39def matched(matches: Bool,geometry: Geometry) -> Maybe<&2,Geometry>:
40  match matches:
41    case True{}: Some{geometry}
42    case False{}: None{}
43
44def geometry(input: Maybe<&2,InputShape>,weights: List<&2,U32>,bias: List<&2,U32>,stride: U32,padding: U32) -> Maybe<&2,Geometry>:
45  match input weights bias:
46    case Some{InputShape{batch,batched,+channels,height,width}} Con{+rows,Con{weight_channels,Con{+kernel,Con{kernel_width,Nil{}}}}} Con{bias_channels,Nil{}}:
47      matched(U32.is_eq(channels,weight_channels) && U32.is_eq(kernel,kernel_width) && U32.is_eq(rows,bias_channels),
48        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)}})
49    case input weights bias: None{}
50
51def output_shape(batched: Bool,batch: U32,rows: U32,height: U32,width: U32) -> List<&2,U32>:
52  match batched:
53    case False{}: [rows,height,width]
54    case True{}: [batch,rows,height,width]
55
56def selected_depth(plan: KernelPlanning.Plan) -> U32:
57  KernelPlanning.Plan{packed,depth,workspace,time} = plan
58  depth
59
60def configured_depth(work: Pipeline.Work,input_capacity: U32,bias_capacity: U32,workers: U32,+activation: Activations.Activation,+config: Cost.Configuration) -> U32:
61  Pipeline.Work{rows,height,width,+kernel,stride,padding,output_width,positions,+inner} = work
62  # As in the pipeline split, input channels are the reduction length over kernel taps.
63  +geometry = {KernelPlanning.Geometry{(inner / (kernel * kernel) : U32),height,width,kernel,stride,padding,output_width} : KernelPlanning.Geometry}
64  selected_depth(KernelPlanning.select(KernelPlanning.Work{positions,inner,rows,KernelPlanning.Owners{input_capacity,bias_capacity},geometry,activation},workers,None{},config))
65
66# The input tensor's workers bound the split, as a thread pool bounds PyTorch's parallel_for.
67def depth(work: Pipeline.Work,input_capacity: U32,bias_capacity: U32,workers: U32,+activation: Activations.Activation) -> U32:
68  configured_depth(work,input_capacity,bias_capacity,workers,activation,DefaultCost.configuration())
69
70def copied_output(pair: Storage.Buffer<F32> & Storage.Buffer<F32>) -> Storage.Buffer<F32>:
71  (source,output) = pair
72  output
73
74def image(~custom: F32 -> F32,+activation: Activations.Activation,depth: U32,work: Pipeline.Work,index: U32,+input_size: U32,input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) ->
75  Storage.Buffer<F32> & Storage.Buffer<F32>:
76  Storage.bind_pair(~Storage.Buffer<F32>,~Storage.Buffer<F32>,~(Storage.Buffer<F32> & Storage.Buffer<F32>),
77    Copy.range(~F32,input_size,(index * input_size : U32),0,input,Tensors.buffer(~F32,0.0,input_size)),source => image =>
78      (source,Pipeline.run_packed(~custom,activation,depth,work,image,weights,bias)))
79
80def batch(~custom: F32 -> F32,+activation: Activations.Activation,count: Nat,+index: U32,+depth: U32,+work: Pipeline.Work,+input_size: U32,+output_size: U32,
81  input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,output: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
82  match count:
83    case 0n: output
84    case 1n:
85      Storage.bind_pair(~Storage.Buffer<F32>,~Storage.Buffer<F32>,~Storage.Buffer<F32>,
86        image(~custom,activation,depth,work,index,input_size,input,weights,bias),source => result =>
87          copied_output(Copy.range(~F32,output_size,0,(index * output_size : U32),result,output)))
88    case 2n+rest:
89      Storage.bind_pair(~Storage.Buffer<Tile.Tile8>,~Storage.Buffer<Tile.Tile8>,~Storage.Buffer<F32>,Clone.clone(~Tile.Tile8,weights),current_weights => remaining_weights =>
90        Storage.bind_pair(~Storage.Buffer<F32>,~Storage.Buffer<F32>,~Storage.Buffer<F32>,Clone.clone(~F32,bias),current_bias => remaining_bias =>
91          Storage.bind_pair(~Storage.Buffer<F32>,~Storage.Buffer<F32>,~Storage.Buffer<F32>,
92            image(~custom,activation,depth,work,index,input_size,input,current_weights,current_bias),source => result =>
93              batch(~custom,activation,1n+rest,(index + 1 : U32),depth,work,input_size,output_size,source,remaining_weights,remaining_bias,
94                copied_output(Copy.range(~F32,output_size,0,(index * output_size : U32),result,output))))))
95
96def images_selected(~custom: F32 -> F32,+activation: Activations.Activation,empty: Bool,single: Bool,+count: U32,depth: U32,work: Pipeline.Work,input_size: U32,+output_size: U32,
97  input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
98  match empty:
99    case True{}: Tensors.buffer(~F32,0.0,0)
100    case False{}:
101      match single:
102        case True{}: Pipeline.run_packed(~custom,activation,depth,work,input,weights,bias)
103        case False{}:
104          batch(~custom,activation,U32.to_nat(count),0,depth,work,input_size,output_size,
105            input,weights,bias,Tensors.buffer(~F32,0.0,(count * output_size : U32)))
106
107def images(~custom: F32 -> F32,+activation: Activations.Activation,+count: U32,depth: U32,work: Pipeline.Work,input_size: U32,+output_size: U32,
108  input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
109  images_selected(~custom,activation,U32.is_eq(count,0),U32.is_eq(count,1),count,depth,work,input_size,output_size,input,weights,bias)
110
111def execute_dense_at_depth(~custom: F32 -> F32,+activation: Activations.Activation,+image_count: U32,+selected_depth: U32,+work: Pipeline.Work,+input_size: U32,+output_size: U32,+output_extents: List<&2,U32>,
112  input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,+workers: U32) -> Tensors.Tensor<F32>:
113  Tensors.Tensor{images(~custom,activation,image_count,selected_depth,work,input_size,output_size,input,weights,bias),
114    output_extents,0.0,Tensors.Usable{},workers}
115
116def dense_arguments(geometry: Geometry,+bias_depth: Nat,+workers: U32,+activation: Activations.Activation) -> DenseArguments:
117  Geometry{+count,batched,KernelShape.Dimensions{+ci,+rows,+height,+width,+kernel,+stride,+padding}} = geometry
118  +oh = U32.from_nat(KernelShape.axis_size(height,kernel,stride,padding))
119  +ow = U32.from_nat(KernelShape.axis_size(width,kernel,stride,padding))
120  +input_size = U32.from_nat(Nat.mul(ci,Nat.mul(height,width)))
121  +output_size = (U32.from_nat(rows) * oh * ow : U32)
122  +work = {Pipeline.Work{U32.from_nat(rows),U32.from_nat(height),U32.from_nat(width),U32.from_nat(kernel),U32.from_nat(stride),U32.from_nat(padding),ow,(oh * ow : U32),U32.from_nat(Nat.mul(ci,Nat.mul(kernel,kernel)))} : Pipeline.Work}
123  DenseArguments{count,work,input_size,output_size,output_shape(batched,count,U32.from_nat(rows),oh,ow),
124    depth(work,U32.from_nat(KernelPlanning.capacity(input_size)),U32.shln(1,bias_depth),workers,activation)}
125
126def execute_dense_arguments(~custom: F32 -> F32,+activation: Activations.Activation,arguments: DenseArguments,input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,+workers: U32) -> Tensors.Tensor<F32>:
127  DenseArguments{count,work,input_size,output_size,output_extents,selected_depth} = arguments
128  execute_dense_at_depth(~custom,activation,count,selected_depth,work,input_size,output_size,output_extents,input,weights,bias,workers)
129
130def execute_dense(~custom: F32 -> F32,+activation: Activations.Activation,geometry: Geometry,input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,+workers: U32) -> Tensors.Tensor<F32>:
131  Storage.Buffer{bias,bias_length,Certificate.Certificate{+bias_depth,bias_equation}} = bias
132  execute_dense_arguments(~custom,activation,dense_arguments(geometry,bias_depth,workers,activation),input,weights,Storage.Buffer{bias,bias_length,Certificate.Certificate{bias_depth,bias_equation}},workers)
133
134def band_context(geometry: Geometry,parallel: Bool,+activation: Activations.Activation) -> BandPipeline.Context:
135  Geometry{count,batched,KernelShape.Dimensions{+ci,rows,+height,+width,+kernel,+stride,+padding}} = geometry
136  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),
137    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}
138
139def band_work(context: BandPipeline.Context,bias_capacity: U32) -> KernelPlanning.Work:
140  BandPipeline.Context{BandWindow.Geometry{Bands.Geometry{batch,+channels,+height,+width},kernel,stride,padding,oh,+ow},rows,inner,parallel,activation} = context
141  KernelPlanning.Work{(oh * ow : U32),inner,rows,KernelPlanning.Owners{U32.from_nat(KernelPlanning.capacity((channels * height * width : U32))),bias_capacity},
142    KernelPlanning.Geometry{channels,height,width,kernel,stride,padding,ow},activation}
143
144def context_window(context: BandPipeline.Context) -> BandWindow.Geometry:
145  BandPipeline.Context{window,rows,inner,parallel,activation} = context
146  window
147
148def 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>:
149  Geometry{+count,batched,KernelShape.Dimensions{ci,+rows,height,width,+kernel,+stride,+padding}} = geometry
150  +oh = U32.from_nat(KernelShape.axis_size(height,kernel,stride,padding))
151  +ow = U32.from_nat(KernelShape.axis_size(width,kernel,stride,padding))
152  Tensors.Banded{BandPipeline.execute(~custom,band_context(geometry,parallel,activation),input,weights,bias),
153    Bands.Geometry{count,U32.from_nat(rows),oh,ow},output_shape(batched,count,U32.from_nat(rows),oh,ow),0.0,Tensors.Usable{},workers}
154
155def start_bands(~custom: F32 -> F32,+activation: Activations.Activation,depth: U32,parallel: Bool,+geometry: Geometry,input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,workers: U32) -> Tensors.Tensor<F32>:
156  Geometry{count,batched,KernelShape.Dimensions{ci,rows,h,w,k,s,p}} = geometry
157  band_output(~custom,activation,geometry,parallel,Bands.split(~F32,U32.to_nat(depth),Bands.Geometry{count,U32.from_nat(ci),U32.from_nat(h),U32.from_nat(w)},0.0,input),weights,bias,workers)
158
159def initial_layout_selected(~custom: F32 -> F32,+activation: Activations.Activation,zero: Bool,depth: U32,parallel: Bool,+geometry: Geometry,input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,workers: U32) -> Tensors.Tensor<F32>:
160  match zero:
161    case True{}: Tensors.with_workers(~F32,execute_dense(~custom,activation,geometry,input,weights,bias,1),workers)
162    case False{}: start_bands(~custom,activation,depth,parallel,geometry,input,weights,bias,workers)
163
164def initial_layout(~custom: F32 -> F32,+activation: Activations.Activation,plan: BandPlanning.Plan,+geometry: Geometry,input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,workers: U32) -> Tensors.Tensor<F32>:
165  BandPlanning.Plan{+depth,parallel,time} = plan
166  initial_layout_selected(~custom,activation,U32.is_eq(depth,0),depth,parallel,geometry,input,weights,bias,workers)
167
168def execute_many(~custom: F32 -> F32,+activation: Activations.Activation,+geometry: Geometry,input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,+workers: U32) -> Tensors.Tensor<F32>:
169  Storage.Buffer{bias,length,Certificate.Certificate{+depth,equation}} = bias
170  +context = band_context(geometry,True{},activation)
171  initial_layout(~custom,activation,BandPlanning.select(band_work(context,U32.shln(1,depth)),context_window(context),workers,DefaultCost.configuration()),
172    geometry,input,weights,Storage.Buffer{bias,length,Certificate.Certificate{depth,equation}},workers)
173
174# Padding-only output windows and an empty batch use the existing dense path.
175# This is a representation choice, not a restriction on valid convolution shapes.
176def uses_band_windows(geometry: Geometry) -> Bool:
177  Geometry{count,batched,KernelShape.Dimensions{ci,rows,h,w,k,s,p}} = geometry
178  U32.is_lt(0,count) && Nat.is_lt(p,k)
179
180def execute_layout(~custom: F32 -> F32,+activation: Activations.Activation,band_windows: Bool,+geometry: Geometry,input: Storage.Buffer<F32>,
181  weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,workers: U32) -> Tensors.Tensor<F32>:
182  match band_windows:
183    case False{}: execute_dense(~custom,activation,geometry,input,weights,bias,workers)
184    case True{}: execute_many(~custom,activation,geometry,input,weights,bias,workers)
185
186def execute_selected(~custom: F32 -> F32,+activation: Activations.Activation,single_worker: Bool,+geometry: Geometry,input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,workers: U32) -> Tensors.Tensor<F32>:
187  match single_worker:
188    case True{}: execute_dense(~custom,activation,geometry,input,weights,bias,1)
189    case False{}: execute_layout(~custom,activation,uses_band_windows(geometry),geometry,input,weights,bias,workers)
190
191def execute(~custom: F32 -> F32,+activation: Activations.Activation,+geometry: Geometry,input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,+workers: U32) -> Tensors.Tensor<F32>:
192  execute_selected(~custom,activation,U32.is_eq(workers,1),geometry,input,weights,bias,workers)
193
194# The weights in the order the kernel reads them: once, after the checks.
195def reordered(dimensions: KernelShape.Dimensions,weights: Storage.Buffer<F32>) -> Storage.Buffer<Tile.Tile8>:
196  KernelShape.Dimensions{ci,rows,h,w,+k,s,p} = dimensions
197  GemmBuffer.reorder(U32.from_nat(rows),U32.from_nat(Nat.mul(ci,Nat.mul(k,k))),weights)
198
199def packed(geometry: Geometry,weights: Storage.Buffer<F32>) -> Storage.Buffer<Tile.Tile8>:
200  Geometry{count,batched,dimensions} = geometry
201  reordered(dimensions,weights)
202
203def machine(~custom: F32 -> F32,+activation: Activations.Activation,supported: Bool,+geometry: Geometry,input: Storage.Buffer<F32>,weights: Storage.Buffer<F32>,bias: Storage.Buffer<F32>,+shape: List<&2,U32>,fill: F32,+workers: U32) -> Tensors.Tensor<F32>:
204  match supported:
205    case True{}: execute(~custom,activation,geometry,input,packed(geometry,weights),bias,workers)
206    case False{}: Tensors.Tensor{input,shape,fill,Tensors.Failed{Tensors.Problem{"conv2d",Tensors.InvalidShape{shape}}},workers}
207
208def machine_support(geometry: Geometry,+limits: Machine.Limits) -> Bool:
209  Geometry{+count,batched,+dimensions} = geometry
210  KernelShape.Dimensions{+ci,+rows,+h,+w,+k,+s,+p} = dimensions
211  KernelShape.supports_packed(limits,dimensions,Nat.mul(ci,Nat.mul(h,w)),Nat.mul(rows,Nat.mul(ci,Nat.mul(k,k)))) &&
212    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)))
213
214def supported(~custom: F32 -> F32,+activation: Activations.Activation,valid: Bool,+geometry: Geometry,input: Storage.Buffer<F32>,weights: Storage.Buffer<F32>,bias: Storage.Buffer<F32>,+shape: List<&2,U32>,fill: F32,+workers: U32) -> Tensors.Tensor<F32>:
215  match valid:
216    case False{}: Tensors.Tensor{input,shape,fill,Tensors.Failed{Tensors.Problem{"conv2d",Tensors.InvalidArgument{shape}}},workers}
217    case True{}:
218      machine(~custom,activation,machine_support(geometry,Default.limits()),geometry,input,weights,bias,shape,fill,workers)
219
220# What a dense call needs of its three buffers: lengths that fit the shape and
221# the capacities that hold them.
222def dense_call_valid(+dimensions: KernelShape.Dimensions,+count: U32,+input_length: Nat,+input_depth: Nat,+weight_length: Nat,+weight_depth: Nat,
223  +bias_length: Nat,+bias_depth: Nat) -> Bool:
224  KernelShape.Dimensions{+ci,rows,+h,+w,k,s,p} = dimensions
225  KernelShape.valid(dimensions,Nat.mul(ci,Nat.mul(h,w)),weight_length,bias_length) &&
226    Nat.is_eq(input_length,Nat.mul(U32.to_nat(count),Nat.mul(ci,Nat.mul(h,w)))) &&
227    Nat.is_le(input_length,U32.to_nat(U32.shln(1,input_depth))) &&
228    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)))
229
230def checked(~custom: F32 -> F32,+activation: Activations.Activation,geometry: Maybe<&2,Geometry>,input: Storage.Buffer<F32>,weights: Storage.Buffer<F32>,bias: Storage.Buffer<F32>,+shape: List<&2,U32>,fill: F32,+workers: U32) -> Tensors.Tensor<F32>:
231  match geometry:
232    case None{}: Tensors.Tensor{input,shape,fill,Tensors.Failed{Tensors.Problem{"conv2d",Tensors.InvalidArgument{shape}}},workers}
233    case Some{+geometry}:
234      Geometry{+count,batched,+dimensions} = geometry
235      Storage.Buffer{input,+input_length,Certificate.Certificate{+input_depth,input_equation}} = input
236      Storage.Buffer{weights,+weight_length,Certificate.Certificate{+weight_depth,weight_equation}} = weights
237      Storage.Buffer{bias,+bias_length,Certificate.Certificate{+bias_depth,bias_equation}} = bias
238      supported(~custom,activation,
239        dense_call_valid(dimensions,count,input_length,input_depth,weight_length,weight_depth,bias_length,bias_depth),geometry,
240        Storage.Buffer{input,input_length,Certificate.Certificate{input_depth,input_equation}},
241        Storage.Buffer{weights,weight_length,Certificate.Certificate{weight_depth,weight_equation}},
242        Storage.Buffer{bias,bias_length,Certificate.Certificate{bias_depth,bias_equation}},shape,fill,workers)
243
244def status(~custom: F32 -> F32,+activation: Activations.Activation,status: Tensors.Status,input: Storage.Buffer<F32>,+input_extents: List<&2,U32>,fill: F32,
245  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>:
246  match status:
247    case Tensors.Failed{problem}: Tensors.Tensor{input,input_extents,fill,Tensors.Failed{problem},workers}
248    case Tensors.Usable{}: checked(~custom,activation,geometry(input_shape(input_extents),weight_shape,bias_shape,stride,padding),input,weights,bias,input_extents,fill,workers)
249
250def replanned(~custom: F32 -> F32,+activation: Activations.Activation,+geometry: Geometry,selection: BandPlanning.LayoutPlan,input: Bands.Bands(~F32),weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,workers: U32) -> Tensors.Tensor<F32>:
251  BandPlanning.LayoutPlan{layout,+parallel,time,repartition} = selection
252  band_output(~custom,activation,geometry,parallel,Bands.align(~F32,BandWindow.input_geometry(context_window(band_context(geometry,parallel,Activations.Identity{}))),0.0,layout,input),weights,bias,workers)
253
254def existing_layout(~custom: F32 -> F32,+activation: Activations.Activation,+geometry: Geometry,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,+workers: U32,
255  owners: Bands.Bands(~F32) & Bands.Plan) -> Tensors.Tensor<F32>:
256  Storage.Buffer{bias,length,Certificate.Certificate{+depth,equation}} = bias
257  (input,plan) = owners
258  +context = band_context(geometry,True{},activation)
259  replanned(~custom,activation,geometry,BandPlanning.replan(plan,band_work(context,U32.shln(1,depth)),context_window(context),workers,DefaultCost.configuration()),
260    input,weights,Storage.Buffer{bias,length,Certificate.Certificate{depth,equation}},workers)
261
262def execute_existing_layout(~custom: F32 -> F32,+activation: Activations.Activation,band_windows: Bool,+geometry: Geometry,+layout: Bands.Geometry,
263  input: Bands.Bands(~F32),weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,workers: U32) -> Tensors.Tensor<F32>:
264  match band_windows:
265    case True{}: existing_layout(~custom,activation,geometry,weights,bias,workers,Bands.plan(~F32,input))
266    case False{}:
267      Bands.Geometry{+images,+channels,+height,+width} = layout
268      execute_dense(~custom,activation,geometry,
269        Bands.gather(~F32,layout,input,Tensors.buffer(~F32,0.0,U32.from_nat(Nat.mul(U32.to_nat(images),Nat.mul(U32.to_nat(channels),Nat.mul(U32.to_nat(height),U32.to_nat(width))))))),
270        weights,bias,workers)
271
272def band_machine(~custom: F32 -> F32,+activation: Activations.Activation,supported: Bool,+geometry: Geometry,layout: Bands.Geometry,input: Bands.Bands(~F32),
273  weights: Storage.Buffer<F32>,bias: Storage.Buffer<F32>,+shape: List<&2,U32>,fill: F32,workers: U32) -> Tensors.Tensor<F32>:
274  match supported:
275    case True{}: execute_existing_layout(~custom,activation,uses_band_windows(geometry),geometry,layout,input,packed(geometry,weights),bias,workers)
276    case False{}: Tensors.Banded{input,layout,shape,fill,Tensors.Failed{Tensors.Problem{"conv2d",Tensors.InvalidShape{shape}}},workers}
277
278def band_supported(~custom: F32 -> F32,+activation: Activations.Activation,valid: Bool,+geometry: Geometry,layout: Bands.Geometry,input: Bands.Bands(~F32),
279  weights: Storage.Buffer<F32>,bias: Storage.Buffer<F32>,+shape: List<&2,U32>,fill: F32,workers: U32) -> Tensors.Tensor<F32>:
280  match valid:
281    case False{}: Tensors.Banded{input,layout,shape,fill,Tensors.Failed{Tensors.Problem{"conv2d",Tensors.InvalidArgument{shape}}},workers}
282    case True{}:
283      band_machine(~custom,activation,machine_support(geometry,Default.limits()),geometry,layout,input,weights,bias,shape,fill,workers)
284
285# What a call on bands needs: bands that cover the images of the geometry, and
286# parameters whose lengths fit the shape and their capacities.
287def band_call_valid(+dimensions: KernelShape.Dimensions,+count: U32,layout: Bands.Geometry,inspection: Bands.Inspection,+weight_length: Nat,+weight_depth: Nat,
288  +bias_length: Nat,+bias_depth: Nat) -> Bool:
289  KernelShape.Dimensions{+ci,rows,+h,+w,k,s,p} = dimensions
290  Bands.Geometry{batch,channels,height,width} = layout
291  Bands.Inspection{valid,first,end,leaves} = inspection
292  valid && U32.is_zero(first) && Nat.is_eq(U32.to_nat(end),Nat.mul(U32.to_nat(count),h)) &&
293    U32.is_eq(batch,count) && Nat.is_eq(U32.to_nat(channels),ci) && Nat.is_eq(U32.to_nat(height),h) && Nat.is_eq(U32.to_nat(width),w) &&
294    KernelShape.valid(dimensions,Nat.mul(ci,Nat.mul(h,w)),weight_length,bias_length) &&
295    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)))
296
297def band_inspected(~custom: F32 -> F32,+activation: Activations.Activation,+geometry: Geometry,+layout: Bands.Geometry,weights: Storage.Buffer<F32>,bias: Storage.Buffer<F32>,shape: List<&2,U32>,fill: F32,workers: U32,
298  inspected: Bands.Bands(~F32) & Bands.Inspection) -> Tensors.Tensor<F32>:
299  Geometry{+count,batched,+dimensions} = geometry
300  Storage.Buffer{weights,+weight_length,Certificate.Certificate{+weight_depth,weight_equation}} = weights
301  Storage.Buffer{bias,+bias_length,Certificate.Certificate{+bias_depth,bias_equation}} = bias
302  (input,inspection) = inspected
303  band_supported(~custom,activation,band_call_valid(dimensions,count,layout,inspection,weight_length,weight_depth,bias_length,bias_depth),
304    geometry,layout,input,Storage.Buffer{weights,weight_length,Certificate.Certificate{weight_depth,weight_equation}},
305    Storage.Buffer{bias,bias_length,Certificate.Certificate{bias_depth,bias_equation}},shape,fill,workers)
306
307def band_checked(~custom: F32 -> F32,+activation: Activations.Activation,geometry: Maybe<&2,Geometry>,+layout: Bands.Geometry,input: Bands.Bands(~F32),
308  weights: Storage.Buffer<F32>,bias: Storage.Buffer<F32>,+shape: List<&2,U32>,fill: F32,workers: U32) -> Tensors.Tensor<F32>:
309  match geometry:
310    case None{}: Tensors.Banded{input,layout,shape,fill,Tensors.Failed{Tensors.Problem{"conv2d",Tensors.InvalidArgument{shape}}},workers}
311    case Some{geometry}: band_inspected(~custom,activation,geometry,layout,weights,bias,shape,fill,workers,Bands.inspect(~F32,layout,input))
312
313def band_status(~custom: F32 -> F32,+activation: Activations.Activation,status: Tensors.Status,layout: Bands.Geometry,input: Bands.Bands(~F32),+shape: List<&2,U32>,fill: F32,
314  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>:
315  match status:
316    case Tensors.Failed{problem}: Tensors.Banded{input,layout,shape,fill,Tensors.Failed{problem},workers}
317    case Tensors.Usable{}: band_checked(~custom,activation,geometry(input_shape(shape),weight_shape,bias_shape,stride,padding),layout,input,weights,bias,shape,fill,workers)
318
319def 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>,
320  parameter_status: Tensors.Status,parameter_workers: U32,stride: U32,padding: U32) -> Tensors.Tensor<F32>:
321  match input:
322    case Tensors.Tensor{input,shape,fill,input_status,workers}:
323      status(~custom,activation,Tensors.first_failure(input_status,parameter_status),input,shape,fill,weights,weight_shape,bias,bias_shape,stride,padding,U32.max(workers,parameter_workers))
324    case Tensors.Banded{input,layout,shape,fill,input_status,workers}:
325      band_status(~custom,activation,Tensors.first_failure(input_status,parameter_status),layout,input,shape,fill,weights,weight_shape,bias,bias_shape,stride,padding,U32.max(workers,parameter_workers))
326
327def input_selected(~custom: F32 -> F32,+activation: Activations.Activation,dense: Bool,input: Tensors.Tensor<F32>,weights: Storage.Buffer<F32>,weight_shape: List<&2,U32>,bias: Storage.Buffer<F32>,bias_shape: List<&2,U32>,
328  +parameter_status: Tensors.Status,+parameter_workers: U32,+stride: U32,+padding: U32) -> Tensors.Tensor<F32>:
329  match dense:
330    case False{}: input_layout(~custom,activation,input,weights,weight_shape,bias,bias_shape,parameter_status,parameter_workers,stride,padding)
331    case True{}:
332      Tensors.with_dense(~F32,~Tensors.Tensor<F32>,input,input => shape => fill => input_status => workers =>
333        status(~custom,activation,Tensors.first_failure(input_status,parameter_status),input,shape,fill,weights,weight_shape,bias,bias_shape,stride,padding,U32.max(workers,parameter_workers)))
334
335# Parameter representations select one public path before any element work.
336# Dense parameters retain the existing incoming-band convolution leaf path.
337def conv2d_selected(~custom: F32 -> F32,+activation: Activations.Activation,+dense: Bool,input: Tensors.Tensor<F32>,weights: Tensors.Tensor<F32>,bias: Tensors.Tensor<F32>,+stride: U32,+padding: U32) -> Tensors.Tensor<F32>:
338  Tensors.with_dense(~F32,~Tensors.Tensor<F32>,weights,weights => weight_shape => weight_fill => weight_status => weight_workers =>
339    Tensors.with_dense(~F32,~Tensors.Tensor<F32>,bias,bias => bias_shape => bias_fill => bias_status => bias_workers =>
340      input_selected(~custom,activation,dense,input,weights,weight_shape,bias,bias_shape,Tensors.first_failure(weight_status,bias_status),U32.max(weight_workers,bias_workers),stride,padding)))
341
342def conv2d_bias_layout(~custom: F32 -> F32,+activation: Activations.Activation,dense: Bool,input: Tensors.Tensor<F32>,weights: Tensors.Tensor<F32>,bias: Tensors.Tensor<F32>,stride: U32,padding: U32) -> Tensors.Tensor<F32>:
343  match bias:
344    case Tensors.Tensor{storage,shape,fill,status,workers}:
345      conv2d_selected(~custom,activation,dense,input,weights,Tensors.Tensor{storage,shape,fill,status,workers},stride,padding)
346    case Tensors.Banded{bands,geometry,shape,fill,status,workers}:
347      conv2d_selected(~custom,activation,True{},input,weights,Tensors.Banded{bands,geometry,shape,fill,status,workers},stride,padding)
348
349def 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>:
350  match weights:
351    case Tensors.Tensor{storage,shape,fill,status,workers}:
352      conv2d_bias_layout(~custom,activation,False{},input,Tensors.Tensor{storage,shape,fill,status,workers},bias,stride,padding)
353    case Tensors.Banded{bands,geometry,shape,fill,status,workers}:
354      conv2d_bias_layout(~custom,activation,True{},input,Tensors.Banded{bands,geometry,shape,fill,status,workers},bias,stride,padding)