Skip to content

Bufferization

bufferization

Bufferization = Dialect('bufferization', [AllocTensorOp, CloneOp, ToTensorOp, ToBufferOp, MaterializeInDestinationOp], []) module-attribute

TensorFromMemRefConstraint dataclass

Bases: AttrConstraint[TensorType[Attribute] | UnrankedTensorType[Attribute]]

Converts an input memref constraint to the corresponding tensor constraint, i.e. the constraints on element type and shape are the same as the input constraint, but the attribute is verified to be a tensor instead of a memref.

Source code in xdsl/dialects/bufferization.py 
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
92
93
94
95
96
97
@dataclass(frozen=True)
class TensorFromMemRefConstraint(
    AttrConstraint[TensorType[Attribute] | UnrankedTensorType[Attribute]]
):
    """
    Converts an input memref constraint to the corresponding tensor constraint, i.e. the constraints
    on element type and shape are the same as the input constraint, but the attribute is verified to be
    a tensor instead of a memref.
    """

    memref_constraint: AttrConstraint[MemRefType | UnrankedMemRefType]

    @staticmethod
    def tensor_to_memref(
        tensor: TensorType | UnrankedTensorType,
    ) -> MemRefType | UnrankedMemRefType:
        if isinstance(tensor, TensorType):
            return MemRefType(tensor.element_type, tensor.shape)
        else:
            return UnrankedMemRefType.from_type(tensor.element_type)

    @staticmethod
    def memref_to_tensor(
        memref: MemRefType | UnrankedMemRefType,
    ) -> TensorType | UnrankedTensorType:
        if isinstance(memref, MemRefType):
            return TensorType(memref.element_type, memref.shape)
        else:
            return UnrankedTensorType(memref.element_type)

    def can_infer(self, var_constraint_names: AbstractSet[str]) -> bool:
        return self.memref_constraint.can_infer(var_constraint_names)

    def infer(
        self, context: ConstraintContext
    ) -> TensorType[Attribute] | UnrankedTensorType[Attribute]:
        memref_type = self.memref_constraint.infer(context)
        return self.memref_to_tensor(memref_type)

    def verify(self, attr: Attribute, constraint_context: ConstraintContext) -> None:
        if isa(attr, TensorType | UnrankedTensorType):
            memref_type = self.tensor_to_memref(attr)
        else:
            raise VerifyException(
                f"Expected tensor or unranked tensor type, got {attr}"
            )
        return self.memref_constraint.verify(memref_type, constraint_context)

    def get_bases(self) -> set[type[Attribute]] | None:
        return {TensorType, UnrankedTensorType}

    def mapping_type_vars(
        self, type_var_mapping: Mapping[TypeVar, AttrConstraint | IntConstraint]
    ) -> TensorFromMemRefConstraint:
        return TensorFromMemRefConstraint(
            self.memref_constraint.mapping_type_vars(type_var_mapping)
        )

memref_constraint: AttrConstraint[MemRefType | UnrankedMemRefType] instance-attribute

__init__(memref_constraint: AttrConstraint[MemRefType | UnrankedMemRefType]) -> None

tensor_to_memref(tensor: TensorType | UnrankedTensorType) -> MemRefType | UnrankedMemRefType staticmethod

Source code in xdsl/dialects/bufferization.py 
53
54
55
56
57
58
59
60
@staticmethod
def tensor_to_memref(
    tensor: TensorType | UnrankedTensorType,
) -> MemRefType | UnrankedMemRefType:
    if isinstance(tensor, TensorType):
        return MemRefType(tensor.element_type, tensor.shape)
    else:
        return UnrankedMemRefType.from_type(tensor.element_type)

memref_to_tensor(memref: MemRefType | UnrankedMemRefType) -> TensorType | UnrankedTensorType staticmethod

Source code in xdsl/dialects/bufferization.py 
62
63
64
65
66
67
68
69
@staticmethod
def memref_to_tensor(
    memref: MemRefType | UnrankedMemRefType,
) -> TensorType | UnrankedTensorType:
    if isinstance(memref, MemRefType):
        return TensorType(memref.element_type, memref.shape)
    else:
        return UnrankedTensorType(memref.element_type)

can_infer(var_constraint_names: AbstractSet[str]) -> bool

Source code in xdsl/dialects/bufferization.py 
71
72
def can_infer(self, var_constraint_names: AbstractSet[str]) -> bool:
    return self.memref_constraint.can_infer(var_constraint_names)

infer(context: ConstraintContext) -> TensorType[Attribute] | UnrankedTensorType[Attribute]

