Extracts tensors from logic plan, replacing them with aliases.
Source code in src/finch/autoschedule/executor.py
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47 | def extract_tensors(
root: lgc.LogicStatement,
bindings: dict[lgc.Alias, Tensor],
) -> tuple[lgc.LogicStatement, dict[lgc.Alias, Tensor]]:
"""
Extracts tensors from logic plan, replacing them with aliases.
"""
bindings = bindings.copy()
# ids is a dictionary that has key value as memory_address : Alias
ids: dict[int, lgc.Alias] = {id(val): key for key, val in bindings.items()}
spc = Namespace(root)
for alias in bindings:
# Reserving the Alias names that already exist
spc.freshen(alias.name)
def rule_0(node):
match node:
# Case where we have table with actual tensor
case lgc.Table(lgc.Literal(tns), idxs):
if id(tns) in ids:
var = ids[id(tns)]
return lgc.Table(var, idxs)
# If we don't have an Alias for the tensor we just found we create one
var = lgc.Alias(spc.freshen("A"))
# Updating the ids and bindings
ids[id(tns)] = var
bindings[var] = tns
return lgc.Table(var, idxs)
root = Rewrite(PostWalk(rule_0))(root)
return root, bindings
|