Skip to content

X86 scf

x86_scf

X86_Scf = Dialect('x86_scf', [ForOp, RofOp, YieldOp]) module-attribute

YieldOp dataclass

Bases: AbstractYieldOperation[X86RegisterType]

Source code in xdsl/dialects/x86_scf.py
52
53
54
55
56
57
58
59
60
61
62
@irdl_op_definition
class YieldOp(AbstractYieldOperation[X86RegisterType]):
    name = "x86_scf.yield"

    traits = lazy_traits_def(
        lambda: (
            IsTerminator(),
            HasParent(ForRofOperation),
            NoMemoryEffect(),
        )
    )

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

traits = lazy_traits_def(lambda: (IsTerminator(), HasParent(ForRofOperation), NoMemoryEffect())) class-attribute instance-attribute

ForRofOperation

Bases: X86HasRegisterConstraints, ABC

Source code in xdsl/dialects/x86_scf.py
 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
 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
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
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
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
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
class ForRofOperation(X86HasRegisterConstraints, ABC):
    lb = operand_def(GeneralRegisterType)
    ub_val = opt_operand_def(GeneralRegisterType)
    ub_attr = opt_prop_def(IntegerAttr[SI32])
    step_val = opt_operand_def(GeneralRegisterType)
    step_attr = opt_prop_def(IntegerAttr[SI32])

    iter_args = var_operand_def(X86RegisterType)

    lb_end = result_def(GeneralRegisterType)
    """Final value of the lower-bound / induction-variable register (inout with `lb`)."""

    res = var_result_def(X86RegisterType)

    body = region_def("single_block")

    traits = traits_def(SingleBlockImplicitTerminator(YieldOp), RecursiveMemoryEffect())
    irdl_options = (AttrSizedOperandSegments(as_property=True),)

    @property
    def ub(self) -> IntegerAttr[SI32] | SSAValue:
        """Static upper bound (typed integer) or dynamic register SSA value."""
        if self.ub_attr is not None:
            return self.ub_attr
        else:
            assert self.ub_val is not None, (
                "Exactly one of ub_attr or ub_val must be set"
            )
            return self.ub_val

    @property
    def step(self) -> IntegerAttr[SI32] | SSAValue:
        """Static step (typed integer) or dynamic register SSA value."""
        if self.step_attr is not None:
            return self.step_attr
        else:
            assert self.step_val is not None, (
                "Exactly one of step_attr or step_val must be set"
            )
            return self.step_val

    def __init__(
        self,
        lb: SSAValue | Operation,
        ub: SSAValue | Operation | IntegerAttr,
        step: SSAValue | Operation | IntegerAttr,
        iter_args: Sequence[SSAValue],
        body: Region | Sequence[Operation] | Sequence[Block] | Block | None = None,
    ):
        lb = SSAValue.get(lb)
        if body is None:
            body = Region(
                Block(arg_types=(lb.type, *(iter_arg.type for iter_arg in iter_args)))
            )

        if isinstance(body, Block):
            body = [body]

        if isinstance(ub, IntegerAttr):
            ub_attr = ub
            ub_val = None
        else:
            ub_attr = None
            ub_val = ub

        if isinstance(step, IntegerAttr):
            step_attr = step
            step_val = None
        else:
            step_attr = None
            step_val = step

        super().__init__(
            operands=[lb, ub_val, step_val, iter_args],
            properties={"ub_attr": ub_attr, "step_attr": step_attr},
            result_types=[lb.type, [SSAValue.get(a).type for a in iter_args]],
            regions=[body],
        )

    def verify_(self):
        if (self.ub_attr is None) == (self.ub_val is None):
            raise VerifyException(
                "Exactly one of ub_attr (static) or ub_val (dynamic) must be set, "
                f"got ub_attr={self.ub_attr}, ub_val={self.ub_val}"
            )
        if (self.step_attr is None) == (self.step_val is None):
            raise VerifyException(
                "Exactly one of step_attr (static) or step_val (dynamic) must be set, "
                f"got step_attr={self.step_attr}, step_val={self.step_val}"
            )
        if (len(self.iter_args) + 1) != len(self.body.block.args):
            raise VerifyException(
                f"Wrong number of block arguments, expected {len(self.iter_args) + 1}, got "
                f"{len(self.body.block.args)}. The body must have the induction "
                f"variable and loop-carried variables as arguments."
            )
        if self.body.block.args and (iter_var := self.body.block.args[0]):
            if not isinstance(iter_var.type, GeneralRegisterType):
                raise VerifyException(
                    f"The first block argument of the body is of type {iter_var.type}"
                    " instead of x86 GeneralRegisterType"
                )
            if iter_var.type != self.lb.type:
                raise VerifyException(
                    f"Expected induction var to be same type as lb, "
                    f"got {iter_var.type} and {self.lb.type}"
                )
            if iter_var.type != self.lb_end.type:
                raise VerifyException(
                    f"Expected induction var to be same type as lb_end result, "
                    f"got {iter_var.type} and {self.lb_end.type}"
                )
        for idx, (arg, block_arg) in enumerate(
            zip(self.iter_args, self.body.block.args[1:])
        ):
            if block_arg.type != arg.type:
                raise VerifyException(
                    f"Block argument {idx + 1} has wrong type, expected {arg.type}, "
                    f"got {block_arg.type}. Arguments after the "
                    f"induction variable must match the carried variables."
                )
        if len(self.body.ops) > 0 and isinstance(
            yieldop := self.body.block.last_op, YieldOp
        ):
            if len(yieldop.arguments) != len(self.iter_args):
                raise VerifyException(
                    f"Expected {len(self.iter_args)} args, got {len(yieldop.arguments)}. "
                    f"The riscv_scf.for must yield its carried variables."
                )
            for iter_arg, yield_arg in zip(self.iter_args, yieldop.arguments):
                if iter_arg.type != yield_arg.type:
                    raise VerifyException(
                        f"Expected {iter_arg.type}, got {yield_arg.type}. The "
                        f"riscv_scf.for's riscv_scf.yield must match carried"
                        f"variables types."
                    )

    def allocate_registers(self, allocator: BlockAllocator) -> None:
        """Allocate loop-carried and IV registers, then the body under those reservations."""
        # Allocate values used inside the body but defined outside.
        # Their scope lasts for the whole body execution scope
        live_ins = allocator.live_ins_per_block[self.body.block]
        for live_in in live_ins:
            allocator.allocate_value(live_in)

        yield_op = self.body.block.last_op
        assert yield_op is not None, (
            "last op of riscv_scf.ForOp is guaranteed to be riscv_scf.Yield"
        )
        block_args = self.body.block.args

        # The loop-carried variables are trickier
        # The for op operand, block arg, and yield operand must have the same type
        for block_arg, operand, yield_operand, op_result in zip(
            block_args[1:], self.iter_args, yield_op.operands, self.res, strict=True
        ):
            allocator.allocate_values_same_reg(
                (block_arg, operand, yield_operand, op_result)
            )

        allocator.allocate_values_same_reg((block_args[0], self.lb, self.lb_end))

        # ub and step are used throughout loop when dynamic
        if self.ub_val is not None:
            allocator.allocate_value(self.ub_val)
        if self.step_val is not None:
            allocator.allocate_value(self.step_val)

        # Reserve the loop carried variables for allocation within the body
        regs = self.iter_args.types
        assert all(isinstance(reg, X86RegisterType) for reg in regs)
        regs = cast(tuple[X86RegisterType, ...], regs)
        with allocator.available_registers.reserve_registers(regs):
            allocator.allocate_block(self.body.block)

    def get_register_constraints(self) -> RegisterConstraints:
        """`lb` and each iter_arg are inout; dynamic `ub`/`step` are in-only."""
        ins: list[SSAValue] = []
        if self.ub_val is not None:
            ins.append(self.ub_val)
        if self.step_val is not None:
            ins.append(self.step_val)
        inouts = ((self.lb, self.lb_end), *zip(self.iter_args, self.res, strict=True))
        return RegisterConstraints(ins, (), inouts)

    def _body_live_outs(self, live_after: AbstractSet[SSAValue]) -> set[SSAValue]:
        """
        Values live at the end of the loop body.
        """
        block = self.body.block
        res = set(live_after)
        # The body runs repeatedly, so every value defined outside it and used inside
        # must survive a whole iteration. This covers the dynamic bounds: the loop
        # re-reads them on the back edge, and a clobber in the body is itself a use.
        res.update(live_ins_per_block(block)[block])
        # The induction variable is a block argument rather than a live-in, but the
        # loop reads it on the back edge to compute the next value.
        res.add(block.args[0])
        return res

    def update_liveness(self, ctx: LivenessContext) -> None:
        # Create a new context to use inside the loop
        body_ctx = ctx.copy(self._body_live_outs(ctx.alive))
        body_ctx.process_block(self.body.block)
        # Update the outer context with all the values that are alive coming into the
        # body
        ctx.alive.update(body_ctx.alive)
        # HasRegisterConstraints default implementation
        super().update_liveness(ctx)

