tensor_indexed.bend

stelliferous@0.0.2.0 · 80 lines · 14 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.2.0/traversal_array.bend as Traversal
5import stelliferous@0.0.2.0/traversal_loop.bend as Loop
6import stelliferous@0.0.2.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)