~/jendhub

tensor_window.bend

stelliferous@0.0.2.0 · 120 lines · 16 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.2.0/storage_buffer.bend as Storage
8import stelliferous@0.0.2.0/tensor_shape.bend as Shapes
9import stelliferous@0.0.2.0/tensor_access.bend as Access
10import stelliferous@0.0.2.0/tensor_operations.bend as Operations
11import stelliferous@0.0.2.0/tensor.bend as Tensors
12import stelliferous@0.0.2.0/tensor_layout.bend as Layout
13import stelliferous@0.0.2.0/tensor_band.bend as Bands
14import stelliferous@0.0.2.0/tensor_band_window.bend as BandWindow
15import stelliferous@0.0.2.0/tensor_band_pool.bend as BandPool
16
17def out_extent(+extent: U32,+kernel: U32,+stride: U32,+padding: U32) -> U32:
18  ((extent + 2 * padding - kernel : U32) / U32.max(stride,1) + 1 : U32)
19
20def geometry(+shape: List<&2,U32>,+kernel: U32,+stride: U32,+padding: U32) -> Access.Pool:
21  +rank = List.length(&2,U32,shape)
22  +height = Layout.extent_at(shape,Nat.sub(rank,2n))
23  +width = Layout.extent_at(shape,Nat.sub(rank,1n))
24  Access.Pool{height,width,kernel,stride,padding,out_extent(height,kernel,stride,padding),out_extent(width,kernel,stride,padding)}
25
26def leading(ranked: Bool,+shape: List<&2,U32>) -> Maybe<&2,U32>:
27  match ranked:
28    case True{}: Shapes.product(List.take(&2,U32,shape,Nat.sub(List.length(&2,U32,shape),2n)),Some{1})
29    case False{}: None{}
30
31# The number of planes before the last two axes; none below rank two.
32def planes(+shape: List<&2,U32>) -> Maybe<&2,U32>:
33  leading(Nat.is_le(2n,List.length(&2,U32,shape)),shape)
34
35def pooled(~Element: Data,name: String,+shape: List<&2,U32>,+out_shape: List<&2,U32>,fill: Element,+workers: U32,
36  result: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) -> Tensors.Tensor<Element>:
37  match result:
38    case Operations.Accepted{(input,output)}: Tensors.Tensor{output,out_shape,fill,Tensors.Usable{},workers}
39    case Operations.Rejected{(input,output)}: Tensors.Tensor{output,out_shape,fill,Tensors.Failed{Tensors.Problem{name,Tensors.InvalidArgument{shape}}},workers}
40
41def pool_middle(~Element: Data,~operation: Element -> Element -> Element,name: String,+pad: Element,middle: Maybe<&2,U32>,count: U32,+channels: U32,+pool: Access.Pool,
42  storage: Storage.Buffer<Element>,+shape: List<&2,U32>,+out_shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
43  match middle:
44    case Some{middle}:
45      pooled(~Element,name,shape,out_shape,fill,workers,
46        Operations.pool(~Element,~operation,pad,storage,pool,channels,Tensors.buffer(~Element,fill,middle),Tensors.buffer(~Element,fill,count)))
47    case None{}: Tensors.failed(~Element,name,Tensors.InvalidArgument{shape},shape,fill,workers)
48
49# The column-swept planes keep the rows of the input and have the columns of
50# the output.
51def middle_shape(+shape: List<&2,U32>,pool: Access.Pool) -> List<&2,U32>:
52  Access.Pool{height,width,kernel,stride,padding,out_height,out_width} = pool
53  List.append(&2,U32,List.take(&2,U32,shape,Nat.sub(List.length(&2,U32,shape),2n)),[height,out_width])
54
55def pool_counts(~Element: Data,~operation: Element -> Element -> Element,name: String,+pad: Element,channels: Maybe<&2,U32>,count: Maybe<&2,U32>,+pool: Access.Pool,
56  storage: Storage.Buffer<Element>,+shape: List<&2,U32>,+out_shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
57  match channels count:
58    case Some{+channels} Some{count}:
59      pool_middle(~Element,~operation,name,pad,Tensors.count_of(middle_shape(shape,pool)),count,channels,pool,storage,shape,out_shape,fill,workers)
60    case channels count: Tensors.failed(~Element,name,Tensors.InvalidArgument{shape},shape,fill,workers)
61
62def pool_planes(~Element: Data,~operation: Element -> Element -> Element,name: String,+pad: Element,channels: Maybe<&2,U32>,+pool: Access.Pool,
63  storage: Storage.Buffer<Element>,+shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
64  Access.Pool{height,width,kernel,stride,padding,+out_height,+out_width} = pool
65  +out_shape = List.append(&2,U32,List.take(&2,U32,shape,Nat.sub(List.length(&2,U32,shape),2n)),[out_height,out_width])
66  pool_counts(~Element,~operation,name,pad,channels,Tensors.count_of(out_shape),pool,storage,shape,out_shape,fill,workers)
67
68def pool_status(~Element: Data,~operation: Element -> Element -> Element,name: String,status: Tensors.Status,+pad: Element,+kernel: U32,+stride: U32,+padding: U32,
69  storage: Storage.Buffer<Element>,+shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
70  match status:
71    case Tensors.Usable{}:
72      pool_planes(~Element,~operation,name,pad,planes(shape),geometry(shape,kernel,stride,padding),storage,shape,fill,workers)
73    case Tensors.Failed{problem}: Tensors.Tensor{storage,shape,fill,Tensors.Failed{problem},workers}
74
75# pad is both the value outside the input and the initial accumulator.
76def band_shape(shape: List<&2,U32>,geometry: Bands.Geometry) -> Bool:
77  match shape geometry:
78    case [c,h,w] Bands.Geometry{n,ci,hi,wi}: U32.is_eq(n,1) && U32.is_eq(c,ci) && U32.is_eq(h,hi) && U32.is_eq(w,wi)
79    case [n,c,h,w] Bands.Geometry{ni,ci,hi,wi}: U32.is_eq(n,ni) && U32.is_eq(c,ci) && U32.is_eq(h,hi) && U32.is_eq(w,wi)
80    case shape geometry: False{}
81
82def band_execute(~Element: Data,~operation: Element -> Element -> Element,+pad: Element,pool: Access.Pool,+geometry: Bands.Geometry,bands: Bands.Bands(~Element),
83  out_shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
84  match pool geometry:
85    case Access.Pool{height,width,kernel,stride,padding,+oh,+ow} Bands.Geometry{batch,channels,ih,iw}:
86      Tensors.Banded{BandPool.execute(~Element,~operation,BandPool.Context{BandWindow.Geometry{geometry,kernel,stride,padding,oh,ow},pad,fill,workers},bands),
87        Bands.Geometry{batch,channels,oh,ow},out_shape,fill,Tensors.Usable{},workers}
88
89# The band route needs the same counts as the dense route: the planes swept
90# along the columns and the output.
91def band_count(~Element: Data,~operation: Element -> Element -> Element,name: String,valid: Bool,planes: Maybe<&2,U32>,count: Maybe<&2,U32>,+pad: Element,+pool: Access.Pool,+geometry: Bands.Geometry,bands: Bands.Bands(~Element),
92  +shape: List<&2,U32>,out_shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
93  match valid planes count pool:
94    case True{} Some{planes} Some{count} pool: band_execute(~Element,~operation,pad,pool,geometry,bands,out_shape,fill,workers)
95    case valid planes count Access.Pool{height,width,kernel,stride,padding,oh,ow}:
96      Tensors.with_dense(~Element,~Tensors.Tensor<Element>,Tensors.Banded{bands,geometry,shape,fill,Tensors.Usable{},workers},storage => shape => fill => status => workers =>
97        pool_status(~Element,~operation,name,status,pad,kernel,stride,padding,storage,shape,fill,workers))
98
99def band_geometry(~Element: Data,~operation: Element -> Element -> Element,name: String,pad: Element,+pool: Access.Pool,+geometry: Bands.Geometry,bands: Bands.Bands(~Element),
100  +shape: List<&2,U32>,fill: Element,workers: U32) -> Tensors.Tensor<Element>:
101  Access.Pool{+height,+width,+kernel,+stride,+padding,+out_height,+out_width} = pool
102  +pool = {Access.Pool{height,width,kernel,stride,padding,out_height,out_width} : Access.Pool}
103  +out_shape = List.append(&2,U32,List.take(&2,U32,shape,Nat.sub(List.length(&2,U32,shape),2n)),[out_height,out_width])
104  band_count(~Element,~operation,name,band_shape(shape,geometry) && Operations.pool_geometry(pool) && U32.is_lt(0,height) && U32.is_lt(0,width) &&
105    Bool.not(Shapes.has_zero(out_shape)),Tensors.count_of(middle_shape(shape,pool)),Tensors.count_of(out_shape),pad,pool,geometry,bands,shape,out_shape,fill,workers)
106
107def band_status(~Element: Data,~operation: Element -> Element -> Element,name: String,status: Tensors.Status,pad: Element,+kernel: U32,+stride: U32,+padding: U32,
108  layout: Bands.Geometry,bands: Bands.Bands(~Element),+shape: List<&2,U32>,fill: Element,workers: U32) -> Tensors.Tensor<Element>:
109  match status:
110    case Tensors.Failed{+problem}:
111      Tensors.with_dense(~Element,~Tensors.Tensor<Element>,Tensors.Banded{bands,layout,shape,fill,Tensors.Failed{problem},workers},
112        storage => shape => fill => status => workers =>
113          pool_status(~Element,~operation,name,status,pad,kernel,stride,padding,storage,shape,fill,workers))
114    case Tensors.Usable{}: band_geometry(~Element,~operation,name,pad,geometry(shape,kernel,stride,padding),layout,bands,shape,fill,workers)
115
116def pool2d(~Element: Data,~operation: Element -> Element -> Element,name: String,+pad: Element,tensor: Tensors.Tensor<Element>,
117  +kernel: U32,+stride: U32,+padding: U32) -> Tensors.Tensor<Element>:
118  match tensor:
119    case Tensors.Tensor{storage,shape,fill,status,workers}: pool_status(~Element,~operation,name,status,pad,kernel,stride,padding,storage,shape,fill,workers)
120    case Tensors.Banded{bands,layout,shape,fill,status,workers}: band_status(~Element,~operation,name,status,pad,kernel,stride,padding,layout,bands,shape,fill,workers)