lb = operand_def(GeneralRegisterType) class-attribute instance-attribute

ub_val = opt_operand_def(GeneralRegisterType) class-attribute instance-attribute

ub_attr = opt_prop_def(IntegerAttr[SI32]) class-attribute instance-attribute

step_val = opt_operand_def(GeneralRegisterType) class-attribute instance-attribute

step_attr = opt_prop_def(IntegerAttr[SI32]) class-attribute instance-attribute

iter_args = var_operand_def(X86RegisterType) class-attribute instance-attribute

lb_end = result_def(GeneralRegisterType) class-attribute instance-attribute

Final value of the lower-bound / induction-variable register (inout with lb).

res = var_result_def(X86RegisterType) class-attribute instance-attribute

body = region_def('single_block') class-attribute instance-attribute

traits = traits_def(SingleBlockImplicitTerminator(YieldOp), RecursiveMemoryEffect()) class-attribute instance-attribute

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

ub: IntegerAttr[SI32] | SSAValue property

Static upper bound (typed integer) or dynamic register SSA value.

step: IntegerAttr[SI32] | SSAValue property

Static step (typed integer) or dynamic register SSA value.

__init__(lb: SSAValue | Operation, ub: SSAValue | Operation | IntegerAttr, step: SSAValue | Operation | IntegerAttr, iter_args: Sequence[SSAValue], body: Region | Sequence[Operation] | Sequence[Block] | Block | None = None)

