tensor_indexed.bend source

tensor_indexed.bend on the hub · documented module

# Indexed mapping with a runtime context and a closed index template. The# source and destination traversal are shared with checked tensor operations.import Baseimport ./traversal_array.bend as Traversalimport ./traversal_loop.bend as Loopimport ./default_machine_limits.bend as Defaultdef value(~Element: Data,~Context: Data,~index: Context -> U32 -> U32,~operation: U32 -> Element -> Element,  context: Context,position: U32,result: Array<Element> & Element) -> Array<Element> & Element:  (source,element) = result  (source,operation(index(context,position),element))def read(~Element: Data,~Context: Data,~index: Context -> U32 -> U32,~operation: U32 -> Element -> Element,  source: Array<Element>,context: Context,position: U32) -> Array<Element> & Element:  +position = position  value(~Element,~Context,~index,~operation,context,position,Array.get(Element,source,position))def destination(~Context: Data,context: Context,position: U32) -> U32:  positiontype Call<-Context: Data> is Data:  Call{count: Nat,context: Context}def run(~Element: Data,~Context: Data,~index: Context -> U32 -> U32,~operation: U32 -> Element -> Element,  call: Call<Context>,input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:  Call{count,context} = call  Traversal.finish(~Array<Element>,~Element,~Context,    Traversal.run(~Array<Element>,~Element,~Context,~read(~Element,~Context,~index,~operation),~destination(~Context),count,      Traversal.State{input,output,0,context}))# The same mapping over blocks of consecutive positions. A block starts where# the previous one ended and gives position p the index base + (p - start),# so an index that is a quotient and a remainder of the position costs one# base per block instead of a division per element.type Block is Data:  Block{base: U32,start: U32}def block_index(block: Block,position: U32) -> U32:  Block{base,start} = block  (base + (position - start) : U32)type Blocks<-Context: Data> is Data:  Blocks{block: U32,size: U32,context: Context}def next_block(~Element: Data,~Context: Data,block: U32,size: U32,context: Context,  state: Traversal.State<&1,&1,Array<Element>,Array<Element>,Block>) -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Blocks<Context>>:  Traversal.State{input,output,position,inner} = state  Traversal.State{input,output,position,Blocks{(block + 1 : U32),size,context}}def block_step(~Element: Data,~Context: Data,~base: Context -> U32 -> U32,~operation: U32 -> Element -> Element,  state: Traversal.State<&1,&1,Array<Element>,Array<Element>,Blocks<Context>>) -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Blocks<Context>>:  Traversal.State{input,output,+position,Blocks{+block,+size,+context}} = state  next_block(~Element,~Context,block,size,context,    Traversal.run(~Array<Element>,~Element,~Block,~read(~Element,~Block,~block_index,~operation),~destination(~Block),U32.to_nat(size),      Traversal.State{input,output,position,Block{base(context,block),position}}))def run_blocks(~Element: Data,~Context: Data,~base: Context -> U32 -> U32,~operation: U32 -> Element -> Element,  blocks: Nat,size: U32,context: Context,input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:  Traversal.finish(~Array<Element>,~Element,~Blocks<Context>,    Loop.run(~Traversal.State<&1,&1,Array<Element>,Array<Element>,Blocks<Context>>,~block_step(~Element,~Context,~base,~operation),blocks,      Traversal.State{input,output,0,Blocks{0,size,context}}))def run_whole(~Element: Data,~Context: Data,~index: Context -> U32 -> U32,~base: Context -> U32 -> U32,~operation: U32 -> Element -> Element,  whole: Bool,blocks: U32,size: U32,total: U32,context: Context,input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:  match whole:    case True{}: run_blocks(~Element,~Context,~base,~operation,U32.to_nat(blocks),size,context,input,output)    case False{}: run(~Element,~Context,~index,~operation,Call{U32.to_nat(total),context},input,output)# Blocks are used only up to the machine capacity, where the block size# divides a position in U32 without overflow. No storage is larger.def block_limit() -> U32:  Default.capacity()# total positions in blocks(context) blocks of size(context) positions. The# blocks are used only when they are exactly the total and the total is# within the limit; any other count keeps the index function.def run_tiled(~Element: Data,~Context: Data,~index: Context -> U32 -> U32,~base: Context -> U32 -> U32,~blocks: Context -> U32,~size: Context -> U32,  ~operation: U32 -> Element -> Element,+total: U32,+context: Context,input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:  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)),    blocks(context),size(context),total,context,input,output)