1 # Shape validation and affine tensor ownership for the shared reduction kernel. 2 import Base 3 import ./tensor.bend as Tensors 4 import ./storage_buffer.bend as Storage 5 import ./kernel_parameter_gradient.bend as Kernel 6 import ./kernel_parameter_gradient_buffer.bend as KernelBuffer 7 import ./kernel_weight_gradient.bend as WeightKernel 8 import ./kernel_shape.bend as Shape 9 10 type ImageGeometry is Data: 11 ImageGeometry{images: U32,channels: U32,height: U32,width: U32} 12 13 def 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 19 def 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 23 def 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 28 def fits(count: Maybe<&2,U32>) -> Bool: 29 match count: 30 case None{}: False{} 31 case Some{count}: True{} 32 33 def 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 42 def 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 & Storage.Buffer) & Storage.Buffer) -> Tensors.Tensor & (Tensors.Tensor & Tensors.Tensor): 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 48 def weight_selected(accepted: Bool,+xi: ImageGeometry,+di: ImageGeometry,+k: U32,s: U32,p: U32,g: U32, 49 x: Storage.Buffer,+xs: List<&2,U32>,xf: F32,xs_status: Tensors.Status,+xw: U32, 50 d: Storage.Buffer,ds: List<&2,U32>,df: F32,ds_status: Tensors.Status,+dw: U32) -> Tensors.Tensor & (Tensors.Tensor & Tensors.Tensor): 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. 62 def 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,+xs: List<&2,U32>,xf: F32,+xw: U32, 64 d: Storage.Buffer,ds: List<&2,U32>,df: F32,dw: U32) -> Tensors.Tensor & (Tensors.Tensor & Tensors.Tensor): 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 71 def weight(x: Tensors.Tensor,d: Tensors.Tensor,+k: U32,+s: U32,+p: U32,+g: U32) -> Tensors.Tensor & (Tensors.Tensor & Tensors.Tensor): 72 Tensors.with_whole(~F32,~(Tensors.Tensor & (Tensors.Tensor & Tensors.Tensor)),x,x => xs => xf => status => xw => 73 +xs = xs 74 Tensors.with_whole(~F32,~(Tensors.Tensor & (Tensors.Tensor & Tensors.Tensor)),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 78 def 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 82 def bias_finished(+shape: List<&2,U32>,fill: F32,status: Tensors.Status,+workers: U32,+outputs: U32,result: Storage.Buffer & Storage.Buffer) -> Tensors.Tensor & Tensors.Tensor: 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]. 87 def bias_counted(count: Maybe<&2,U32>,+di: ImageGeometry,+g: U32,d: Storage.Buffer,+shape: List<&2,U32>,fill: F32,status: Tensors.Status,+workers: U32) -> 88 Tensors.Tensor & Tensors.Tensor: 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 97 def bias_selected(accepted: Bool,+di: ImageGeometry,g: U32,d: Storage.Buffer,+shape: List<&2,U32>,fill: F32,status: Tensors.Status,+workers: U32) -> Tensors.Tensor & Tensors.Tensor: 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 104 def bias_shaped(di: Maybe<&2,ImageGeometry>,+g: U32,d: Storage.Buffer,+shape: List<&2,U32>,fill: F32,+status: Tensors.Status,workers: U32) -> Tensors.Tensor & Tensors.Tensor: 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 109 def bias(d: Tensors.Tensor,+g: U32) -> Tensors.Tensor & Tensors.Tensor: 110 Tensors.with_whole(~F32,~(Tensors.Tensor & Tensors.Tensor),d,d => shape => fill => status => workers => 111 +shape = shape 112 bias_shaped(image(shape),g,d,shape,fill,status,workers))