Source code in xdsl/dialects/x86_scf.py
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
def __init__(
    self,
    lb: SSAValue | Operation,
    ub: SSAValue | Operation | IntegerAttr,
    step: SSAValue | Operation | IntegerAttr,
    iter_args: Sequence[SSAValue],
    body: Region | Sequence[Operation] | Sequence[Block] | Block | None = None,
):
    lb = SSAValue.get(lb)
    if body is None:
        body = Region(
            Block(arg_types=(lb.type, *(iter_arg.type for iter_arg in iter_args)))
        )

    if isinstance(body, Block):
        body = [body]

    if isinstance(ub, IntegerAttr):
        ub_attr = ub
        ub_val = None
    else:
        ub_attr = None
        ub_val = ub

    if isinstance(step, IntegerAttr):
        step_attr = step
        step_val = None
    else:
        step_attr = None
        step_val = step

    super().__init__(
        operands=[lb, ub_val, step_val, iter_args],
        properties={"ub_attr": ub_attr, "step_attr": step_attr},
        result_types=[lb.type, [SSAValue.get(a).type for a in iter_args]],
        regions=[body],
    )

verify_()

Source code in xdsl/dialects/x86_scf.py
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
def verify_(self):
    if (self.ub_attr is None) == (self.ub_val is None):
        raise VerifyException(
            "Exactly one of ub_attr (static) or ub_val (dynamic) must be set, "
            f"got ub_attr={self.ub_attr}, ub_val={self.ub_val}"
        )
    if (self.step_attr is None) == (self.step_val is None):
        raise VerifyException(
            "Exactly one of step_attr (static) or step_val (dynamic) must be set, "
            f"got step_attr={self.step_attr}, step_val={self.step_val}"
        )
    if (len(self.iter_args) + 1) != len(self.body.block.args):
        raise VerifyException(
            f"Wrong number of block arguments, expected {len(self.iter_args) + 1}, got "
            f"{len(self.body.block.args)}. The body must have the induction "
            f"variable and loop-carried variables as arguments."
        )
    if self.body.block.args and (iter_var := self.body.block.args[0]):
        if not isinstance(iter_var.type, GeneralRegisterType):
            raise VerifyException(
                f"The first block argument of the body is of type {iter_var.type}"
                " instead of x86 GeneralRegisterType"
            )
        if iter_var.type != self.lb.type:
            raise VerifyException(
                f"Expected induction var to be same type as lb, "
                f"got {iter_var.type} and {self.lb.type}"
            )
        if iter_var.type != self.lb_end.type:
            raise VerifyException(
                f"Expected induction var to be same type as lb_end result, "
                f"got {iter_var.type} and {self.lb_end.type}"
            )
    for idx, (arg, block_arg) in enumerate(
        zip(self.iter_args, self.body.block.args[1:])
    ):
        if block_arg.type != arg.type:
            raise VerifyException(
                f"Block argument {idx + 1} has wrong type, expected {arg.type}, "
                f"got {block_arg.type}. Arguments after the "
                f"induction variable must match the carried variables."
            )
    if len(self.body.ops) > 0 and isinstance(
        yieldop := self.body.block.last_op, YieldOp
    ):
        if len(yieldop.arguments) != len(self.iter_args):
            raise VerifyException(
                f"Expected {len(self.iter_args)} args, got {len(yieldop.arguments)}. "
                f"The riscv_scf.for must yield its carried variables."
            )
        for iter_arg, yield_arg in zip(self.iter_args, yieldop.arguments):
            if iter_arg.type != yield_arg.type:
                raise VerifyException(
                    f"Expected {iter_arg.type}, got {yield_arg.type}. The "
                    f"riscv_scf.for's riscv_scf.yield must match carried"
                    f"variables types."
                )

