annoscript 0.1.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- annoscript/__init__.py +1 -0
- annoscript/_backend_profiles.py +17 -0
- annoscript/_client_common.py +287 -0
- annoscript/_data.py +37 -0
- annoscript/_dispatch_pool.py +47 -0
- annoscript/_env.py +19 -0
- annoscript/_examples.py +34 -0
- annoscript/_expressions.py +93 -0
- annoscript/_plan_viewer.html +953 -0
- annoscript/_remote_jobs.py +211 -0
- annoscript/_request_adapters.py +71 -0
- annoscript/_speculation.py +321 -0
- annoscript/_test_artifacts.py +33 -0
- annoscript/_test_evaluation_http.py +117 -0
- annoscript/_test_ir_runtime.py +138 -0
- annoscript/_test_job_process.py +166 -0
- annoscript/_test_plan_runtime.py +331 -0
- annoscript/_test_remote_client.py +46 -0
- annoscript/_test_remote_server.py +75 -0
- annoscript/_test_synthetic.py +67 -0
- annoscript/anno.py +270 -0
- annoscript/backends/__init__.py +13 -0
- annoscript/backends/_batch_wire.py +26 -0
- annoscript/backends/_bounded_grammar.py +97 -0
- annoscript/backends/_example_layout.py +23 -0
- annoscript/backends/_generation_policy.py +116 -0
- annoscript/backends/_grammar.py +83 -0
- annoscript/backends/_preview.py +13 -0
- annoscript/backends/_prompt.py +115 -0
- annoscript/backends/_shared_rules.py +38 -0
- annoscript/backends/adapters/__init__.py +17 -0
- annoscript/backends/adapters/chat.py +39 -0
- annoscript/backends/adapters/decision.py +119 -0
- annoscript/backends/base.py +427 -0
- annoscript/backends/config.py +67 -0
- annoscript/backends/decision.py +150 -0
- annoscript/backends/llm.py +338 -0
- annoscript/client.py +247 -0
- annoscript/client_tests.py +17 -0
- annoscript/compiler.py +708 -0
- annoscript/cross_version_tests.py +68 -0
- annoscript/fixtures/request-adapters-v1.json +4741 -0
- annoscript/formatter.py +171 -0
- annoscript/grammar_tests.py +176 -0
- annoscript/inspect.py +77 -0
- annoscript/ir.py +458 -0
- annoscript/job_cross_version_tests.py +118 -0
- annoscript/job_protocol.py +180 -0
- annoscript/jobs.py +255 -0
- annoscript/metrics.py +153 -0
- annoscript/p0_grammar_tests.py +147 -0
- annoscript/passes/__init__.py +40 -0
- annoscript/passes/_analysis.py +178 -0
- annoscript/passes/control_flow.py +68 -0
- annoscript/passes/dead_assignments.py +44 -0
- annoscript/passes/folding.py +104 -0
- annoscript/passes/propagation.py +53 -0
- annoscript/passes/temporaries.py +45 -0
- annoscript/plan.py +542 -0
- annoscript/plan_formatter.py +452 -0
- annoscript/plan_lowering.py +571 -0
- annoscript/preview.py +115 -0
- annoscript/remote.py +304 -0
- annoscript/runtime.py +861 -0
- annoscript/server/__init__.py +1 -0
- annoscript/server/__main__.py +49 -0
- annoscript/server/_plan_cache.py +79 -0
- annoscript/server/api.py +146 -0
- annoscript/server/metrics.py +140 -0
- annoscript/server/migrate.py +61 -0
- annoscript/server/releases.py +143 -0
- annoscript/server/store.py +429 -0
- annoscript/server/worker.py +290 -0
- annoscript/substring_tests.py +79 -0
- annoscript/test_advisory_lengths.py +115 -0
- annoscript/test_backends.py +181 -0
- annoscript/test_client.py +349 -0
- annoscript/test_compact_batch.py +95 -0
- annoscript/test_compiler.py +444 -0
- annoscript/test_decision_adapters.py +410 -0
- annoscript/test_decision_wire.py +190 -0
- annoscript/test_diagnostics.py +97 -0
- annoscript/test_engineering_p0.py +191 -0
- annoscript/test_environment.py +65 -0
- annoscript/test_example_backends.py +251 -0
- annoscript/test_examples.py +149 -0
- annoscript/test_field_names.py +221 -0
- annoscript/test_formatter.py +153 -0
- annoscript/test_fusion_scheduling.py +109 -0
- annoscript/test_generation_policy.py +180 -0
- annoscript/test_ir.py +148 -0
- annoscript/test_jobs.py +371 -0
- annoscript/test_jobs_server.py +511 -0
- annoscript/test_llm_batch.py +428 -0
- annoscript/test_metrics.py +59 -0
- annoscript/test_metrics_server.py +261 -0
- annoscript/test_minimal_fusion.py +59 -0
- annoscript/test_nullable_generate.py +169 -0
- annoscript/test_passes.py +408 -0
- annoscript/test_plan.py +522 -0
- annoscript/test_plan_constraints.py +305 -0
- annoscript/test_plan_formatter.py +272 -0
- annoscript/test_prompt_layout.py +140 -0
- annoscript/test_qwen_backend.py +181 -0
- annoscript/test_qwen_pilot_regression.py +180 -0
- annoscript/test_release_client.py +20 -0
- annoscript/test_remote.py +308 -0
- annoscript/test_request_adapters.py +295 -0
- annoscript/test_requirements.py +228 -0
- annoscript/test_rule_dedup.py +67 -0
- annoscript/test_runtime.py +314 -0
- annoscript/test_runtime_client.py +133 -0
- annoscript/test_source_fields.py +111 -0
- annoscript/test_speculation.py +458 -0
- annoscript/test_substring_backend.py +81 -0
- annoscript/test_synthetic_example.py +61 -0
- annoscript-0.1.0.dist-info/METADATA +145 -0
- annoscript-0.1.0.dist-info/RECORD +128 -0
- annoscript-0.1.0.dist-info/WHEEL +4 -0
- annoscript-0.1.0.dist-info/licenses/LICENSE +21 -0
- annoscript_examples/__init__.py +1 -0
- annoscript_examples/fixtures.py +3 -0
- annoscript_examples/request.execution_plan.json +3017 -0
- annoscript_examples/request.json +858 -0
- annoscript_examples/request.py +37 -0
- annoscript_examples/request_preview.decision.json +23 -0
- annoscript_examples/request_preview.generation.json +20 -0
- annoscript_examples/run.py +24 -0
|
@@ -0,0 +1,571 @@
|
|
|
1
|
+
"""Lower structured AnnoScript IR into demand-driven dependency templates.
|
|
2
|
+
|
|
3
|
+
No source inspection, model execution, or optimization of the input Program.
|
|
4
|
+
Alias analysis is deliberately allocation-based, not index-disjointness based:
|
|
5
|
+
unknown accesses retain may-dependences, including dictionary insertion order.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
from dataclasses import dataclass, fields, replace
|
|
11
|
+
import inspect
|
|
12
|
+
|
|
13
|
+
from annoscript import ir
|
|
14
|
+
from annoscript import plan as p
|
|
15
|
+
from annoscript.anno import Anno
|
|
16
|
+
from annoscript.passes._analysis import (block_writes, definitions, direct_reads,
|
|
17
|
+
expression_reads, literal_key, walk_statements)
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class PlanLoweringError(ValueError):
|
|
21
|
+
def __init__(self, message, statement=None):
|
|
22
|
+
self.source = statement.source if statement else None
|
|
23
|
+
suffix = ""
|
|
24
|
+
if statement:
|
|
25
|
+
s = statement.source.span
|
|
26
|
+
suffix = f" at {s.module}:{s.line}:{s.column} ({s.function})"
|
|
27
|
+
for call in statement.source.calls:
|
|
28
|
+
suffix += f" <- {call.function}:{call.line}:{call.column}"
|
|
29
|
+
super().__init__(message + suffix)
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def _names(target):
|
|
33
|
+
if isinstance(target, ir.NameTarget):
|
|
34
|
+
return [target.name]
|
|
35
|
+
return [n for item in target.items for n in _names(item)]
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def _entry_reads(body, bound=frozenset()):
|
|
39
|
+
"""Names read before a definite local definition, and definite definitions."""
|
|
40
|
+
bound, reads = set(bound), set()
|
|
41
|
+
for s in body:
|
|
42
|
+
reads |= set(direct_reads(s)) - bound
|
|
43
|
+
if isinstance(s, ir.If):
|
|
44
|
+
a, ad = _entry_reads(s.then_body, bound)
|
|
45
|
+
b, bd = _entry_reads(s.else_body, bound)
|
|
46
|
+
reads |= a | b
|
|
47
|
+
bound = ad & bd
|
|
48
|
+
elif isinstance(s, ir.For):
|
|
49
|
+
a, _ = _entry_reads(s.body, bound | definitions(s))
|
|
50
|
+
reads |= a
|
|
51
|
+
else:
|
|
52
|
+
bound |= definitions(s)
|
|
53
|
+
return reads, bound
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def _expr_children(expr):
|
|
57
|
+
for f in fields(expr):
|
|
58
|
+
def visit(v):
|
|
59
|
+
if isinstance(v, ir.Expr):
|
|
60
|
+
yield v
|
|
61
|
+
elif isinstance(v, tuple):
|
|
62
|
+
for x in v:
|
|
63
|
+
yield from visit(x)
|
|
64
|
+
yield from visit(getattr(expr, f.name))
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def _has_identity(program):
|
|
68
|
+
def visit(e):
|
|
69
|
+
if isinstance(e, ir.Compare) and any(op in ("is", "is_not") for op in e.ops):
|
|
70
|
+
# Singleton identity is a value test, not allocation observation.
|
|
71
|
+
operands = (e.left,) + e.comparators
|
|
72
|
+
if any(op in ("is", "is_not") and not any(isinstance(x, ir.Literal) and (x.value is None or type(x.value) is bool)
|
|
73
|
+
for x in operands[i:i+2])
|
|
74
|
+
for i, op in enumerate(e.ops)):
|
|
75
|
+
return True
|
|
76
|
+
return any(visit(c) for c in _expr_children(e))
|
|
77
|
+
return any(visit(v) for s in walk_statements(program.body) for f in fields(s)
|
|
78
|
+
for v in (getattr(s, f.name),) if isinstance(v, ir.Expr))
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
@dataclass
|
|
82
|
+
class _Flow:
|
|
83
|
+
env: dict
|
|
84
|
+
tail: str | None = None
|
|
85
|
+
barrier: str | None = None
|
|
86
|
+
gates: tuple = ()
|
|
87
|
+
falls: bool = True
|
|
88
|
+
returns: bool = False
|
|
89
|
+
|
|
90
|
+
def copy(self):
|
|
91
|
+
return replace(self, env=dict(self.env))
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
class _Lowerer:
|
|
95
|
+
def __init__(self, program):
|
|
96
|
+
self.program = program
|
|
97
|
+
self.tasks = {}
|
|
98
|
+
self.regions = {}
|
|
99
|
+
self.values = {}
|
|
100
|
+
self.edges = []
|
|
101
|
+
self.identity = _has_identity(program)
|
|
102
|
+
self.serial = 0
|
|
103
|
+
self.requests = []
|
|
104
|
+
|
|
105
|
+
def fresh(self, prefix):
|
|
106
|
+
self.serial += 1
|
|
107
|
+
return f"{prefix}.{self.serial}"
|
|
108
|
+
|
|
109
|
+
def region(self, parent, kind, owner):
|
|
110
|
+
rid = self.fresh("region")
|
|
111
|
+
self.regions[rid] = p.Region(rid, parent, kind, owner, ())
|
|
112
|
+
return rid
|
|
113
|
+
|
|
114
|
+
def path(self, task):
|
|
115
|
+
return p.loop_path(self.tasks[task].region, self.regions)
|
|
116
|
+
|
|
117
|
+
def edge(self, source, target, kind, *, carried=None, distance=0, **kwargs):
|
|
118
|
+
d = p.Dependency(source, target, kind,
|
|
119
|
+
p.InstanceMap(self.path(source), self.path(target), carried, distance), **kwargs)
|
|
120
|
+
if d not in self.edges:
|
|
121
|
+
self.edges.append(d)
|
|
122
|
+
|
|
123
|
+
def value_edge(self, value, target, **kwargs):
|
|
124
|
+
self.edge(self.values[value].producer, target, "data", value=value, **kwargs)
|
|
125
|
+
|
|
126
|
+
def bound(self, expr, flow, s):
|
|
127
|
+
names = sorted(expression_reads(expr))
|
|
128
|
+
for name in names:
|
|
129
|
+
if name not in flow.env:
|
|
130
|
+
raise PlanLoweringError(f"variable {name!r} is not definitely assigned", s)
|
|
131
|
+
return p.BoundExpr(expr, tuple(p.Binding(n, flow.env[n]) for n in names))
|
|
132
|
+
|
|
133
|
+
def add(self, cls, s, region, flow, *, names=(), tid=None, sequence=True, **kwargs):
|
|
134
|
+
tid = tid or self.fresh(s.id)
|
|
135
|
+
outputs = tuple(self.fresh("value") for _ in names)
|
|
136
|
+
task = cls(id=tid, region=region, origin=s.id, source=s.source,
|
|
137
|
+
outputs=outputs, gates=flow.gates, **kwargs)
|
|
138
|
+
self.tasks[tid] = task
|
|
139
|
+
r = self.regions[region]
|
|
140
|
+
self.regions[region] = replace(r, tasks=r.tasks + (tid,))
|
|
141
|
+
for name, value in zip(names, outputs):
|
|
142
|
+
self.values[value] = p.Value(value, tid, name)
|
|
143
|
+
for expr in p.expressions(task):
|
|
144
|
+
for binding in expr.bindings:
|
|
145
|
+
self.value_edge(binding.value, tid)
|
|
146
|
+
for gate in flow.gates:
|
|
147
|
+
self.edge(gate.task, tid, "control", port=gate.port)
|
|
148
|
+
if flow.barrier and flow.barrier != tid:
|
|
149
|
+
self.edge(flow.barrier, tid, "constraint", strict_only=True)
|
|
150
|
+
if sequence and flow.tail:
|
|
151
|
+
self.edge(flow.tail, tid, "failure_order")
|
|
152
|
+
if sequence:
|
|
153
|
+
flow.tail = tid
|
|
154
|
+
return task
|
|
155
|
+
|
|
156
|
+
def block(self, body, region, incoming):
|
|
157
|
+
flow = incoming.copy()
|
|
158
|
+
for s in body:
|
|
159
|
+
if not flow.falls:
|
|
160
|
+
break
|
|
161
|
+
if isinstance(s, ir.Assign):
|
|
162
|
+
t = self.add(p.Eval, s, region, flow, names=(s.target,), expression=self.bound(s.value, flow, s))
|
|
163
|
+
flow.env[s.target] = t.outputs[0]
|
|
164
|
+
elif isinstance(s, ir.Unpack):
|
|
165
|
+
names = _names(s.target)
|
|
166
|
+
t = self.add(p.Unpack, s, region, flow, names=names,
|
|
167
|
+
expression=self.bound(s.value, flow, s), target=s.target)
|
|
168
|
+
flow.env.update(zip(names, t.outputs))
|
|
169
|
+
elif isinstance(s, ir.SetItem):
|
|
170
|
+
self.add(p.Store, s, region, flow, container=self.bound(s.value, flow, s),
|
|
171
|
+
index=self.bound(s.index, flow, s), item=self.bound(s.item, flow, s))
|
|
172
|
+
elif isinstance(s, ir.Infer):
|
|
173
|
+
args = ir.TupleExpr((ir.TupleExpr(s.args), ir.TupleExpr(tuple(v for _, v in s.kwargs))))
|
|
174
|
+
a = self.add(p.Eval, s, region, flow, names=(None,), expression=self.bound(args, flow, s))
|
|
175
|
+
t = self.add(p.Infer, s, region, flow, names=(s.target,), request="pending",
|
|
176
|
+
arguments=a.outputs[0], keywords=tuple(k for k, _ in s.kwargs))
|
|
177
|
+
self.value_edge(a.outputs[0], t.id)
|
|
178
|
+
flow.env[s.target] = t.outputs[0]
|
|
179
|
+
elif isinstance(s, ir.Constrain):
|
|
180
|
+
t = self.add(p.Check, s, region, flow, condition=self.bound(s.condition, flow, s),
|
|
181
|
+
message=self.bound(s.message, flow, s) if s.message is not None else None)
|
|
182
|
+
# Strict checks wait for prior active work. Eager checks need
|
|
183
|
+
# only their actual data, control and memory prerequisites.
|
|
184
|
+
for previous in self.regions[region].tasks[:-1]:
|
|
185
|
+
self.edge(previous, t.id, "constraint", strict_only=True)
|
|
186
|
+
flow.barrier = t.id
|
|
187
|
+
elif isinstance(s, ir.Return):
|
|
188
|
+
if p.loop_path(region, self.regions):
|
|
189
|
+
raise PlanLoweringError("return inside a loop is unsupported", s)
|
|
190
|
+
self.add(p.Return, s, region, flow, expression=self.bound(s.value, flow, s))
|
|
191
|
+
flow.falls = False
|
|
192
|
+
flow.returns = True
|
|
193
|
+
elif isinstance(s, ir.If):
|
|
194
|
+
self.branch(s, region, flow)
|
|
195
|
+
elif isinstance(s, ir.For):
|
|
196
|
+
self.loop(s, region, flow)
|
|
197
|
+
else:
|
|
198
|
+
raise PlanLoweringError(f"unsupported statement {type(s).__name__}", s)
|
|
199
|
+
return flow
|
|
200
|
+
|
|
201
|
+
def branch(self, s, region, flow):
|
|
202
|
+
tid = self.fresh(s.id)
|
|
203
|
+
yes = self.region(region, "then", tid)
|
|
204
|
+
no = self.region(region, "else", tid)
|
|
205
|
+
t = self.add(p.Branch, s, region, flow, tid=tid, condition=self.bound(s.condition, flow, s),
|
|
206
|
+
then_region=yes, else_region=no)
|
|
207
|
+
children = []
|
|
208
|
+
for body, rid, port in ((s.then_body, yes, "true"), (s.else_body, no, "false")):
|
|
209
|
+
child = flow.copy()
|
|
210
|
+
child.returns = False
|
|
211
|
+
child.gates += (p.Gate(t.id, port),)
|
|
212
|
+
children.append(self.block(body, rid, child))
|
|
213
|
+
a, b = children
|
|
214
|
+
join = self.add(p.Join, s, region, flow, branch=t.id)
|
|
215
|
+
self.edge(t.id, join.id, "control")
|
|
216
|
+
# Join itself is a completion token: the chosen child must finish.
|
|
217
|
+
if a.barrier != flow.barrier or b.barrier != flow.barrier:
|
|
218
|
+
flow.barrier = join.id
|
|
219
|
+
if a.returns or b.returns:
|
|
220
|
+
flow.gates += (p.Gate(join.id, "continue"),)
|
|
221
|
+
flow.returns |= a.returns or b.returns
|
|
222
|
+
flow.falls = a.falls or b.falls
|
|
223
|
+
if not flow.falls:
|
|
224
|
+
return
|
|
225
|
+
available = (a.env.keys() & b.env.keys()) if a.falls and b.falls else (a.env.keys() if a.falls else b.env.keys())
|
|
226
|
+
merged = {}
|
|
227
|
+
for name in sorted(available):
|
|
228
|
+
av, bv = (a.env.get(name) if a.falls else None), (b.env.get(name) if b.falls else None)
|
|
229
|
+
if av == bv:
|
|
230
|
+
merged[name] = av
|
|
231
|
+
continue
|
|
232
|
+
m = self.add(p.Merge, s, region, flow, sequence=False, names=(name,), branch=t.id,
|
|
233
|
+
then_value=av, else_value=bv)
|
|
234
|
+
self.edge(t.id, m.id, "control")
|
|
235
|
+
for value, port in ((av, "true"), (bv, "false")):
|
|
236
|
+
if value is not None:
|
|
237
|
+
self.value_edge(value, m.id, condition=p.Gate(t.id, port))
|
|
238
|
+
merged[name] = m.outputs[0]
|
|
239
|
+
flow.env = merged
|
|
240
|
+
|
|
241
|
+
def loop(self, s, region, flow):
|
|
242
|
+
tid = self.fresh(s.id)
|
|
243
|
+
rid = self.region(region, "loop", tid)
|
|
244
|
+
t = self.add(p.Loop, s, region, flow, tid=tid, iterable=self.bound(s.iterable, flow, s),
|
|
245
|
+
body=rid, exports=())
|
|
246
|
+
old = dict(flow.env)
|
|
247
|
+
writes = block_writes(s.body) | definitions(s)
|
|
248
|
+
reads, definite = _entry_reads(s.body, definitions(s))
|
|
249
|
+
carry_names = sorted((writes & old.keys()) & (reads | (writes - definite)))
|
|
250
|
+
child = _Flow(dict(old), gates=flow.gates + (p.Gate(tid, "enter"),), barrier=flow.barrier)
|
|
251
|
+
carries = []
|
|
252
|
+
for name in carry_names:
|
|
253
|
+
c = self.add(p.Carry, s, rid, child, names=(name,), sequence=False,
|
|
254
|
+
loop=tid, initial=old[name], update="pending")
|
|
255
|
+
child.env[name] = c.outputs[0]
|
|
256
|
+
carries.append((name, c))
|
|
257
|
+
names = _names(s.target)
|
|
258
|
+
item = self.add(p.Iteration, s, rid, child, names=names, loop=tid, target=s.target)
|
|
259
|
+
child.env.update(zip(names, item.outputs))
|
|
260
|
+
child = self.block(s.body, rid, child)
|
|
261
|
+
exports = []
|
|
262
|
+
for name in sorted(writes & old.keys()):
|
|
263
|
+
value = self.fresh("value")
|
|
264
|
+
self.values[value] = p.Value(value, tid, name)
|
|
265
|
+
exports.append(p.LoopExport(value, old[name], child.env[name]))
|
|
266
|
+
flow.env[name] = value
|
|
267
|
+
self.value_edge(old[name], tid)
|
|
268
|
+
self.tasks[tid] = replace(t, outputs=tuple(e.output for e in exports), exports=tuple(exports))
|
|
269
|
+
for name, c in carries:
|
|
270
|
+
update = child.env[name]
|
|
271
|
+
self.tasks[c.id] = replace(c, update=update)
|
|
272
|
+
self.value_edge(update, c.id, carried=tid, distance=1, seed=c.initial)
|
|
273
|
+
if child.tail:
|
|
274
|
+
self.edge(child.tail, item.id, "failure_order", carried=tid, distance=1)
|
|
275
|
+
# Checks preserve source order in strict mode, but do not alone force
|
|
276
|
+
# eager iteration sequencing. Memory analysis may still require it.
|
|
277
|
+
has_check = any(isinstance(x, ir.Constrain) for x in walk_statements(s.body))
|
|
278
|
+
if has_check:
|
|
279
|
+
if child.tail:
|
|
280
|
+
self.edge(child.tail, item.id, "constraint", carried=tid, distance=1, strict_only=True)
|
|
281
|
+
flow.barrier = tid
|
|
282
|
+
|
|
283
|
+
def finish(self):
|
|
284
|
+
root = self.region(None, "root", None)
|
|
285
|
+
flow = self.block(self.program.body, root, _Flow({}))
|
|
286
|
+
if flow.falls:
|
|
287
|
+
raise PlanLoweringError("program has a reachable path without Return")
|
|
288
|
+
self.analyze_memory()
|
|
289
|
+
self.share_requests()
|
|
290
|
+
result = p.ExecutionPlan(self.program.entry, self.program.config, self.program.source,
|
|
291
|
+
tuple(self.regions.values()), tuple(self.tasks.values()),
|
|
292
|
+
tuple(self.values.values()), tuple(self.edges), tuple(self.requests))
|
|
293
|
+
result.validate()
|
|
294
|
+
return result
|
|
295
|
+
|
|
296
|
+
def analyze_memory(self):
|
|
297
|
+
"""Monotone reachability sets; allocation IDs retain iteration scopes."""
|
|
298
|
+
self.aliases = {v: set() for v in self.values}
|
|
299
|
+
self.points = {v: set() for v in self.values}
|
|
300
|
+
self.allocations = {}
|
|
301
|
+
self.mutable = set()
|
|
302
|
+
|
|
303
|
+
def reach(e, bindings, allocation):
|
|
304
|
+
if isinstance(e, ir.Var):
|
|
305
|
+
return set(self.aliases[bindings[e.name]])
|
|
306
|
+
if isinstance(e, ir.State):
|
|
307
|
+
return {"state"}
|
|
308
|
+
if isinstance(e, (ir.ListExpr, ir.DictExpr, ir.TupleExpr)):
|
|
309
|
+
result = {allocation}
|
|
310
|
+
for c in _expr_children(e):
|
|
311
|
+
result |= reach(c, bindings, allocation)
|
|
312
|
+
return result
|
|
313
|
+
if isinstance(e, (ir.GetItem, ir.Field)):
|
|
314
|
+
return reach(e.value, bindings, allocation)
|
|
315
|
+
if isinstance(e, ir.Binary):
|
|
316
|
+
return {allocation} | set().union(*(reach(c, bindings, allocation) for c in _expr_children(e)))
|
|
317
|
+
if isinstance(e, ir.Intrinsic) and e.name == "enumerate":
|
|
318
|
+
return {allocation} | set().union(*(reach(c, bindings, allocation) for c in e.args))
|
|
319
|
+
return set()
|
|
320
|
+
|
|
321
|
+
self.reach = reach
|
|
322
|
+
|
|
323
|
+
def points(e, bindings, allocation):
|
|
324
|
+
if isinstance(e, ir.Var):
|
|
325
|
+
return set(self.points[bindings[e.name]])
|
|
326
|
+
if isinstance(e, (ir.ListExpr, ir.DictExpr, ir.TupleExpr, ir.Binary)):
|
|
327
|
+
return {allocation}
|
|
328
|
+
return reach(e, bindings, allocation)
|
|
329
|
+
|
|
330
|
+
for t in self.tasks.values():
|
|
331
|
+
for v in t.outputs:
|
|
332
|
+
self.allocations[v] = self.path(t.id)
|
|
333
|
+
changed = True
|
|
334
|
+
while changed:
|
|
335
|
+
changed = False
|
|
336
|
+
for t in self.tasks.values():
|
|
337
|
+
results = []
|
|
338
|
+
direct = []
|
|
339
|
+
if isinstance(t, (p.Eval, p.Unpack)):
|
|
340
|
+
b = t.expression
|
|
341
|
+
results = [reach(b.expression, {x.name: x.value for x in b.bindings}, t.outputs[0])]
|
|
342
|
+
direct = [points(b.expression, {x.name: x.value for x in b.bindings}, t.outputs[0])]
|
|
343
|
+
elif isinstance(t, p.Infer):
|
|
344
|
+
results = [{f"response:{t.id}"}]
|
|
345
|
+
elif isinstance(t, p.Merge):
|
|
346
|
+
results = [set().union(*(self.aliases[v] for v in (t.then_value, t.else_value) if v))]
|
|
347
|
+
direct = [set().union(*(self.points[v] for v in (t.then_value, t.else_value) if v))]
|
|
348
|
+
elif isinstance(t, p.Carry):
|
|
349
|
+
results = [self.aliases[t.initial] | self.aliases[t.update]]
|
|
350
|
+
direct = [self.points[t.initial] | self.points[t.update]]
|
|
351
|
+
elif isinstance(t, p.Loop):
|
|
352
|
+
results = [self.aliases[e.initial] | self.aliases[e.update] for e in t.exports]
|
|
353
|
+
direct = [self.points[e.initial] | self.points[e.update] for e in t.exports]
|
|
354
|
+
elif isinstance(t, p.Iteration):
|
|
355
|
+
loop = self.tasks[t.loop]
|
|
356
|
+
results = [reach(loop.iterable.expression, {x.name: x.value for x in loop.iterable.bindings}, t.outputs[0])]
|
|
357
|
+
for i, v in enumerate(t.outputs):
|
|
358
|
+
new = results[min(i, len(results)-1)] if results else set()
|
|
359
|
+
if not new <= self.aliases[v]:
|
|
360
|
+
self.aliases[v] |= new
|
|
361
|
+
changed = True
|
|
362
|
+
new_direct = direct[min(i, len(direct)-1)] if direct else new
|
|
363
|
+
if not new_direct <= self.points[v]:
|
|
364
|
+
self.points[v] |= new_direct
|
|
365
|
+
changed = True
|
|
366
|
+
# A store changes containment, not the identity of its target.
|
|
367
|
+
# Repeat value analysis so subsequent subscript-derived aliases
|
|
368
|
+
# also see objects inserted after the original allocation.
|
|
369
|
+
for t in self.tasks.values():
|
|
370
|
+
if not isinstance(t, p.Store):
|
|
371
|
+
continue
|
|
372
|
+
b = t.container
|
|
373
|
+
targets = points(b.expression, {x.name: x.value for x in b.bindings}, t.id)
|
|
374
|
+
b = t.item
|
|
375
|
+
inserted = reach(b.expression, {x.name: x.value for x in b.bindings}, t.id)
|
|
376
|
+
for aliases in self.aliases.values():
|
|
377
|
+
if aliases & targets and not inserted <= aliases:
|
|
378
|
+
aliases |= inserted
|
|
379
|
+
changed = True
|
|
380
|
+
self.effects = {}
|
|
381
|
+
for t in self.tasks.values():
|
|
382
|
+
reads, writes = set(), set()
|
|
383
|
+
for b in p.expressions(t):
|
|
384
|
+
bindings = {x.name: x.value for x in b.bindings}
|
|
385
|
+
|
|
386
|
+
def accesses(e):
|
|
387
|
+
result = set()
|
|
388
|
+
if isinstance(e, (ir.GetItem, ir.Field, ir.Compare, ir.BoolOp, ir.Unary, ir.Binary, ir.Intrinsic, ir.FormatValue)):
|
|
389
|
+
for c in _expr_children(e):
|
|
390
|
+
result |= reach(c, bindings, t.outputs[0] if t.outputs else t.id)
|
|
391
|
+
for c in _expr_children(e):
|
|
392
|
+
result |= accesses(c)
|
|
393
|
+
return result
|
|
394
|
+
reads |= accesses(b.expression)
|
|
395
|
+
# Check truth-tests the result even when its expression is a
|
|
396
|
+
# bare container Var. Protect that read (and its message) from
|
|
397
|
+
# later writes now that eager checks are not global barriers.
|
|
398
|
+
if isinstance(t, (p.Return, p.Loop, p.Unpack, p.Check)):
|
|
399
|
+
reads |= reach(b.expression, bindings, t.id)
|
|
400
|
+
if isinstance(t, p.Infer):
|
|
401
|
+
reads |= self.aliases[t.arguments]
|
|
402
|
+
if isinstance(t, p.Store):
|
|
403
|
+
b = t.container
|
|
404
|
+
writes = points(b.expression, {x.name: x.value for x in b.bindings}, t.id)
|
|
405
|
+
if not writes or "state" in writes or any(a.startswith("response:") for a in writes):
|
|
406
|
+
raise PlanLoweringError("write target is not a provably local container", t)
|
|
407
|
+
self.mutable |= writes
|
|
408
|
+
# enumerate and iteration consume an object, even without SetItem.
|
|
409
|
+
if isinstance(t, (p.Loop, p.Unpack)):
|
|
410
|
+
writes |= {a for a in reads if self.is_iterator(a)}
|
|
411
|
+
self.effects[t.id] = (reads, writes)
|
|
412
|
+
self.mutable |= set().union(*(writes for _, writes in self.effects.values()))
|
|
413
|
+
self.carried_allocations = {}
|
|
414
|
+
for t in self.tasks.values():
|
|
415
|
+
if isinstance(t, p.Carry):
|
|
416
|
+
self.carried_allocations.setdefault(t.loop, set()).update(self.aliases[t.outputs[0]])
|
|
417
|
+
tasks = list(self.tasks.values())
|
|
418
|
+
for i, a in enumerate(tasks):
|
|
419
|
+
for b in tasks[i+1:]:
|
|
420
|
+
if self.exclusive(a.region, b.region):
|
|
421
|
+
continue
|
|
422
|
+
self.conflicts(a, b)
|
|
423
|
+
# Loop-carried hazards exist only for storage shared between iterations.
|
|
424
|
+
for a in tasks:
|
|
425
|
+
for b in tasks:
|
|
426
|
+
common = []
|
|
427
|
+
for x, y in zip(self.path(a.id), self.path(b.id)):
|
|
428
|
+
if x != y:
|
|
429
|
+
break
|
|
430
|
+
common.append(x)
|
|
431
|
+
for loop in common:
|
|
432
|
+
self.conflicts(a, b, loop)
|
|
433
|
+
for t in tuple(self.tasks.values()):
|
|
434
|
+
if isinstance(t, p.Loop):
|
|
435
|
+
read = self.effects[t.id][0]
|
|
436
|
+
if any(read & writes for task, (_, writes) in self.effects.items()
|
|
437
|
+
if t.id in self.path(task)):
|
|
438
|
+
self.tasks[t.id] = replace(t, serial_modes=("eager", "strict"))
|
|
439
|
+
|
|
440
|
+
def is_iterator(self, allocation):
|
|
441
|
+
value = self.values.get(allocation)
|
|
442
|
+
t = self.tasks.get(value.producer) if value else None
|
|
443
|
+
return isinstance(t, p.Eval) and isinstance(t.expression.expression, ir.Intrinsic) and t.expression.expression.name == "enumerate"
|
|
444
|
+
|
|
445
|
+
def exclusive(self, left, right):
|
|
446
|
+
def choices(region):
|
|
447
|
+
result = {}
|
|
448
|
+
while region:
|
|
449
|
+
r = self.regions[region]
|
|
450
|
+
if r.kind in ("then", "else"):
|
|
451
|
+
result[r.owner] = r.kind
|
|
452
|
+
region = r.parent
|
|
453
|
+
return result
|
|
454
|
+
a, b = choices(left), choices(right)
|
|
455
|
+
return any(k in b and b[k] != v for k, v in a.items())
|
|
456
|
+
|
|
457
|
+
def conflicts(self, a, b, loop=None):
|
|
458
|
+
# Loop iteration effects are represented by serial_modes, not by a
|
|
459
|
+
# self-deadlocking edge from a loop's completion into its own body.
|
|
460
|
+
if (isinstance(a, p.Loop) and a.id in self.path(b.id)) or (isinstance(b, p.Loop) and b.id in self.path(a.id)):
|
|
461
|
+
return
|
|
462
|
+
ar, aw = self.effects[a.id]
|
|
463
|
+
br, bw = self.effects[b.id]
|
|
464
|
+
for hazard, overlap in (("RAW", aw & br), ("WAR", ar & bw), ("WAW", aw & bw)):
|
|
465
|
+
if loop:
|
|
466
|
+
overlap = {x for x in overlap if loop not in self.allocations.get(x, ())
|
|
467
|
+
or x in self.carried_allocations.get(loop, ())}
|
|
468
|
+
if overlap:
|
|
469
|
+
self.edge(a.id, b.id, "memory", carried=loop, distance=1 if loop else 0,
|
|
470
|
+
hazard=hazard, certainty="conservative")
|
|
471
|
+
|
|
472
|
+
def share_requests(self):
|
|
473
|
+
"""Value-number requests without executing expressions or user code."""
|
|
474
|
+
memo = {}
|
|
475
|
+
|
|
476
|
+
def value_key(v, stack=frozenset()):
|
|
477
|
+
if v in memo:
|
|
478
|
+
return memo[v]
|
|
479
|
+
if v in stack or self.aliases[v] & self.mutable or any(self.is_iterator(a) for a in self.aliases[v]):
|
|
480
|
+
return ("version", v)
|
|
481
|
+
t = self.tasks[self.values[v].producer]
|
|
482
|
+
if self.effects[t.id][0] & self.mutable:
|
|
483
|
+
result = ("version", v)
|
|
484
|
+
elif isinstance(t, p.Eval):
|
|
485
|
+
result = expr_key(t.expression.expression, {b.name: b.value for b in t.expression.bindings}, stack | {v})
|
|
486
|
+
else:
|
|
487
|
+
result = ("version", v)
|
|
488
|
+
memo[v] = result
|
|
489
|
+
return result
|
|
490
|
+
|
|
491
|
+
def expr_key(e, bindings, stack=frozenset()):
|
|
492
|
+
if isinstance(e, ir.Var):
|
|
493
|
+
return value_key(bindings[e.name], stack)
|
|
494
|
+
if isinstance(e, ir.Literal):
|
|
495
|
+
ty, val = literal_key(e.value)
|
|
496
|
+
return ("literal", ty.__name__, val)
|
|
497
|
+
if isinstance(e, ir.State):
|
|
498
|
+
return ("state",)
|
|
499
|
+
def key(x):
|
|
500
|
+
if isinstance(x, ir.Expr):
|
|
501
|
+
return expr_key(x, bindings, stack)
|
|
502
|
+
if isinstance(x, tuple):
|
|
503
|
+
return tuple(key(v) for v in x)
|
|
504
|
+
return x
|
|
505
|
+
return (e.kind,) + tuple(key(getattr(e, f.name)) for f in fields(e))
|
|
506
|
+
|
|
507
|
+
def scope_for(key):
|
|
508
|
+
paths = []
|
|
509
|
+
def walk(x):
|
|
510
|
+
if isinstance(x, tuple):
|
|
511
|
+
if len(x) == 2 and x[0] == "version":
|
|
512
|
+
paths.append(self.path(self.values[x[1]].producer))
|
|
513
|
+
else:
|
|
514
|
+
for v in x:
|
|
515
|
+
walk(v)
|
|
516
|
+
walk(key)
|
|
517
|
+
return max(paths, key=len, default=())
|
|
518
|
+
|
|
519
|
+
groups = {}
|
|
520
|
+
for t in tuple(self.tasks.values()):
|
|
521
|
+
if not isinstance(t, p.Infer):
|
|
522
|
+
continue
|
|
523
|
+
a = self.tasks[self.values[t.arguments].producer]
|
|
524
|
+
expr = a.expression.expression
|
|
525
|
+
bindings = {b.name: b.value for b in a.expression.bindings}
|
|
526
|
+
args = tuple(expr_key(x, bindings) for x in expr.items[0].items)
|
|
527
|
+
kwargs = {n: expr_key(x, bindings) for n, x in zip(t.keywords, expr.items[1].items)}
|
|
528
|
+
# Signature binding is static metadata inspection, never a call.
|
|
529
|
+
original = next(s for s in walk_statements(self.program.body) if s.id == t.origin)
|
|
530
|
+
operation = original.operation
|
|
531
|
+
signature = inspect.signature(getattr(Anno, operation))
|
|
532
|
+
try:
|
|
533
|
+
bound = signature.bind(None, *args, **kwargs)
|
|
534
|
+
except TypeError as exc:
|
|
535
|
+
raise PlanLoweringError(f"invalid {operation} arguments: {exc}", original) from exc
|
|
536
|
+
canonical = []
|
|
537
|
+
for name, parameter in signature.parameters.items():
|
|
538
|
+
if name == "self":
|
|
539
|
+
continue
|
|
540
|
+
if name in bound.arguments:
|
|
541
|
+
canonical.append((name, bound.arguments[name]))
|
|
542
|
+
else:
|
|
543
|
+
canonical.append((name, expr_key(ir.Literal(parameter.default), {})))
|
|
544
|
+
key = (operation, tuple(canonical))
|
|
545
|
+
# A mutable argument's static binding can refer to different
|
|
546
|
+
# contents at each demand; conservatively keep these sites apart.
|
|
547
|
+
unsafe = self.identity or bool(self.aliases[t.arguments] & self.mutable)
|
|
548
|
+
unsafe |= any(self.is_iterator(x) for x in self.aliases[t.arguments])
|
|
549
|
+
scope = self.path(t.id) if unsafe else scope_for(key)
|
|
550
|
+
if unsafe:
|
|
551
|
+
key = (key, t.id)
|
|
552
|
+
group_key = (key, scope)
|
|
553
|
+
if group_key not in groups:
|
|
554
|
+
groups[group_key] = [self.fresh("request"), operation, scope, [], not unsafe]
|
|
555
|
+
group = groups[group_key]
|
|
556
|
+
group[3].append(t.id)
|
|
557
|
+
self.tasks[t.id] = replace(t, request=group[0])
|
|
558
|
+
self.requests = [p.InferRequest(qid, operation, scope, tuple(sites), reusable)
|
|
559
|
+
for qid, operation, scope, sites, reusable in groups.values()]
|
|
560
|
+
|
|
561
|
+
|
|
562
|
+
def lower_to_plan(program: ir.Program) -> p.ExecutionPlan:
|
|
563
|
+
"""Build a validated plan for either failure mode without running Python."""
|
|
564
|
+
if not isinstance(program, ir.Program):
|
|
565
|
+
raise TypeError("lower_to_plan expects an ir.Program")
|
|
566
|
+
try:
|
|
567
|
+
# Also validate constructed Programs, not only deserialized ones.
|
|
568
|
+
checked = ir.Program.from_dict(program.to_dict())
|
|
569
|
+
return _Lowerer(checked).finish()
|
|
570
|
+
except ir.IRDecodeError as exc:
|
|
571
|
+
raise PlanLoweringError(f"invalid AnnoScript IR: {exc}") from exc
|