tensor_window.bend
stelliferous@0.0.3.0 · 91 lines · 0 definitions
1# Two-dimensional pooling over the last two axes, as two ordered sweeps. Along
2# the columns, each row of the window folds its k samples in increasing kx
3# from pad: line = op(...op(op(pad, x0), x1)..., x(k-1)). Along the rows, the
4# output folds the k lines in increasing ky from pad. A sample outside the
5# input reads pad, so pad should be the identity of the operation.
6import Base
7import stelliferous@0.0.3.0/tensor_shape.bend as Shapes
8import stelliferous@0.0.3.0/tensor_access.bend as Access
9import stelliferous@0.0.3.0/tensor_operations.bend as Operations
10import stelliferous@0.0.3.0/tensor.bend as Tensors
11import stelliferous@0.0.3.0/tensor_layout.bend as Layout
12import stelliferous@0.0.3.0/tensor_band.bend as Bands
13import stelliferous@0.0.3.0/tensor_band_window.bend as BandWindow
14import stelliferous@0.0.3.0/tensor_band_pool.bend as BandPool
15import stelliferous@0.0.3.0/tensor_band_schedule.bend as BandSchedule
16import stelliferous@0.0.3.0/tensor_view.bend as Views
17
18def out_extent(+extent: U32,+kernel: U32,+stride: U32,+padding: U32) -> U32:
19 ((extent + 2 * padding - kernel : U32) / U32.max(stride,1) + 1 : U32)
20
21def geometry(+shape: List<&2,U32>,+kernel: U32,+stride: U32,+padding: U32) -> Access.Pool:
22 +rank = List.length(&2,U32,shape)
23 +height = Layout.extent_at(shape,Nat.sub(rank,2n))
24 +width = Layout.extent_at(shape,Nat.sub(rank,1n))
25 Access.Pool{height,width,kernel,stride,padding,out_extent(height,kernel,stride,padding),out_extent(width,kernel,stride,padding)}
26
27# The column-swept planes keep the rows of the input and have the columns of
28# the output.
29def middle_shape(+shape: List<&2,U32>,pool: Access.Pool) -> List<&2,U32>:
30 Access.Pool{height,width,kernel,stride,padding,out_height,out_width} = pool
31 List.append(&2,U32,List.take(&2,U32,shape,Nat.sub(List.length(&2,U32,shape),2n)),[height,out_width])
32
33def output_shape(+shape: List<&2,U32>,pool: Access.Pool) -> List<&2,U32>:
34 Access.Pool{height,width,kernel,stride,padding,out_height,out_width} = pool
35 List.append(&2,U32,List.take(&2,U32,shape,Nat.sub(List.length(&2,U32,shape),2n)),[out_height,out_width])
36
37# pad is both the value outside the input and the initial accumulator.
38# The pooled bands are the result when every band's pool accepted and the
39# bands cover the output geometry.
40def band_execute(~Element: Data,~operation: Element -> Element -> Element,name: String,+pad: Element,pool: Access.Pool,+geometry: Bands.Geometry,bands: Bands.Bands(~Element),
41 shape: List<&2,U32>,out_shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
42 match pool geometry:
43 case Access.Pool{height,width,kernel,stride,padding,+oh,+ow} Bands.Geometry{+batch,+channels,ih,iw}:
44 Tensors.transformed(~Element,name,Tensors.InvalidArgument{shape},Bands.Geometry{batch,channels,oh,ow},out_shape,fill,workers,
45 BandPool.execute(~Element,~operation,BandPool.Context{BandWindow.Geometry{geometry,kernel,stride,padding,oh,ow},pad,fill,workers},bands))
46
47# The input arrives in the layout of its planes, and the band pool runs there.
48def planned(~Element: Data,~operation: Element -> Element -> Element,name: String,+pad: Element,+pool: Access.Pool,tensor: Tensors.Tensor<Element>,
49 +shape: List<&2,U32>,+fill: Element,+workers: U32,target: BandSchedule.Target) -> Tensors.Tensor<Element>:
50 BandSchedule.Target{+geometry,layout,split} = target
51 band_execute(~Element,~operation,name,pad,pool,geometry,Tensors.arrive(~Element,geometry,layout,tensor),shape,output_shape(shape,pool),fill,workers)
52
53def offered(~Element: Data,~operation: Element -> Element -> Element,name: String,+pad: Element,+pool: Access.Pool,+planes: U32,
54 +shape: List<&2,U32>,+fill: Element,+workers: U32,owners: Tensors.Tensor<Element> & BandSchedule.Target) -> Tensors.Tensor<Element>:
55 Access.Pool{+height,+width,kernel,stride,padding,out_height,out_width} = pool
56 (tensor,target) = owners
57 planned(~Element,~operation,name,pad,pool,tensor,shape,fill,workers,BandSchedule.plane_target(shape,planes,height,width,target))
58
59# An empty output has no element to pool.
60def nonempty(~Element: Data,~operation: Element -> Element -> Element,name: String,empty: Bool,+pad: Element,+pool: Access.Pool,+planes: U32,
61 tensor: Tensors.Tensor<Element>,+shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
62 match empty:
63 case True{}: Tensors.whole(~Element,Tensors.buffer(~Element,fill,0),output_shape(shape,pool),fill,Tensors.Usable{},workers)
64 case False{}: offered(~Element,~operation,name,pad,pool,planes,shape,fill,workers,Tensors.offer(~Element,tensor))
65
66def checked(~Element: Data,~operation: Element -> Element -> Element,name: String,valid: Bool,planes: Maybe<&2,U32>,+pad: Element,+pool: Access.Pool,
67 tensor: Tensors.Tensor<Element>,+shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
68 match valid planes:
69 case True{} Some{+planes}:
70 nonempty(~Element,~operation,name,Shapes.has_zero(output_shape(shape,pool)),pad,pool,planes,tensor,shape,fill,workers)
71 case valid planes: Tensors.failed(~Element,name,Tensors.InvalidArgument{shape},shape,fill,workers)
72
73# A pool needs planes before the last two axes, positive extents and windows
74# that fit them, and counts of the column-swept planes and of the output.
75def valid(+shape: List<&2,U32>,+pool: Access.Pool) -> Bool:
76 Access.Pool{+height,+width,kernel,stride,padding,out_height,out_width} = pool
77 Operations.pool_geometry(pool) && U32.is_lt(0,height) && U32.is_lt(0,width) &&
78 Views.has_count(Tensors.count_of(middle_shape(shape,pool))) && Views.has_count(Tensors.count_of(output_shape(shape,pool)))
79
80def pool_status(~Element: Data,~operation: Element -> Element -> Element,name: String,status: Tensors.Status,+pad: Element,+kernel: U32,+stride: U32,+padding: U32,
81 bands: Bands.Bands(~Element),layout: Bands.Geometry,+shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
82 match status:
83 case Tensors.Failed{problem}: Tensors.Tensor{bands,layout,shape,fill,Tensors.Failed{problem},workers}
84 case Tensors.Usable{}:
85 +pool = geometry(shape,kernel,stride,padding)
86 checked(~Element,~operation,name,valid(shape,pool),Layout.planes(shape),pad,pool,Tensors.Tensor{bands,layout,shape,fill,Tensors.Usable{},workers},shape,fill,workers)
87
88def pool2d(~Element: Data,~operation: Element -> Element -> Element,name: String,+pad: Element,tensor: Tensors.Tensor<Element>,
89 +kernel: U32,+stride: U32,+padding: U32) -> Tensors.Tensor<Element>:
90 Tensors.Tensor{bands,layout,+shape,+fill,status,+workers} = tensor
91 pool_status(~Element,~operation,name,status,pad,kernel,stride,padding,bands,layout,shape,fill,workers)