Skip to content

Equivalence

equivalence

An embedding of equivalence classes in IR, for use in equality saturation with non-destructive rewrites.

See the overview notebook.

EQSAT_COST_LABEL = 'eqsat_cost' module-attribute

Key used to store the cost of computing the result of an operation.

AnyClassOp = ClassOp | ConstantClassOp module-attribute

A type representing either a regular e-class operation or a constant e-class operation.

Equivalence = Dialect('equivalence', [ClassOp, ConstantClassOp, YieldOp, GraphOp]) module-attribute

ConstantClassOp

Bases: IRDLOperation, ConstantLikeInterface

An e-class representing a known constant value. For non-constant e-classes, use ClassOp.

Source code in xdsl/dialects/equivalence.py
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
71
72
73
74
75
76
77
78
@irdl_op_definition
class ConstantClassOp(IRDLOperation, ConstantLikeInterface):
    """An e-class representing a known constant value.
    For non-constant e-classes, use [ClassOp][xdsl.dialects.equivalence.ClassOp].
    """

    T: ClassVar = VarConstraint("T", AnyAttr())

    name = "equivalence.const_class"

    assembly_format = (
        "$arguments ` ` `(` `constant` `=` $value `)` attr-dict `:` type($result)"
    )
    traits = traits_def(Pure())

    arguments = var_operand_def(T)
    result = result_def(T)
    value = prop_def()
    min_cost_index = opt_attr_def(IntAttr)

    def get_constant_value(self):
        return self.value

    def __init__(self, const_arg: OpResult):
        if (trait := const_arg.owner.get_trait(ConstantLike)) is None:
            raise DiagnosticException(
                "The argument of a ConstantClass must be a constant-like operation."
            )
        value = trait.get_constant_value(const_arg.owner)
        super().__init__(
            operands=[const_arg],
            result_types=[const_arg.type],
            properties={"value": value},
        )

T: ClassVar = VarConstraint('T', AnyAttr()) class-attribute instance-attribute

name = 'equivalence.const_class' class-attribute instance-attribute

assembly_format = '$arguments ` ` `(` `constant` `=` $value `)` attr-dict `:` type($result)' class-attribute instance-attribute

traits = traits_def(Pure()) class-attribute instance-attribute

arguments = var_operand_def(T) class-attribute instance-attribute

result = result_def(T) class-attribute instance-attribute

value = prop_def() class-attribute instance-attribute

min_cost_index = opt_attr_def(IntAttr) class-attribute instance-attribute

get_constant_value()

Source code in xdsl/dialects/equivalence.py
65
66
def get_constant_value(self):
    return self.value

__init__(const_arg: OpResult)

Source code in xdsl/dialects/equivalence.py
68
69
70
71
72
73
74
75
76
77
78
def __init__(self, const_arg: OpResult):
    if (trait := const_arg.owner.get_trait(ConstantLike)) is None:
        raise DiagnosticException(
            "The argument of a ConstantClass must be a constant-like operation."
        )
    value = trait.get_constant_value(const_arg.owner)
    super().__init__(
        operands=[const_arg],
        result_types=[const_arg.type],
        properties={"value": value},
    )

ClassOp

Bases: IRDLOperation

An e-class representing a set of equivalent values. E-classes that represent a constant value can instead be represented by ConstantClassOp.

Source code in xdsl/dialects/equivalence.py
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
@irdl_op_definition
class ClassOp(IRDLOperation):
    """An e-class representing a set of equivalent values.
    E-classes that represent a constant value can instead
    be represented by [ConstantClassOp][xdsl.dialects.equivalence.ConstantClassOp].
    """

    T: ClassVar = VarConstraint("T", AnyAttr())

    name = "equivalence.class"
    arguments = var_operand_def(T)
    result = result_def(T)
    min_cost_index = opt_attr_def(IntAttr)
    traits = traits_def(Pure())

    assembly_format = "$arguments attr-dict `:` type($result)"

    def __init__(
        self,
        *arguments: SSAValue,
        min_cost_index: IntAttr | None = None,
        res_type: Attribute | None = None,
    ):
        if not arguments:
            raise DiagnosticException("E-class op must have at least one operand")
        if res_type is None:
            res_type = arguments[0].type

        super().__init__(
            operands=[arguments],
            result_types=[res_type],
            attributes={"min_cost_index": min_cost_index},
        )

    def verify_(self) -> None:
        if not self.operands:
            raise VerifyException("E-class operations must have at least one operand.")

        for operand in self.operands:
            if isinstance(operand.owner, ClassOp):
                # The two ops should have been merged into one.
                raise VerifyException(
                    "A result of an e-class operation cannot be used as an operand of "
                    "another e-class."
                )

            if not operand.has_one_use():
                if len(set(use.operation for use in operand.uses)) == 1:
                    raise VerifyException(
                        "E-class operands must only be used once by the e-class."
                    )
                else:
                    raise VerifyException(
                        "E-class operands must only be used by the e-class."
                    )

