1	# Gradients of conv2d in a fixed FP32 order. Every function gives the
2	# tensors it only reads back to its caller.
3	import Base
4	import ./tensor.bend as Tensors
5	import ./tensor_band.bend as Bands
6	import ./tensor_convolution.bend as TensorConvolution
7	import ./kernel_input_gradient.bend as InputKernel
8	import ./kernel_shape.bend as KernelShape
9	import ./storage_buffer.bend as Storage
10	import ./default_machine_limits.bend as Default
11	import ./machine_limits.bend as Machine
12	import ./tensor_parameter_gradient.bend as ParameterGradient
13	
14	def Tensor() -> Type:
15	  Tensors.Tensor<F32>
16	
17	type Arguments is Data:
18	  Arguments{geometry: InputKernel.Geometry,batch: U32,batched: Bool}
19	
20	def input_dimensions(arguments: Arguments) -> KernelShape.Dimensions:
21	  Arguments{InputKernel.Geometry{c,o,h,w,k,s,p,gh,gw},batch,batched} = arguments
22	  KernelShape.Dimensions{U32.to_nat(c),U32.to_nat(o),U32.to_nat(h),U32.to_nat(w),U32.to_nat(k),U32.to_nat(s),U32.to_nat(p)}
23	
24	# The bands of the input gradient: rows of every image of the batch.
25	def input_band_geometry(arguments: Arguments) -> Bands.Geometry:
26	  Arguments{InputKernel.Geometry{c,o,h,w,k,s,p,gh,gw},batch,batched} = arguments
27	  Bands.Geometry{batch,c,h,w}
28	
29	def input_shape(arguments: Arguments) -> List<&2,U32>:
30	  Arguments{InputKernel.Geometry{c,o,h,w,k,s,p,gh,gw},batch,batched} = arguments
31	  TensorConvolution.output_shape(batched,batch,c,h,w)
32	
33	# The limits are an argument so that a proof can keep them symbolic.
34	def input_valid_with(+arguments: Arguments,d_length: Nat,+weight_length: Nat,+limits: Machine.Limits) -> Bool:
35	  Arguments{+geometry,+batch,batched} = arguments
36	  InputKernel.Geometry{+c,+o,+h,+w,+k,+s,+p,+gh,+gw} = geometry
37	  +dimensions = input_dimensions(arguments)
38	  +input_length = Nat.mul(U32.to_nat(c),Nat.mul(U32.to_nat(h),U32.to_nat(w)))
39	  KernelShape.valid(dimensions,input_length,weight_length,U32.to_nat(o)) &&
40	    KernelShape.supports_direct(limits,dimensions,input_length,weight_length) &&
41	    Nat.is_eq(U32.to_nat(gh),KernelShape.axis_size(U32.to_nat(h),U32.to_nat(k),U32.to_nat(s),U32.to_nat(p))) &&
42	    Nat.is_eq(U32.to_nat(gw),KernelShape.axis_size(U32.to_nat(w),U32.to_nat(k),U32.to_nat(s),U32.to_nat(p))) &&
43	    Nat.is_eq(d_length,Nat.mul(U32.to_nat(batch),Nat.mul(U32.to_nat(o),Nat.mul(U32.to_nat(gh),U32.to_nat(gw))))) &&
44	    InputKernel.fits(geometry,Machine.capacity(limits))
45	
46	def input_valid(arguments: Arguments,d_length: Nat,weight_length: Nat) -> Bool:
47	  input_valid_with(arguments,d_length,weight_length,Default.limits())
48	
49	def input_result(+arguments: Arguments,+workers: U32,+weight_shape: List<&2,U32>,+weight_fill: F32,+weight_workers: U32,result: InputKernel.Outcome) -> Tensor() & Tensor():
50	  InputKernel.Outcome{weights,output} = result
51	  (Tensors.Tensor{output,input_band_geometry(arguments),input_shape(arguments),0.0,Tensors.Usable{},workers},Tensors.whole(~F32,weights,weight_shape,weight_fill,Tensors.Usable{},weight_workers))
52	
53	def input_checked(accepted: Bool,+arguments: Arguments,+count: U32,gradient: Storage.Buffer<F32>,weights: Storage.Buffer<F32>,+workers: U32,+weight_shape: List<&2,U32>,+weight_fill: F32,+weight_workers: U32) -> Tensor() & Tensor():
54	  match accepted:
55	    case False{}:
56	      (Tensors.failed(~F32,"conv2d_input_gradient",Tensors.InvalidArgument{input_shape(arguments)},input_shape(arguments),0.0,workers),Tensors.whole(~F32,weights,weight_shape,weight_fill,Tensors.Usable{},weight_workers))
57	    case True{}:
58	      Arguments{+geometry,batch,batched} = arguments
59	      input_result(arguments,workers,weight_shape,weight_fill,weight_workers,InputKernel.run(geometry,U32.to_nat(count),U32.max(workers,U32.max(weight_workers,1)),gradient,weights))
60	
61	def input_count(count: Maybe<&2,U32>,+arguments: Arguments,gradient: Storage.Buffer<F32>,weights: Storage.Buffer<F32>,+workers: U32,+weight_shape: List<&2,U32>,+weight_fill: F32,+weight_workers: U32) -> Tensor() & Tensor():
62	  match count:
63	    case None{}:
64	      (Tensors.failed(~F32,"conv2d_input_gradient",Tensors.InvalidShape{input_shape(arguments)},input_shape(arguments),0.0,workers),Tensors.whole(~F32,weights,weight_shape,weight_fill,Tensors.Usable{},weight_workers))
65	    case Some{+count}:
66	      Storage.Buffer{d,+d_length,d_certificate} = gradient
67	      Storage.Buffer{w,+weight_length,w_certificate} = weights
68	      input_checked(input_valid(arguments,d_length,weight_length),arguments,count,gradient,weights,workers,weight_shape,weight_fill,weight_workers)
69	
70	def input_geometry(+stride: U32,+padding: U32,+height: U32,+width: U32,shape: Maybe<&2,TensorConvolution.InputShape>,+weight_shape: List<&2,U32>,gradient: Storage.Buffer<F32>,weights: Storage.Buffer<F32>,+workers: U32,+weight_fill: F32,+weight_workers: U32) -> Tensor() & Tensor():
71	  match shape weight_shape:
72	    case Some{TensorConvolution.InputShape{+batch,+batched,+outputs,+gh,+gw}} Con{wo,Con{+channels,Con{+kernel,Con{kw,Nil{}}}}}:
73	      +arguments = {Arguments{InputKernel.Geometry{channels,outputs,height,width,kernel,stride,padding,gh,gw},batch,batched} : Arguments}
74	      input_count(Bool.pick(Maybe<&2,U32>,U32.is_eq(outputs,wo) && U32.is_eq(kernel,kw),Tensors.count_of(input_shape(arguments)),None{}),arguments,gradient,weights,workers,weight_shape,weight_fill,weight_workers)
75	    case shape weight_shape:
76	      (Tensors.failed(~F32,"conv2d_input_gradient",Tensors.InvalidShape{weight_shape},weight_shape,0.0,workers),Tensors.whole(~F32,weights,weight_shape,weight_fill,Tensors.Usable{},weight_workers))
77	
78	def input_status(status: Tensors.Status,+stride: U32,+padding: U32,+height: U32,+width: U32,gradient: Storage.Buffer<F32>,+d_shape: List<&2,U32>,+d_fill: F32,+workers: U32,weights: Storage.Buffer<F32>,+weight_shape: List<&2,U32>,+weight_fill: F32,+weight_status: Tensors.Status,+weight_workers: U32) -> Tensor() & Tensor():
79	  match status:
80	    case Tensors.Failed{problem}:
81	      (Tensors.whole(~F32,gradient,d_shape,d_fill,Tensors.Failed{problem},workers),Tensors.whole(~F32,weights,weight_shape,weight_fill,weight_status,weight_workers))
82	    case Tensors.Usable{}:
83	      input_geometry(stride,padding,height,width,TensorConvolution.input_shape(d_shape),weight_shape,gradient,weights,workers,weight_fill,weight_workers)
84	
85	# d [N,O,Ho,Wo] or [O,Ho,Wo] and weights [O,C,k,k] of a convolution of an
86	# input of the given height and width: (the gradient of that input, the
87	# weights). Several input sizes give the same output size when the stride is
88	# above 1, so height and width are given.
89	def input(d: Tensor(),weights: Tensor(),+stride: U32,+padding: U32,+height: U32,+width: U32) -> Tensor() & Tensor():
90	  Tensors.with_whole(~F32,~(Tensor() & Tensor()),d,gradient => d_shape => d_fill => d_status => workers =>
91	    Tensors.with_whole(~F32,~(Tensor() & Tensor()),weights,weights => weight_shape => weight_fill => weight_status => weight_workers =>
92	      +weight_status = weight_status
93	      input_status(Tensors.first_failure(d_status,weight_status),stride,padding,height,width,gradient,d_shape,d_fill,workers,weights,weight_shape,weight_fill,weight_status,weight_workers)))
94	
95	# The group length of the weight and bias gradient contract.
96	def group_length() -> U32:
97	  256
98	
99	# x [N,C,H,W], d [N,O,Ho,Wo], or both without their batch axes.
100	# Returns (weight gradient [O,C,k,k], (x,d)); all sources are read only.
101	def weight(x: Tensor(),d: Tensor(),kernel: U32,stride: U32,padding: U32) -> Tensor() & (Tensor() & Tensor()):
102	  ParameterGradient.weight(x,d,kernel,stride,padding,group_length())
103	
104	# d [N,O,Ho,Wo] or [O,Ho,Wo]: (bias gradient [O], d).
105	def bias(d: Tensor()) -> Tensor() & Tensor():
106	  ParameterGradient.bias(d,group_length())