allocate_registers(allocator: BlockAllocator) -> None

Allocate loop-carried and IV registers, then the body under those reservations.

Source code in xdsl/dialects/x86_scf.py
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
def allocate_registers(self, allocator: BlockAllocator) -> None:
    """Allocate loop-carried and IV registers, then the body under those reservations."""
    # Allocate values used inside the body but defined outside.
    # Their scope lasts for the whole body execution scope
    live_ins = allocator.live_ins_per_block[self.body.block]
    for live_in in live_ins:
        allocator.allocate_value(live_in)

    yield_op = self.body.block.last_op
    assert yield_op is not None, (
        "last op of riscv_scf.ForOp is guaranteed to be riscv_scf.Yield"
    )
    block_args = self.body.block.args

    # The loop-carried variables are trickier
    # The for op operand, block arg, and yield operand must have the same type
    for block_arg, operand, yield_operand, op_result in zip(
        block_args[1:], self.iter_args, yield_op.operands, self.res, strict=True
    ):
        allocator.allocate_values_same_reg(
            (block_arg, operand, yield_operand, op_result)
        )

    allocator.allocate_values_same_reg((block_args[0], self.lb, self.lb_end))

    # ub and step are used throughout loop when dynamic
    if self.ub_val is not None:
        allocator.allocate_value(self.ub_val)
    if self.step_val is not None:
        allocator.allocate_value(self.step_val)

    # Reserve the loop carried variables for allocation within the body
    regs = self.iter_args.types
    assert all(isinstance(reg, X86RegisterType) for reg in regs)
    regs = cast(tuple[X86RegisterType, ...], regs)
    with allocator.available_registers.reserve_registers(regs):
        allocator.allocate_block(self.body.block)

get_register_constraints() -> RegisterConstraints

lb and each iter_arg are inout; dynamic ub/step are in-only.

Source code in xdsl/dialects/x86_scf.py
240
241
242
243
244
245
246
247
248
def get_register_constraints(self) -> RegisterConstraints:
    """`lb` and each iter_arg are inout; dynamic `ub`/`step` are in-only."""
    ins: list[SSAValue] = []
    if self.ub_val is not None:
        ins.append(self.ub_val)
    if self.step_val is not None:
        ins.append(self.step_val)
    inouts = ((self.lb, self.lb_end), *zip(self.iter_args, self.res, strict=True))
    return RegisterConstraints(ins, (), inouts)

update_liveness(ctx: LivenessContext) -> None

Source code in xdsl/dialects/x86_scf.py
265
266
267
268
269
270
271
272
273
def update_liveness(self, ctx: LivenessContext) -> None:
    # Create a new context to use inside the loop
    body_ctx = ctx.copy(self._body_live_outs(ctx.alive))
    body_ctx.process_block(self.body.block)
    # Update the outer context with all the values that are alive coming into the
    # body
    ctx.alive.update(body_ctx.alive)
    # HasRegisterConstraints default implementation
    super().update_liveness(ctx)

ForOp dataclass

Bases: ForRofOperation

A for loop, counting up from lb to ub by step each iteration.

Source code in xdsl/dialects/x86_scf.py
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
@irdl_op_definition
class ForOp(ForRofOperation):
    """
    A for loop, counting up from lb to ub by step each iteration.
    """

    name = "x86_scf.for"

    def print(self, printer: Printer):
        print_for_op_like(
            printer,
            self.lb,
            self.ub,
            self.step,
            self.iter_args,
            self.body,
        )

    @classmethod
    def parse(cls, parser: Parser) -> Self:
        lb, ub, step, iter_arg_operands, body = parse_for_op_like(
            parser, allow_static_upper_bound=True, allow_static_step=True
        )
        _, *iter_args = body.block.args

        for_op = cls(lb, ub, step, iter_arg_operands, body)

        if not iter_args:
            for trait in for_op.get_traits_of_type(SingleBlockImplicitTerminator):
                ensure_terminator(for_op, trait)

        return for_op