T: ClassVar = VarConstraint('T', AnyAttr()) class-attribute instance-attribute

name = 'equivalence.class' class-attribute instance-attribute

arguments = var_operand_def(T) class-attribute instance-attribute

result = result_def(T) class-attribute instance-attribute

min_cost_index = opt_attr_def(IntAttr) class-attribute instance-attribute

traits = traits_def(Pure()) class-attribute instance-attribute

assembly_format = '$arguments attr-dict `:` type($result)' class-attribute instance-attribute

__init__(*arguments: SSAValue, min_cost_index: IntAttr | None = None, res_type: Attribute | None = None)

Source code in xdsl/dialects/equivalence.py
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def __init__(
    self,
    *arguments: SSAValue,
    min_cost_index: IntAttr | None = None,
    res_type: Attribute | None = None,
):
    if not arguments:
        raise DiagnosticException("E-class op must have at least one operand")
    if res_type is None:
        res_type = arguments[0].type

    super().__init__(
        operands=[arguments],
        result_types=[res_type],
        attributes={"min_cost_index": min_cost_index},
    )

verify_() -> None

Source code in xdsl/dialects/equivalence.py
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
def verify_(self) -> None:
    if not self.operands:
        raise VerifyException("E-class operations must have at least one operand.")

    for operand in self.operands:
        if isinstance(operand.owner, ClassOp):
            # The two ops should have been merged into one.
            raise VerifyException(
                "A result of an e-class operation cannot be used as an operand of "
                "another e-class."
            )

        if not operand.has_one_use():
            if len(set(use.operation for use in operand.uses)) == 1:
                raise VerifyException(
                    "E-class operands must only be used once by the e-class."
                )
            else:
                raise VerifyException(
                    "E-class operands must only be used by the e-class."
                )

GraphOp

Bases: IRDLOperation

Source code in xdsl/dialects/equivalence.py
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
@irdl_op_definition
class GraphOp(IRDLOperation):
    name = "equivalence.graph"

    outputs = var_result_def()
    body = region_def()

    traits = lazy_traits_def(lambda: (SingleBlockImplicitTerminator(YieldOp),))

    assembly_format = "`->` type($outputs) $body attr-dict"

    def __init__(
        self,
        result_types: Sequence[Attribute] | None,
        body: Region,
    ):
        super().__init__(
            result_types=(result_types,),
            regions=[body],
        )

name = 'equivalence.graph' class-attribute instance-attribute

outputs = var_result_def() class-attribute instance-attribute

body = region_def() class-attribute instance-attribute

traits = lazy_traits_def(lambda: (SingleBlockImplicitTerminator(YieldOp),)) class-attribute instance-attribute

assembly_format = '`->` type($outputs) $body attr-dict' class-attribute instance-attribute

__init__(result_types: Sequence[Attribute] | None, body: Region)

Source code in xdsl/dialects/equivalence.py
156
157
158
159
160
161
162
163
164
def __init__(
    self,
    result_types: Sequence[Attribute] | None,
    body: Region,
):
    super().__init__(
        result_types=(result_types,),
        regions=[body],
    )

YieldOp

Bases: IRDLOperation

Source code in xdsl/dialects/equivalence.py
167
168
169
170
171
172
173
174
175
176
177
178
179
180
@irdl_op_definition
class YieldOp(IRDLOperation):
    name = "equivalence.yield"
    values = var_operand_def()

    traits = traits_def(HasParent(GraphOp), IsTerminator())

    assembly_format = "$values `:` type($values) attr-dict"

    def __init__(
        self,
        *values: SSAValue,
    ):
        super().__init__(operands=[values])

name = 'equivalence.yield' class-attribute instance-attribute

values = var_operand_def() class-attribute instance-attribute

traits = traits_def(HasParent(GraphOp), IsTerminator()) class-attribute instance-attribute

assembly_format = '$values `:` type($values) attr-dict' class-attribute instance-attribute

__init__(*values: SSAValue)

Source code in xdsl/dialects/equivalence.py
176
177
178
179
180
def __init__(
    self,
    *values: SSAValue,
):
    super().__init__(operands=[values])