src/Batch.bend
bend-edit-distance@0.2.1.0 · 189 lines · 0 definitions
1import Base
2import bend-edit-distance@0.2.1.0/EdlibTypes.bend as E
3import bend-edit-distance@0.2.1.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)))