Skip to content

Folder

folder

Folder dataclass

Source code in xdsl/folder.py
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
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
79
80
81
82
83
84
85
86
87
88
89
90
91
@dataclass
class Folder:
    context: Context

    def try_fold(self, op: Operation) -> tuple[list[SSAValue], list[Operation]] | None:
        """
        Try to fold the given operation.
        Returns a tuple with the list of SSAValues that replace the results of the operation,
        and a list of constant operations that were created from the constant attributes generated by the folder.
        If the operation could not be folded, returns None.

        Note that while this folds only one operation, multiple new operations can be created.
        Each of the results of the original operation might be replaced by a new constant operation.
        """

        if (trait := op.get_trait(HasFolder)) is None:
            return None
        folded = trait.fold(op)
        if folded is None:
            return None
        results: list[SSAValue] = []
        new_ops: list[Operation] = []
        for val, original_result in zip(folded, op.results, strict=True):
            if isinstance(val, SSAValue):
                results.append(val)
            else:
                assert isinstance(val, Attribute)
                dialect = self.context.get_dialect(op.dialect_name())
                interface = dialect.get_interface(ConstantMaterializationInterface)
                if interface is None:
                    return None
                assert isinstance(type := original_result.type, TypeAttribute)
                new_op = interface.materialize_constant(val, type)
                if new_op is None:
                    return None
                new_ops.append(new_op)
                results.append(new_op.results[0])
        return results, new_ops

    def insert_with_fold(
        self, op: Operation, builder: Builder
    ) -> Sequence[SSAValue] | None:
        """
        Inserts the operation using the provided builder, trying to fold it first.
        If folding is successful, the folded results are returned, otherwise None is returned.
        """
        results = self.try_fold(op)
        if results is None:
            builder.insert(op)
            return op.results
        else:
            if op.parent:
                raise ValueError(
                    "Can't insert_with_fold fold an operation that already has a parent."
                )
            values, new_ops = results
            builder.insert_op(new_ops)
            return values

    def replace_with_fold(
        self, op: Operation, rewriter: PatternRewriter, safe_erase: bool = True
    ) -> Sequence[SSAValue] | None:
        """
        Replaces the operation with its folded results.
        If folding is successful, the folded results are returned.
        Otherwise, returns None.
        """
        results = self.try_fold(op)
        if results is None:
            return None
        values, new_ops = results
        rewriter.replace_op(op, new_ops, values, safe_erase)
        return values

context: Context instance-attribute

__init__(context: Context) -> None

try_fold(op: Operation) -> tuple[list[SSAValue], list[Operation]] | None

Try to fold the given operation. Returns a tuple with the list of SSAValues that replace the results of the operation, and a list of constant operations that were created from the constant attributes generated by the folder. If the operation could not be folded, returns None.

Note that while this folds only one operation, multiple new operations can be created. Each of the results of the original operation might be replaced by a new constant operation.

Source code in xdsl/folder.py
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
def try_fold(self, op: Operation) -> tuple[list[SSAValue], list[Operation]] | None:
    """
    Try to fold the given operation.
    Returns a tuple with the list of SSAValues that replace the results of the operation,
    and a list of constant operations that were created from the constant attributes generated by the folder.
    If the operation could not be folded, returns None.

    Note that while this folds only one operation, multiple new operations can be created.
    Each of the results of the original operation might be replaced by a new constant operation.
    """

    if (trait := op.get_trait(HasFolder)) is None:
        return None
    folded = trait.fold(op)
    if folded is None:
        return None
    results: list[SSAValue] = []
    new_ops: list[Operation] = []
    for val, original_result in zip(folded, op.results, strict=True):
        if isinstance(val, SSAValue):
            results.append(val)
        else:
            assert isinstance(val, Attribute)
            dialect = self.context.get_dialect(op.dialect_name())
            interface = dialect.get_interface(ConstantMaterializationInterface)
            if interface is None:
                return None
            assert isinstance(type := original_result.type, TypeAttribute)
            new_op = interface.materialize_constant(val, type)
            if new_op is None:
                return None
            new_ops.append(new_op)
            results.append(new_op.results[0])
    return results, new_ops

insert_with_fold(op: Operation, builder: Builder) -> Sequence[SSAValue] | None

Inserts the operation using the provided builder, trying to fold it first. If folding is successful, the folded results are returned, otherwise None is returned.

Source code in xdsl/folder.py
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def insert_with_fold(
    self, op: Operation, builder: Builder
) -> Sequence[SSAValue] | None:
    """
    Inserts the operation using the provided builder, trying to fold it first.
    If folding is successful, the folded results are returned, otherwise None is returned.
    """
    results = self.try_fold(op)
    if results is None:
        builder.insert(op)
        return op.results
    else:
        if op.parent:
            raise ValueError(
                "Can't insert_with_fold fold an operation that already has a parent."
            )
        values, new_ops = results
        builder.insert_op(new_ops)
        return values

replace_with_fold(op: Operation, rewriter: PatternRewriter, safe_erase: bool = True) -> Sequence[SSAValue] | None

Replaces the operation with its folded results. If folding is successful, the folded results are returned. Otherwise, returns None.

Source code in xdsl/folder.py
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def replace_with_fold(
    self, op: Operation, rewriter: PatternRewriter, safe_erase: bool = True
) -> Sequence[SSAValue] | None:
    """
    Replaces the operation with its folded results.
    If folding is successful, the folded results are returned.
    Otherwise, returns None.
    """
    results = self.try_fold(op)
    if results is None:
        return None
    values, new_ops = results
    rewriter.replace_op(op, new_ops, values, safe_erase)
    return values