Source code in xdsl/dialects/bufferization.py 
74
75
76
77
78
def infer(
    self, context: ConstraintContext
) -> TensorType[Attribute] | UnrankedTensorType[Attribute]:
    memref_type = self.memref_constraint.infer(context)
    return self.memref_to_tensor(memref_type)

verify(attr: Attribute, constraint_context: ConstraintContext) -> None

Source code in xdsl/dialects/bufferization.py 
80
81
82
83
84
85
86
87
def verify(self, attr: Attribute, constraint_context: ConstraintContext) -> None:
    if isa(attr, TensorType | UnrankedTensorType):
        memref_type = self.tensor_to_memref(attr)
    else:
        raise VerifyException(
            f"Expected tensor or unranked tensor type, got {attr}"
        )
    return self.memref_constraint.verify(memref_type, constraint_context)

get_bases() -> set[type[Attribute]] | None

Source code in xdsl/dialects/bufferization.py 
89
90
def get_bases(self) -> set[type[Attribute]] | None:
    return {TensorType, UnrankedTensorType}

mapping_type_vars(type_var_mapping: Mapping[TypeVar, AttrConstraint | IntConstraint]) -> TensorFromMemRefConstraint

Source code in xdsl/dialects/bufferization.py 
92
93
94
95
96
97
def mapping_type_vars(
    self, type_var_mapping: Mapping[TypeVar, AttrConstraint | IntConstraint]
) -> TensorFromMemRefConstraint:
    return TensorFromMemRefConstraint(
        self.memref_constraint.mapping_type_vars(type_var_mapping)
    )

AllocTensorOp

Bases: IRDLOperation

bufferization.alloc_tensor materializes an uninitialized tensor with a given shape (dynamic or static). It always bufferizes to a new buffer allocation of the given shape. The optional copy operand specifies the contents of the tensors. If no copy operand is specified, reading from the result of an alloc_tensor op yields an undefined value.

If copy is specified, no dynamic sizes should be passed, since they are the same as the dynamic sizes of the copy operand.

alloc_tensor is a helper op for bufferization. The operation is provided as an anchor that marks the beginning of a new tensor SSA use-def chain. It can be used to control in-place bufferization decisions during One-Shot Bufferize: The bufferized result of a bufferization.alloc_tensor does not alias with any other buffer, so it can be used to resolve read-after-write conflicts that would have been introduced by the in-place bufferization of another op.

The optional memory_space attribute specifies the memory space when bufferizing this op. The memory space is inferred from copy if specified. If neither copy nor memory_space is specified, the default memory space is used during bufferization.

The optional size_hint operand specifies the number of non-zero elements for sparse tensors. The value of size_hint should be not less than 1 and not larger than the linear size of the corresponding dense tensor type. If this requirement is not met, the behavior of the operator is undefined.

Note: An alloc_tensor with a copy should also be expressed as an alloc_tensor without copy, followed by a copy_tensor.

https://mlir.llvm.org/docs/Dialects/BufferizationOps/#bufferizationalloc_tensor-bufferizationalloctensorop

Source code in xdsl/dialects/bufferization.py 
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
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
@irdl_op_definition
class AllocTensorOp(IRDLOperation):
    """
    `bufferization.alloc_tensor` materializes an uninitialized tensor with a
    given shape (dynamic or static). It always bufferizes to a new buffer
    allocation of the given shape. The optional `copy` operand specifies the
    contents of the tensors. If no `copy` operand is specified, reading from the
    result of an `alloc_tensor` op yields an undefined value.

    If `copy` is specified, no dynamic sizes should be passed, since they are
    the same as the dynamic sizes of the `copy` operand.

    `alloc_tensor` is a helper op for bufferization. The operation is provided
    as an anchor that marks the beginning of a new tensor SSA use-def chain. It
    can be used to control in-place bufferization decisions during One-Shot
    Bufferize: The bufferized result of a `bufferization.alloc_tensor` does not
    alias with any other buffer, so it can be used to resolve read-after-write
    conflicts that would have been introduced by the in-place bufferization of
    another op.

    The optional `memory_space` attribute specifies the memory space when
    bufferizing this op. The memory space is inferred from `copy` if specified.
    If neither `copy` nor `memory_space` is specified, the default memory space
    is used during bufferization.

    The optional `size_hint` operand specifies the number of non-zero elements
    for sparse tensors. The value of `size_hint` should be not less than 1 and
    not larger than the linear size of the corresponding dense tensor type. If
    this requirement is not met, the behavior of the operator is undefined.

    Note: An `alloc_tensor` with a `copy` should also be expressed as an
    `alloc_tensor` without `copy`, followed by a `copy_tensor`.

    https://mlir.llvm.org/docs/Dialects/BufferizationOps/#bufferizationalloc_tensor-bufferizationalloctensorop
    """

    name = "bufferization.alloc_tensor"

    T: ClassVar = VarConstraint("T", AnyTensorTypeConstr | AnyUnrankedTensorTypeConstr)

    dynamic_sizes = var_operand_def(IndexType())
    copy = opt_operand_def(T)
    size_hint = opt_operand_def(IndexType())

    tensor = result_def(T)

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

    assembly_format = "`(` $dynamic_sizes `)` ( `copy` `(` $copy^ `)`)? (`size_hint` `=` $size_hint^)? attr-dict `:` type($tensor)"  # noqa E501

    def __init__(
        self,
        result_type: Attribute,
        dynamic_sizes: Sequence[Operation | SSAValue] | None = None,
        copy: SSAValue | Operation | None = None,
        size_hint: SSAValue | Operation | None = None,
    ):
        super().__init__(
            operands=(dynamic_sizes, copy, size_hint),
            result_types=(result_type,),
        )

