tensor_indexed.bend
stelliferous@0.0.3.0 · 80 lines · 0 definitions
1# Indexed mapping with a runtime context and a closed index template. The
2# source and destination traversal are shared with checked tensor operations.
3import Base
4import stelliferous@0.0.3.0/traversal_array.bend as Traversal
5import stelliferous@0.0.3.0/traversal_loop.bend as Loop
6import stelliferous@0.0.3.0/default_machine_limits.bend as Default
7
8def value(~Element: Data,~Context: Data,~index: Context -> U32 -> U32,~operation: U32 -> Element -> Element,
9 context: Context,position: U32,result: Array<Element> & Element) -> Array<Element> & Element:
10 (source,element) = result
11 (source,operation(index(context,position),element))
12
13def read(~Element: Data,~Context: Data,~index: Context -> U32 -> U32,~operation: U32 -> Element -> Element,
14 source: Array<Element>,context: Context,position: U32) -> Array<Element> & Element:
15 +position = position
16 value(~Element,~Context,~index,~operation,context,position,Array.get(Element,source,position))
17
18def destination(~Context: Data,context: Context,position: U32) -> U32:
19 position
20
21type Call<-Context: Data> is Data:
22 Call{count: Nat,context: Context}
23
24def run(~Element: Data,~Context: Data,~index: Context -> U32 -> U32,~operation: U32 -> Element -> Element,
25 call: Call<Context>,input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
26 Call{count,context} = call
27 Traversal.finish(~Array<Element>,~Element,~Context,
28 Traversal.run(~Array<Element>,~Element,~Context,~read(~Element,~Context,~index,~operation),~destination(~Context),count,
29 Traversal.State{input,output,0,context}))
30
31# The same mapping over blocks of consecutive positions. A block starts where
32# the previous one ended and gives position p the index base + (p - start),
33# so an index that is a quotient and a remainder of the position costs one
34# base per block instead of a division per element.
35type Block is Data:
36 Block{base: U32,start: U32}
37
38def block_index(block: Block,position: U32) -> U32:
39 Block{base,start} = block
40 (base + (position - start) : U32)
41
42type Blocks<-Context: Data> is Data:
43 Blocks{block: U32,size: U32,context: Context}
44
45def next_block(~Element: Data,~Context: Data,block: U32,size: U32,context: Context,
46 state: Traversal.State<&1,&1,Array<Element>,Array<Element>,Block>) -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Blocks<Context>>:
47 Traversal.State{input,output,position,inner} = state
48 Traversal.State{input,output,position,Blocks{(block + 1 : U32),size,context}}
49
50def block_step(~Element: Data,~Context: Data,~base: Context -> U32 -> U32,~operation: U32 -> Element -> Element,
51 state: Traversal.State<&1,&1,Array<Element>,Array<Element>,Blocks<Context>>) -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Blocks<Context>>:
52 Traversal.State{input,output,+position,Blocks{+block,+size,+context}} = state
53 next_block(~Element,~Context,block,size,context,
54 Traversal.run(~Array<Element>,~Element,~Block,~read(~Element,~Block,~block_index,~operation),~destination(~Block),U32.to_nat(size),
55 Traversal.State{input,output,position,Block{base(context,block),position}}))
56
57def run_blocks(~Element: Data,~Context: Data,~base: Context -> U32 -> U32,~operation: U32 -> Element -> Element,
58 blocks: Nat,size: U32,context: Context,input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
59 Traversal.finish(~Array<Element>,~Element,~Blocks<Context>,
60 Loop.run(~Traversal.State<&1,&1,Array<Element>,Array<Element>,Blocks<Context>>,~block_step(~Element,~Context,~base,~operation),blocks,
61 Traversal.State{input,output,0,Blocks{0,size,context}}))
62
63def run_whole(~Element: Data,~Context: Data,~index: Context -> U32 -> U32,~base: Context -> U32 -> U32,~operation: U32 -> Element -> Element,
64 whole: Bool,blocks: U32,size: U32,total: U32,context: Context,input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
65 match whole:
66 case True{}: run_blocks(~Element,~Context,~base,~operation,U32.to_nat(blocks),size,context,input,output)
67 case False{}: run(~Element,~Context,~index,~operation,Call{U32.to_nat(total),context},input,output)
68
69# Blocks are used only up to the machine capacity, where the block size
70# divides a position in U32 without overflow. No storage is larger.
71def block_limit() -> U32:
72 Default.capacity()
73
74# total positions in blocks(context) blocks of size(context) positions. The
75# blocks are used only when they are exactly the total and the total is
76# within the limit; any other count keeps the index function.
77def run_tiled(~Element: Data,~Context: Data,~index: Context -> U32 -> U32,~base: Context -> U32 -> U32,~blocks: Context -> U32,~size: Context -> U32,
78 ~operation: U32 -> Element -> Element,+total: U32,+context: Context,input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
79 run_whole(~Element,~Context,~index,~base,~operation,U32.is_le(total,block_limit()) && Nat.is_eq(Nat.mul(U32.to_nat(blocks(context)),U32.to_nat(size(context))),U32.to_nat(total)),
80 blocks(context),size(context),total,context,input,output)