228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305 | class CostFormatter(LogicFormatter):
"""
CostFormatter chooses the format of each alias from the level options with
the lowest cost under cost_model, memoizing the best format for each depth and
input position count.
"""
def __init__(
self,
loader: LogicLoader | None = None,
cost_model: CostModel[LevelChoice] = STORAGE_COST_MODEL,
options: Sequence[LevelOption] = LEVEL_OPTIONS,
):
super().__init__()
if loader is None:
loader = MockLogicLoader()
self.ctx = loader
self.cost_model = cost_model
self.options = options
def get_tensor_ftype(
self,
alias: lgc.Alias,
term: lgc.LogicStatement,
machine: StatsMachine,
fill_type: AbstractFillFType,
shape_type: tuple[FType, ...],
) -> FiberTensorFType:
lvl = optimize_format(
alias, shape_type, machine, fill_type, self.cost_model, term, self.options
)
return fiber_tensor(lvl)
def lower(
self,
prgm: lgc.LogicStatement,
bindings: dict[lgc.Alias, TensorFType],
stats: dict[lgc.Alias, TS],
stats_factory: StatsFactory,
):
bindings = bindings.copy()
machine = StatsMachine(stats_factory, OrderedDict(stats))
shape_types = prgm.infer_shape_type(
{var: val.shape_type for var, val in bindings.items()}
)
fill_types = prgm.infer_fill_type(
{var: val.fill_type for var, val in bindings.items()}
)
def formatter(node: lgc.LogicStatement):
match node:
case lgc.Plan(bodies):
for body in bodies:
formatter(body)
case lgc.Query(lgc.Table(lgc.Alias() as lhs, _), _) | lgc.QueryInto(
lgc.Table(lgc.Alias() as lhs, _), _, _
):
machine(node)
if lhs not in bindings:
shape_type = tuple(
ftype(dim) if dim is not None else ftypes.intp
for dim in shape_types[lhs]
)
bindings[lhs] = self.get_tensor_ftype(
lhs, prgm, machine, fill_types[lhs], shape_type
)
case lgc.Produces():
return
case _:
raise ValueError(
f"Unsupported logic statement for formatting: {node}"
)
formatter(prgm)
logger.debug(prgm)
return self.ctx(prgm, bindings, machine.bindings, stats_factory)
|