tensor_parameter_gradient.bend
stelliferous@0.0.3.0 · 112 lines · 0 definitions
1# Shape validation and affine tensor ownership for the shared reduction kernel.
2import Base
3import stelliferous@0.0.3.0/tensor.bend as Tensors
4import stelliferous@0.0.3.0/storage_buffer.bend as Storage
5import stelliferous@0.0.3.0/kernel_parameter_gradient.bend as Kernel
6import stelliferous@0.0.3.0/kernel_parameter_gradient_buffer.bend as KernelBuffer
7import stelliferous@0.0.3.0/kernel_weight_gradient.bend as WeightKernel
8import stelliferous@0.0.3.0/kernel_shape.bend as Shape
9
10type ImageGeometry is Data:
11 ImageGeometry{images: U32,channels: U32,height: U32,width: U32}
12
13def image(shape: List<&2,U32>) -> Maybe<&2,ImageGeometry>:
14 match shape:
15 case [n,c,h,w]: Some{ImageGeometry{n,c,h,w}}
16 case [c,h,w]: Some{ImageGeometry{1,c,h,w}}
17 case shape: None{}
18
19def positive(image: ImageGeometry) -> Bool:
20 ImageGeometry{n,c,h,w} = image
21 U32.is_lt(0,n) && U32.is_lt(0,c) && U32.is_lt(0,h) && U32.is_lt(0,w)
22
23def work(x: ImageGeometry,d: ImageGeometry,kernel: U32,stride: U32,padding: U32,group: U32) -> Kernel.Work:
24 ImageGeometry{n,c,h,w} = x
25 ImageGeometry{dn,o,rows,columns} = d
26 Kernel.Work{U32.to_nat(n),U32.to_nat(c),U32.to_nat(o),U32.to_nat(h),U32.to_nat(w),U32.to_nat(kernel),U32.to_nat(stride),U32.to_nat(padding),U32.to_nat(rows),U32.to_nat(columns),U32.to_nat(group)}
27
28def fits(count: Maybe<&2,U32>) -> Bool:
29 match count:
30 case None{}: False{}
31 case Some{count}: True{}
32
33def valid(+x: ImageGeometry,+d: ImageGeometry,+kernel: U32,+stride: U32,+padding: U32,group: U32) -> Bool:
34 ImageGeometry{+n,+c,+h,+w} = x
35 ImageGeometry{dn,+o,rows,columns} = d
36 positive(x) && positive(d) && U32.is_eq(n,dn) && U32.is_lt(0,kernel) && U32.is_lt(0,stride) && U32.is_lt(0,group) &&
37 Nat.is_le(U32.to_nat(kernel),Nat.add(U32.to_nat(h),Nat.mul(2n,U32.to_nat(padding)))) &&
38 Nat.is_le(U32.to_nat(kernel),Nat.add(U32.to_nat(w),Nat.mul(2n,U32.to_nat(padding)))) &&
39 Nat.is_eq(U32.to_nat(rows),Shape.axis_size(U32.to_nat(h),U32.to_nat(kernel),U32.to_nat(stride),U32.to_nat(padding))) &&
40 Nat.is_eq(U32.to_nat(columns),Shape.axis_size(U32.to_nat(w),U32.to_nat(kernel),U32.to_nat(stride),U32.to_nat(padding))) && fits(Tensors.count_of([o,c,kernel,kernel]))
41
42def weight_finished(+xs: List<&2,U32>,xf: F32,xs_status: Tensors.Status,+xw: U32,+ds: List<&2,U32>,df: F32,ds_status: Tensors.Status,+dw: U32,
43 shape: List<&2,U32>,result: (Storage.Buffer<F32> & Storage.Buffer<F32>) & Storage.Buffer<F32>) -> Tensors.Tensor<F32> & (Tensors.Tensor<F32> & Tensors.Tensor<F32>):
44 ((x,d),values) = result
45 (Tensors.whole(~F32,values,shape,0.0,Tensors.Usable{},U32.max(xw,dw)),
46 (Tensors.whole(~F32,x,xs,xf,xs_status,xw),Tensors.whole(~F32,d,ds,df,ds_status,dw)))
47
48def weight_selected(accepted: Bool,+xi: ImageGeometry,+di: ImageGeometry,+k: U32,s: U32,p: U32,g: U32,
49 x: Storage.Buffer<F32>,+xs: List<&2,U32>,xf: F32,xs_status: Tensors.Status,+xw: U32,
50 d: Storage.Buffer<F32>,ds: List<&2,U32>,df: F32,ds_status: Tensors.Status,+dw: U32) -> Tensors.Tensor<F32> & (Tensors.Tensor<F32> & Tensors.Tensor<F32>):
51 match accepted:
52 case False{}:
53 (Tensors.failed(~F32,"conv2d_weight_gradient",Tensors.InvalidArgument{xs},[],0.0,xw),
54 (Tensors.whole(~F32,x,xs,xf,xs_status,xw),Tensors.whole(~F32,d,ds,df,ds_status,dw)))
55 case True{}:
56 ImageGeometry{n,c,h,w} = xi
57 ImageGeometry{dn,o,rows,columns} = di
58 weight_finished(xs,xf,xs_status,xw,ds,df,ds_status,dw,[o,c,k,k],WeightKernel.weight_parts(work(xi,di,k,s,p,g),U32.max(xw,dw),x,d))
59
60# The statuses are matched, not tested and kept: a status used twice would
61# make the constructors of the whole call shared.
62def weight_shaped(xi: Maybe<&2,ImageGeometry>,di: Maybe<&2,ImageGeometry>,xs_status: Tensors.Status,ds_status: Tensors.Status,+k: U32,+s: U32,+p: U32,+g: U32,
63 x: Storage.Buffer<F32>,+xs: List<&2,U32>,xf: F32,+xw: U32,
64 d: Storage.Buffer<F32>,ds: List<&2,U32>,df: F32,dw: U32) -> Tensors.Tensor<F32> & (Tensors.Tensor<F32> & Tensors.Tensor<F32>):
65 match xi di xs_status ds_status:
66 case Some{+xi} Some{+di} Tensors.Usable{} Tensors.Usable{}:
67 weight_selected(valid(xi,di,k,s,p,g),xi,di,k,s,p,g,x,xs,xf,Tensors.Usable{},xw,d,ds,df,Tensors.Usable{},dw)
68 case xi di xs_status ds_status:
69 weight_selected(False{},ImageGeometry{0,0,0,0},ImageGeometry{0,0,0,0},k,s,p,g,x,xs,xf,xs_status,xw,d,ds,df,ds_status,dw)
70
71def weight(x: Tensors.Tensor<F32>,d: Tensors.Tensor<F32>,+k: U32,+s: U32,+p: U32,+g: U32) -> Tensors.Tensor<F32> & (Tensors.Tensor<F32> & Tensors.Tensor<F32>):
72 Tensors.with_whole(~F32,~(Tensors.Tensor<F32> & (Tensors.Tensor<F32> & Tensors.Tensor<F32>)),x,x => xs => xf => status => xw =>
73 +xs = xs
74 Tensors.with_whole(~F32,~(Tensors.Tensor<F32> & (Tensors.Tensor<F32> & Tensors.Tensor<F32>)),d,d => ds => df => d_status => dw =>
75 +ds = ds
76 weight_shaped(image(xs),image(ds),status,d_status,k,s,p,g,x,xs,xf,xw,d,ds,df,dw)))
77
78def bias_work(d: ImageGeometry,group: U32) -> Kernel.Work:
79 ImageGeometry{n,o,rows,columns} = d
80 Kernel.Work{U32.to_nat(n),0n,U32.to_nat(o),0n,0n,1n,1n,0n,U32.to_nat(rows),U32.to_nat(columns),U32.to_nat(group)}
81
82def bias_finished(+shape: List<&2,U32>,fill: F32,status: Tensors.Status,+workers: U32,+outputs: U32,result: Storage.Buffer<F32> & Storage.Buffer<F32>) -> Tensors.Tensor<F32> & Tensors.Tensor<F32>:
83 (d,values) = result
84 (Tensors.whole(~F32,values,[outputs],0.0,Tensors.Usable{},workers),Tensors.whole(~F32,d,shape,fill,status,workers))
85
86# One output per channel, written into storage of the count of [o].
87def bias_counted(count: Maybe<&2,U32>,+di: ImageGeometry,+g: U32,d: Storage.Buffer<F32>,+shape: List<&2,U32>,fill: F32,status: Tensors.Status,+workers: U32) ->
88 Tensors.Tensor<F32> & Tensors.Tensor<F32>:
89 match count:
90 case Some{+count}:
91 ImageGeometry{n,+o,rows,columns} = di
92 bias_finished(shape,fill,status,workers,o,KernelBuffer.bias(bias_work(ImageGeometry{n,o,rows,columns},g),d,Tensors.buffer(~F32,0.0,count)))
93 case None{}:
94 ImageGeometry{n,o,rows,columns} = di
95 (Tensors.failed(~F32,"conv2d_bias_gradient",Tensors.InvalidShape{[o]},[o],0.0,workers),Tensors.whole(~F32,d,shape,fill,status,workers))
96
97def bias_selected(accepted: Bool,+di: ImageGeometry,g: U32,d: Storage.Buffer<F32>,+shape: List<&2,U32>,fill: F32,status: Tensors.Status,+workers: U32) -> Tensors.Tensor<F32> & Tensors.Tensor<F32>:
98 match accepted:
99 case False{}: (Tensors.failed(~F32,"conv2d_bias_gradient",Tensors.InvalidArgument{shape},[],0.0,workers),Tensors.whole(~F32,d,shape,fill,status,workers))
100 case True{}:
101 ImageGeometry{n,+o,rows,columns} = di
102 bias_counted(Tensors.count_of([o]),ImageGeometry{n,o,rows,columns},g,d,shape,fill,status,workers)
103
104def bias_shaped(di: Maybe<&2,ImageGeometry>,+g: U32,d: Storage.Buffer<F32>,+shape: List<&2,U32>,fill: F32,+status: Tensors.Status,workers: U32) -> Tensors.Tensor<F32> & Tensors.Tensor<F32>:
105 match di:
106 case None{}: bias_selected(False{},ImageGeometry{0,0,0,0},g,d,shape,fill,status,workers)
107 case Some{+di}: bias_selected(Tensors.usable(status) && positive(di) && U32.is_lt(0,g),di,g,d,shape,fill,status,workers)
108
109def bias(d: Tensors.Tensor<F32>,+g: U32) -> Tensors.Tensor<F32> & Tensors.Tensor<F32>:
110 Tensors.with_whole(~F32,~(Tensors.Tensor<F32> & Tensors.Tensor<F32>),d,d => shape => fill => status => workers =>
111 +shape = shape
112 bias_shaped(image(shape),g,d,shape,fill,status,workers))