name = 'bufferization.alloc_tensor' class-attribute instance-attribute

T: ClassVar = VarConstraint('T', AnyTensorTypeConstr | AnyUnrankedTensorTypeConstr) class-attribute instance-attribute

dynamic_sizes = var_operand_def(IndexType()) class-attribute instance-attribute

copy = opt_operand_def(T) class-attribute instance-attribute

size_hint = opt_operand_def(IndexType()) class-attribute instance-attribute

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

irdl_options = (AttrSizedOperandSegments(as_property=True),) class-attribute instance-attribute

assembly_format = '`(` $dynamic_sizes `)` ( `copy` `(` $copy^ `)`)? (`size_hint` `=` $size_hint^)? attr-dict `:` type($tensor)' class-attribute instance-attribute

__init__(result_type: Attribute, dynamic_sizes: Sequence[Operation | SSAValue] | None = None, copy: SSAValue | Operation | None = None, size_hint: SSAValue | Operation | None = None)

Source code in xdsl/dialects/bufferization.py 
150
151
152
153
154
155
156
157
158
159
160
def __init__(
    self,
    result_type: Attribute,
    dynamic_sizes: Sequence[Operation | SSAValue] | None = None,
    copy: SSAValue | Operation | None = None,
    size_hint: SSAValue | Operation | None = None,
):
    super().__init__(
        operands=(dynamic_sizes, copy, size_hint),
        result_types=(result_type,),
    )

CloneOp

Bases: IRDLOperation

Source code in xdsl/dialects/bufferization.py 
163
164
165
166
167
168
169
170
171
172
173
174
175
176
@irdl_op_definition
class CloneOp(IRDLOperation):
    name = "bufferization.clone"

    T: ClassVar = VarConstraint("T", MemRefType.constr() | AnyUnrankedMemRefTypeConstr)

    input = operand_def(T)
    output = result_def(T)

    assembly_format = "$input attr-dict `:` type($input) `to` type($output)"

    def __init__(self, input: SSAValue | Operation):
        result_type = SSAValue.get(input).type
        super().__init__(operands=(input,), result_types=(result_type,))

name = 'bufferization.clone' class-attribute instance-attribute

T: ClassVar = VarConstraint('T', MemRefType.constr() | AnyUnrankedMemRefTypeConstr) class-attribute instance-attribute

input = operand_def(T) class-attribute instance-attribute

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

assembly_format = '$input attr-dict `:` type($input) `to` type($output)' class-attribute instance-attribute

__init__(input: SSAValue | Operation)

Source code in xdsl/dialects/bufferization.py 
174
175
176
def __init__(self, input: SSAValue | Operation):
    result_type = SSAValue.get(input).type
    super().__init__(operands=(input,), result_types=(result_type,))

ToTensorOp

Bases: IRDLOperation

Source code in xdsl/dialects/bufferization.py 
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
@irdl_op_definition
class ToTensorOp(IRDLOperation):
    name = "bufferization.to_tensor"

    T: ClassVar = VarConstraint("T", MemRefType.constr() | AnyUnrankedMemRefTypeConstr)

    memref = operand_def(T)
    tensor = result_def(TensorFromMemRefConstraint(T))

    writable = opt_prop_def(UnitAttr)
    restrict = opt_prop_def(UnitAttr)

    assembly_format = "$memref (`restrict` $restrict^)? (`writable` $writable^)? attr-dict `:` type($memref) `to` type($tensor)"

    def __init__(
        self,
        memref: SSAValue | Operation,
        restrict: bool = False,
        writable: bool = False,
    ):
        memref_v = SSAValue.get(memref, type=MemRefType | UnrankedMemRefType)
        properties = dict[str, Attribute]()
        if restrict:
            properties["restrict"] = UnitAttr()
        if writable:
            properties["writable"] = UnitAttr()
        super().__init__(
            operands=(memref,),
            result_types=(TensorFromMemRefConstraint.memref_to_tensor(memref_v.type),),
            properties=properties,
        )