name = 'x86_scf.for' class-attribute instance-attribute

print(printer: Printer)

Source code in xdsl/dialects/x86_scf.py
284
285
286
287
288
289
290
291
292
def print(self, printer: Printer):
    print_for_op_like(
        printer,
        self.lb,
        self.ub,
        self.step,
        self.iter_args,
        self.body,
    )

parse(parser: Parser) -> Self classmethod

Source code in xdsl/dialects/x86_scf.py
294
295
296
297
298
299
300
301
302
303
304
305
306
307
@classmethod
def parse(cls, parser: Parser) -> Self:
    lb, ub, step, iter_arg_operands, body = parse_for_op_like(
        parser, allow_static_upper_bound=True, allow_static_step=True
    )
    _, *iter_args = body.block.args

    for_op = cls(lb, ub, step, iter_arg_operands, body)

    if not iter_args:
        for trait in for_op.get_traits_of_type(SingleBlockImplicitTerminator):
            ensure_terminator(for_op, trait)

    return for_op

RofOp dataclass

Bases: ForRofOperation

Reverse Order For loop.

MLIR's for loops have the constraint of always executing from lb to ub, so in order to express loops that count down from ub to lb, the rof op is needed.

Rof has the semantics of going from ub to lb, decrementing by step each time. The implicit constraints are that lb < ub, and step > 0.

In order to convert a for to a rof, one needs to switch lb and ub. (for the normalized case that (ub - lb) % step == 0)

Source code in xdsl/dialects/x86_scf.py
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
@irdl_op_definition
class RofOp(ForRofOperation):
    """
    Reverse Order For loop.

    MLIR's for loops have the constraint of always executing from lb to ub,
    so in order to express loops that count down from ub to lb, the rof op
    is needed.

    Rof has the semantics of going from ub to lb, decrementing by step each time.
    The implicit constraints are that lb < ub, and step > 0.

    In order to convert a for to a rof, one needs to switch lb and ub.
    (for the normalized case that (ub - lb) % step == 0)
    """

    name = "x86_scf.rof"

    def print(self, printer: Printer):
        print_for_op_like(
            printer,
            self.ub,
            self.lb,
            self.step,
            self.iter_args,
            self.body,
            bound_words=["down", "to"],
        )

    @classmethod
    def parse(cls, parser: Parser) -> Self:
        ub, lb, step, iter_arg_operands, body = parse_for_op_like(
            parser,
            bound_words=["down", "to"],
            allow_static_upper_bound=True,
            allow_static_step=True,
        )
        _, *iter_args = body.block.args

        if isinstance(lb, IntegerAttr):
            parser.raise_error("Expected an operand.")

        rof_op = cls(lb, ub, step, iter_arg_operands, body)

        if not iter_args:
            for trait in rof_op.get_traits_of_type(SingleBlockImplicitTerminator):
                ensure_terminator(rof_op, trait)

        return rof_op

name = 'x86_scf.rof' class-attribute instance-attribute

print(printer: Printer)

Source code in xdsl/dialects/x86_scf.py
328
329
330
331
332
333
334
335
336
337
def print(self, printer: Printer):
    print_for_op_like(
        printer,
        self.ub,
        self.lb,
        self.step,
        self.iter_args,
        self.body,
        bound_words=["down", "to"],
    )

parse(parser: Parser) -> Self classmethod

Source code in xdsl/dialects/x86_scf.py
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
@classmethod
def parse(cls, parser: Parser) -> Self:
    ub, lb, step, iter_arg_operands, body = parse_for_op_like(
        parser,
        bound_words=["down", "to"],
        allow_static_upper_bound=True,
        allow_static_step=True,
    )
    _, *iter_args = body.block.args

    if isinstance(lb, IntegerAttr):
        parser.raise_error("Expected an operand.")

    rof_op = cls(lb, ub, step, iter_arg_operands, body)

    if not iter_args:
        for trait in rof_op.get_traits_of_type(SingleBlockImplicitTerminator):
            ensure_terminator(rof_op, trait)

    return rof_op