def each

stelliferous@0.0.2.0 · tensor_band.bend:174

import Base
import stelliferous@0.0.2.0/storage_buffer.bend as Storage
import stelliferous@0.0.2.0/traversal_partition.bend as Partition

174# The one band recursion. The plan gives every leaf its rows and decides the
175# structure; a tree of another shape is kept. A split inside one image may
176# run its halves in parallel, and a split across images runs them in order.
177# share gives each half its own owner of a read-only resource.
178def each(~Block: Type,~Output: Type,~Resource: Type,~Context: Data,
179  ~share: Resource -> Resource & Resource,
180  ~leaf: Context -> U32 -> U32 -> Block -> Resource -> Partition.Tree<&1,Output,Unit>,
181  ~keep: Partition.Tree<&1,Block,Unit> -> Partition.Tree<&1,Output,Unit>,
182  layout: Plan,mode: Bool,+enabled: Bool,+geometry: Geometry,+context: Context,
183  tree: Partition.Tree<&1,Block,Unit>,resource: Resource) -> Partition.Tree<&1,Output,Unit>:
184  match layout mode tree:
185    case Rows{first,rows} mode tree: leaf_task(~Block,~Output,~Resource,~Context,~leaf,~keep,first,rows,context,tree,resource)
186    case Split{+left_plan,+right_plan} False{} Partition.Fork{left,right}:
187      Storage.bind_pair(~Resource,~Resource,~Partition.Tree<&1,Output,Unit>,share(resource),left_resource => right_resource =>
188        Partition.Fork{each(~Block,~Output,~Resource,~Context,~share,~leaf,~keep,left_plan,enabled && same_image(geometry,left_plan),enabled,geometry,context,left,left_resource),
189          each(~Block,~Output,~Resource,~Context,~share,~leaf,~keep,right_plan,enabled && same_image(geometry,right_plan),enabled,geometry,context,right,right_resource)})
190    case Split{Rows{left_first,left_rows},Rows{right_first,right_rows}} True{} Partition.Fork{left,right}:
191      Storage.bind_pair(~Resource,~Resource,~Partition.Tree<&1,Output,Unit>,share(resource),left_resource => right_resource =>
192        left right = leaf_task(~Block,~Output,~Resource,~Context,~leaf,~keep,left_first,left_rows,context,left,left_resource) leaf_task(~Block,~Output,~Resource,~Context,~leaf,~keep,right_first,right_rows,context,right,right_resource)
193        Partition.Fork{left,right})
194    case Split{Rows{left_first,left_rows},+right_plan} True{} Partition.Fork{left,right}:
195      Storage.bind_pair(~Resource,~Resource,~Partition.Tree<&1,Output,Unit>,share(resource),left_resource => right_resource =>
196        left right = leaf_task(~Block,~Output,~Resource,~Context,~leaf,~keep,left_first,left_rows,context,left,left_resource) each(~Block,~Output,~Resource,~Context,~share,~leaf,~keep,right_plan,enabled && same_image(geometry,right_plan),enabled,geometry,context,right,right_resource)
197        Partition.Fork{left,right})
198    case Split{+left_plan,Rows{right_first,right_rows}} True{} Partition.Fork{left,right}:
199      Storage.bind_pair(~Resource,~Resource,~Partition.Tree<&1,Output,Unit>,share(resource),left_resource => right_resource =>
200        left right = each(~Block,~Output,~Resource,~Context,~share,~leaf,~keep,left_plan,enabled && same_image(geometry,left_plan),enabled,geometry,context,left,left_resource) leaf_task(~Block,~Output,~Resource,~Context,~leaf,~keep,right_first,right_rows,context,right,right_resource)
201        Partition.Fork{left,right})
202    case Split{+left_plan,+right_plan} True{} Partition.Fork{left,right}:
203      Storage.bind_pair(~Resource,~Resource,~Partition.Tree<&1,Output,Unit>,share(resource),left_resource => right_resource =>
204        left right = each(~Block,~Output,~Resource,~Context,~share,~leaf,~keep,left_plan,enabled && same_image(geometry,left_plan),enabled,geometry,context,left,left_resource) each(~Block,~Output,~Resource,~Context,~share,~leaf,~keep,right_plan,enabled && same_image(geometry,right_plan),enabled,geometry,context,right,right_resource)
205        Partition.Fork{left,right})
206    case layout mode tree: keep(tree)