def parallel

stelliferous@0.0.2.0 · tensor_band_pool.bend:30

import Base
import stelliferous@0.0.2.0/tensor_access.bend as Access
import stelliferous@0.0.2.0/tensor_band.bend as Bands
import stelliferous@0.0.2.0/tensor_band_window.bend as BandWindow
import stelliferous@0.0.2.0/tensor_band_schedule.bend as BandSchedule
import stelliferous@0.0.2.0/kernel_cost.bend as Cost
import stelliferous@0.0.2.0/default_kernel_cost.bend as DefaultCost

30# Bytes price loading, storing and the loop that moves them (copy_byte_ns is
31# measured on copy loops); operations count the rest. Per sample
32# (Access.read_line and the operation): an offset and a subtraction, one bound
33# check, a two-operation index, the padding select and the compare-and-select
34# maximum: 8. Each line adds ten for its origin (Access.line_at). An output is
35# one line along the rows; the lines along the columns are one for every
36# window row, stride for every output row.
37def parallel(window: BandWindow.Geometry,workers: U32,size: BandSchedule.Size) -> Bool:
38  BandWindow.Geometry{Bands.Geometry{batch,channels,ih,iw},kernel,stride,padding,oh,ow} = window
39  +config = DefaultCost.configuration()
40  +samples = U32.to_f32(kernel)
41  +lines = (U32.to_f32(stride) + 1.0 : F32)
42  +line_bytes = (samples * 4.0 + 8.0 : F32)
43  +line_operations = (samples * 8.0 + 10.0 : F32)
44  BandSchedule.sized(Bands.Geometry{batch,channels,oh,ow},workers,Cost.element_ns((lines * line_bytes : F32),(lines * line_operations : F32),config),config,size)