src/Batch.bend

bend-edit-distance@0.2.2.0 · 189 lines · 26 definitions

1import Base
2import bend-edit-distance@0.2.2.0/EdlibTypes.bend as E
3import bend-edit-distance@0.2.2.0/src/Myers.bend as M
4
5type WeightedJob is Data:
6  WeightedJob{index: Nat, work: F32, job: E.AlignJob}
7
8type Group is Data:
9  Group{room: Nat, work: F32, jobs: List<&2, WeightedJob>}
10
11# Build the entire plan before entering the alignment fork tree. Recursive
12# argument preparation in a parallel call can drain the native frontier.
13type Plan is Type:
14  Empty{}
15  Leaf{index: Nat, job: E.AlignJob}
16  Fork{left: Plan, right: Plan}
17  Serial{jobs: List<&2, WeightedJob>}
18
19type IndexedResult is Data:
20  IndexedResult{index: Nat, result: E.AlignResult}
21
22def axes(mode: E.Mode, +qn: Nat, +tn: Nat) -> Nat & Nat:
23  match mode:
24    case E.NW{}:
25      (Nat.min(qn, tn), Nat.max(qn, tn))
26    case _:
27      (qn, tn)
28
29def extra_work(mode: E.Mode, task: E.Task, +q: F32, +t: F32,
30               +blocks: F32, pairs: F32) -> F32:
31  match mode task:
32    case _ E.Distance{}:
33      0.0
34    case E.HW{} E.Path{}:
35      (blocks * t * t + q * t * pairs : F32)
36    case E.HW{} _:
37      (blocks * t * t : F32)
38    case _ E.Path{}:
39      (q * t * pairs : F32)
40    case _ _:
41      0.0
42
43def estimate_axes(+mode: E.Mode, task: E.Task, pairs: List<&2, E.EqualityPair>,
44                  sizes: Nat & Nat) -> F32:
45  (+length, columns) = sizes
46  +l = F32.from_nat(length)
47  +c = F32.from_nat(columns)
48  +blocks = F32.from_nat(M.ceil_div(length, 32n))
49  +p = (1.0 + F32.from_nat(List.length(&2, E.EqualityPair, pairs)) : F32)
50  (1.0 + l + c + 256.0 * l * p + blocks * c +
51    extra_work(mode, task, l, c, blocks, p) : F32)
52
53# A scheduling heuristic, not a timing prediction: input scans, 256-symbol
54# profiles, full Myers columns, worst-case HW start searches and PATH DP.
55# F32 avoids wrapping a U32 product or allocating a quadratic unary Nat.
56# Bands, retries, similarity and the selected interval can change actual work.
57def estimate_sizes(qn: Nat, tn: Nat, config: E.AlignConfig) -> F32:
58  match qn tn config:
59    case 0n +t _:
60      (1.0 + F32.from_nat(t) : F32)
61    case +q 0n _:
62      (1.0 + F32.from_nat(q) : F32)
63    case +q +t E.AlignConfig{+mode, task, _, pairs, _}:
64      estimate_axes(mode, task, pairs, axes(mode, q, t))
65
66def estimate(+job: E.AlignJob) -> F32:
67  E.AlignJob{query, target, config} = job
68  # Scans are sequential here; pairs of arbitrary lengths need not balance.
69  estimate_sizes(String.length(query), String.length(target), config)
70
71def indexed(jobs: List<&2, E.AlignJob>, +index: Nat) -> List<&2, WeightedJob>:
72  match jobs:
73    case Nil{}:
74      Nil{}
75    case +job <> tail:
76      WeightedJob{index, estimate(job), job} <> indexed(tail, 1n+index)
77
78def heavier(first: WeightedJob, second: WeightedJob) -> Bool:
79  WeightedJob{_, a, _} = first
80  WeightedJob{_, b, _} = second
81  F32.is_ge(a, b)
82
83def add(group: Group, job: WeightedJob) -> Group:
84  Group{room, work, jobs} = group
85  WeightedJob{index, +cost, value} = job
86  Group{Nat.sub(room, 1n), F32.add(work, cost), WeightedJob{index, cost, value} <> jobs}
87
88def choose(left: Group, right: Group) -> Bool:
89  match left right:
90    case Group{0n, _, _} _:
91      False{}
92    case _ Group{0n, _, _}:
93      True{}
94    case Group{_, a, _} Group{_, b, _}:
95      F32.is_le(a, b)
96
97def place(job: WeightedJob, left: Group, right: Group, use_left: Bool) -> Group & Group:
98  match use_left:
99    case True{}:
100      (add(left, job), right)
101    case False{}:
102      (left, add(right, job))
103
104# Largest estimated jobs go to the lighter branch. Capacity bounds keep both
105# child counts within ceil(n/2), including zero/rounded/equal work estimates.
106def partition(jobs: List<&2, WeightedJob>, groups: Group & Group) -> Group & Group:
107  match jobs:
108    case Nil{}:
109      groups
110    case job <> tail:
111      (+left, +right) = groups
112      partition(tail, place(job, left, right, choose(left, right)))
113
114def group_jobs(group: Group) -> List<&2, WeightedJob>:
115  Group{_, _, jobs} = group
116  List.reverse(&2, WeightedJob, jobs)
117
118def plan_step(groups: Group & Group,
119              next: List<&2, WeightedJob> -> List<&2, WeightedJob> -> Plan) -> Plan:
120  (left, right) = groups
121  next(group_jobs(left), group_jobs(right))
122
123def plan(fuel: Nat, jobs: List<&2, WeightedJob>) -> Plan:
124  match fuel jobs:
125    case _ Nil{}:
126      Empty{}
127    case _ WeightedJob{index, _, job} <> Nil{}:
128      Leaf{index, job}
129    case 0n values:
130      Serial{values}
131    case Succ{+p} +values:
132      +count = List.length(&2, WeightedJob, values)
133      +half = Nat.div(count, 2n)
134      plan_step(partition(values,
135        (Group{half, 0.0, Nil{}}, Group{Nat.sub(count, half), 0.0, Nil{}})),
136        left => right => Fork{plan(p, left), plan(p, right)})
137
138def prepare(+jobs: List<&2, E.AlignJob>) -> Plan:
139  plan(List.length(&2, E.AlignJob, jobs),
140    List.sort(~WeightedJob, ~heavier, indexed(jobs, 0n)))
141
142def single(index: Nat, job: E.AlignJob) -> IndexedResult:
143  E.AlignJob{query, target, config} = job
144  IndexedResult{index, M.edlib_align(query, target, config)}
145
146def serial(jobs: List<&2, WeightedJob>) -> List<&2, IndexedResult>:
147  match jobs:
148    case Nil{}:
149      Nil{}
150    case WeightedJob{index, _, job} <> tail:
151      single(index, job) <> serial(tail)
152
153# Only ready, independent subplans enter parallel calls. A leaf uses the same
154# single-pair implementation; column carries and successive columns stay serial.
155def execute(plan: Plan) -> List<&2, IndexedResult>:
156  match plan:
157    case Empty{}:
158      Nil{}
159    case Leaf{index, job}:
160      [single(index, job)]
161    case Serial{jobs}:
162      serial(jobs)
163    case Fork{left, right}:
164      a b = execute(left) execute(right)
165      List.append(&2, IndexedResult, a, b)
166
167def earlier(first: IndexedResult, second: IndexedResult) -> Bool:
168  IndexedResult{a, _} = first
169  IndexedResult{b, _} = second
170  Nat.is_le(a, b)
171
172def results(values: List<&2, IndexedResult>) -> List<&2, E.AlignResult>:
173  match values:
174    case Nil{}:
175      Nil{}
176    case IndexedResult{_, result} <> tail:
177      result <> results(tail)
178
179def ordered(values: List<&2, IndexedResult>) -> List<&2, E.AlignResult>:
180  results(List.sort(~IndexedResult, ~earlier, values))
181
182def align_batch(jobs: List<&2, E.AlignJob>) -> List<&2, E.AlignResult>:
183  match jobs:
184    case Nil{}:
185      Nil{}
186    case E.AlignJob{query, target, config} <> Nil{}:
187      [M.edlib_align(query, target, config)]
188    case values:
189      ordered(execute(prepare(values)))