name = 'bufferization.to_tensor' class-attribute instance-attribute

T: ClassVar = VarConstraint('T', MemRefType.constr() | AnyUnrankedMemRefTypeConstr) class-attribute instance-attribute

memref = operand_def(T) class-attribute instance-attribute

tensor = result_def(TensorFromMemRefConstraint(T)) class-attribute instance-attribute

writable = opt_prop_def(UnitAttr) class-attribute instance-attribute

restrict = opt_prop_def(UnitAttr) class-attribute instance-attribute

assembly_format = '$memref (`restrict` $restrict^)? (`writable` $writable^)? attr-dict `:` type($memref) `to` type($tensor)' class-attribute instance-attribute

__init__(memref: SSAValue | Operation, restrict: bool = False, writable: bool = False)

Source code in xdsl/dialects/bufferization.py 
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
def __init__(
    self,
    memref: SSAValue | Operation,
    restrict: bool = False,
    writable: bool = False,
):
    memref_v = SSAValue.get(memref, type=MemRefType | UnrankedMemRefType)
    properties = dict[str, Attribute]()
    if restrict:
        properties["restrict"] = UnitAttr()
    if writable:
        properties["writable"] = UnitAttr()
    super().__init__(
        operands=(memref,),
        result_types=(TensorFromMemRefConstraint.memref_to_tensor(memref_v.type),),
        properties=properties,
    )

ToBufferOp dataclass

Bases: IRDLOperation

Source code in xdsl/dialects/bufferization.py 
212
213
214
215
216
217
218
219
220
221
222
@irdl_op_definition
class ToBufferOp(IRDLOperation):
    name = "bufferization.to_buffer"

    T: ClassVar = VarConstraint("T", MemRefType.constr() | AnyUnrankedMemRefTypeConstr)
    tensor = operand_def(TensorFromMemRefConstraint(T))
    memref = result_def(T)

    read_only = opt_prop_def(UnitAttr)

    assembly_format = "$tensor (`read_only` $read_only^)?  `:` attr-dict type($tensor) `to` type($memref)"

name = 'bufferization.to_buffer' class-attribute instance-attribute

T: ClassVar = VarConstraint('T', MemRefType.constr() | AnyUnrankedMemRefTypeConstr) class-attribute instance-attribute

tensor = operand_def(TensorFromMemRefConstraint(T)) class-attribute instance-attribute

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

read_only = opt_prop_def(UnitAttr) class-attribute instance-attribute

assembly_format = '$tensor (`read_only` $read_only^)? `:` attr-dict type($tensor) `to` type($memref)' class-attribute instance-attribute

MaterializeInDestinationOp dataclass

Bases: IRDLOperation

Source code in xdsl/dialects/bufferization.py 
225
226
227
228
229
230
231
232
233
234
235
236
237
@irdl_op_definition
class MaterializeInDestinationOp(IRDLOperation):
    name = "bufferization.materialize_in_destination"

    T: ClassVar = VarConstraint("T", MemRefType.constr() | AnyUnrankedMemRefTypeConstr)
    source = operand_def(TensorFromMemRefConstraint(T))
    dest = operand_def(T | TensorFromMemRefConstraint(T))
    result = opt_result_def(TensorFromMemRefConstraint(T))

    restrict = opt_prop_def(UnitAttr)
    writable = opt_prop_def(UnitAttr)

    assembly_format = "$source `in` (`restrict` $restrict^)? (`writable` $writable^)? $dest attr-dict `:` functional-type(operands, results)"  # noqa: E501

name = 'bufferization.materialize_in_destination' class-attribute instance-attribute

T: ClassVar = VarConstraint('T', MemRefType.constr() | AnyUnrankedMemRefTypeConstr) class-attribute instance-attribute

source = operand_def(TensorFromMemRefConstraint(T)) class-attribute instance-attribute

dest = operand_def(T | TensorFromMemRefConstraint(T)) class-attribute instance-attribute

result = opt_result_def(TensorFromMemRefConstraint(T)) class-attribute instance-attribute

restrict = opt_prop_def(UnitAttr) class-attribute instance-attribute

writable = opt_prop_def(UnitAttr) class-attribute instance-attribute

assembly_format = '$source `in` (`restrict` $restrict^)? (`writable` $writable^)? $dest attr-dict `:` functional-type(operands, results)' class-attribute instance-attribute