~/jendhub

def allocation_elements

stelliferous@0.0.2.0 · tensor_band_schedule.bend:40

import stelliferous@0.0.2.0/tensor_band_schedule.bend as TensorBandSchedule
import Base
import stelliferous@0.0.2.0/tensor_band.bend as Bands
import stelliferous@0.0.2.0/storage_buffer.bend as Storage

40# For equal-shaped operands the logical read/write volume of repartition is
41# identical. The target owners' rounded capacities decide the remaining cost.
42def allocation_elements(layout: Bands.Plan,+geometry: Bands.Geometry) -> Nat:
43  match layout:
44    case Bands.Rows{first,rows}: U32.to_nat(U32.shln(1,Storage.depth_for(Bands.local_count(geometry,rows))))
45    case Bands.Split{left,right}: Nat.add(allocation_elements(left,geometry),allocation_elements(right,geometry))

Uses (8)

Used by (1)