1	# Bias gradient traversal over ordered spatial groups.
2	import Base
3	import ./traversal_grouped.bend as Grouped
4	import ./storage_buffer.bend as Storage
5	
6	type Work is Data:
7	  Work{images: Nat,channels: Nat,outputs: Nat,height: Nat,width: Nat,kernel: Nat,stride: Nat,padding: Nat,rows: Nat,columns: Nat,group: Nat}
8	
9	type Context is Data:
10	  Context{work: Work,image: Nat,element: Nat}
11	
12	def bias_address(context: Context,position: Nat) -> U32:
13	  Context{Work{n,c,+o,h,w,k,s,p,+rows,+columns,g},image,element} = context
14	  U32.from_nat(Nat.add(Nat.mul(Nat.mul(Nat.add(Nat.mul(image,o),element),rows),columns),position))
15	
16	def bias_at(reader: Array<F32>,context: Context,position: Nat) -> Array<F32> & F32:
17	  Array.get(F32,reader,bias_address(context,position))
18	
19	def image_tail(~Reader: Type,pair: Reader & F32,tail: Reader -> Reader & List<&2,F32>) -> Reader & List<&2,F32>:
20	  (reader,value) = pair
21	  Grouped.prepend(~Reader,value,tail(reader))
22	
23	def image_values(~Reader: Type,~read: Reader -> Context -> Nat -> Reader & F32,remaining: Nat,+image: Nat,+element: Nat,+work: Work,
24	  +ranges: List<&2,Grouped.Range>,reader: Reader) -> Reader & List<&2,F32>:
25	  match remaining:
26	    case 0n: (reader,[])
27	    case 1n+rest:
28	      image_tail(~Reader,Grouped.run(~Reader,~Context,~read,ranges,reader,Context{work,image,element}),
29	        reader => image_values(~Reader,~read,rest,1n+image,element,work,ranges,reader))
30	
31	def element(~Reader: Type,~read: Reader -> Context -> Nat -> Reader & F32,+work: Work,index: Nat,+ranges: List<&2,Grouped.Range>,reader: Reader) -> Reader & F32:
32	  Work{images,c,o,h,w,k,s,p,rows,columns,g} = work
33	  Grouped.reduced(~Reader,image_values(~Reader,~read,images,0n,index,work,ranges,reader))
34	
35	# One reduced output written at its position, then the rest of the outputs.
36	def stored(~Reader: Type,+position: U32,result: Storage.Buffer<F32>,pair: Reader & F32,
37	  tail: Reader -> Storage.Buffer<F32> -> Reader & Storage.Buffer<F32>) -> Reader & Storage.Buffer<F32>:
38	  (reader,value) = pair
39	  tail(reader,Storage.write(~F32,position,value,result))
40	
41	def written(~Reader: Type,~read: Reader -> Context -> Nat -> Reader & F32,remaining: Nat,+index: Nat,+position: U32,+work: Work,+ranges: List<&2,Grouped.Range>,
42	  reader: Reader,result: Storage.Buffer<F32>) -> Reader & Storage.Buffer<F32>:
43	  match remaining:
44	    case 0n: (reader,result)
45	    case 1n+rest:
46	      stored(~Reader,position,result,element(~Reader,~read,work,index,ranges,reader),
47	        reader => result => written(~Reader,~read,rest,1n+index,(position + 1 : U32),work,ranges,reader,result))
48	
49	def bias(+work: Work,d: Array<F32>,result: Storage.Buffer<F32>) -> Array<F32> & Storage.Buffer<F32>:
50	  Work{n,c,o,h,w,k,s,p,rows,columns,g} = work
51	  written(~Array<F32>,~bias_at,o,0n,0,work,Grouped.ranges(rows,columns,g),d,result)