Skip to content

Eqsat serialize egraph

eqsat_serialize_egraph

SerializeEGraph dataclass

Bases: ModulePass

Source code in xdsl/transforms/eqsat_serialize_egraph.py 
73
74
75
76
77
class SerializeEGraph(ModulePass):
    name = "eqsat-serialize-egraph"

    def apply(self, ctx: Context, op: builtin.ModuleOp) -> None:
        print(json.dumps({"nodes": serialize_to_egraph(op)}))

name = 'eqsat-serialize-egraph' class-attribute instance-attribute

apply(ctx: Context, op: builtin.ModuleOp) -> None

Source code in xdsl/transforms/eqsat_serialize_egraph.py 
76
77
def apply(self, ctx: Context, op: builtin.ModuleOp) -> None:
    print(json.dumps({"nodes": serialize_to_egraph(op)}))

serialize_to_egraph(mod: builtin.ModuleOp)

Source code in xdsl/transforms/eqsat_serialize_egraph.py 
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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
def serialize_to_egraph(mod: builtin.ModuleOp):
    enode_to_id: defaultdict[Operation | BlockArgument, str] = defaultdict(
        _IDGenerator("enode_")
    )
    eclass_to_id: defaultdict[equivalence.ClassOp, str] = defaultdict(
        _IDGenerator("eclass_")
    )
    nodes: dict[str, dict[str, str | list[str]]] = dict()
    for op in mod.walk(reverse=True):
        if isinstance(op, equivalence.ClassOp):
            for operand in op.operands:
                if isinstance(operand, BlockArgument):
                    nodes[enode_to_id[operand]] = {
                        "op": f"arg {operand.index}",
                        "eclass": eclass_to_id[op],
                        "children": [],
                    }
            continue
        children: list[Any] = []
        eclass_id = None
        for res in op.results:
            for use in res.uses:
                if isinstance(use.operation, equivalence.ClassOp):
                    assert len(op.results) == 1, (
                        "Only single result operations are supported"
                    )
                    assert res.has_one_use(), "Only single use operations are supported"
                    eclass_id = eclass_to_id[use.operation]
        if eclass_id is None:
            continue
        for operand in op.operands:
            if isinstance(operand.owner, equivalence.ClassOp):
                firstoperand = operand.owner.operands[0]
                if isinstance(firstoperand.owner, Operation):
                    canonical_enode = firstoperand.owner
                else:
                    canonical_enode = cast(BlockArgument, firstoperand)
                children.append(enode_to_id[canonical_enode])
        if op.operands:
            name = op.name
        else:
            # If the operation has no operands, we get the full string representation as name for the node.
            # This is useful for operations such as `arith.constant 42`.
            name = str(op).split("=")[1].split(":")[0].strip()
        nodes[enode_to_id[op]] = {
            "op": name,
            "eclass": eclass_id,
            "children": children,
        }
    return nodes