Function for_loop

python/cuda_tile/dialects/cuda_tile_ops.py:3937–4001

Constructs a for loop with the provided body. The body is a function taking as argument the iteration variables and building the operations within the body (including continue and break). By default, only the induction variable is created. If initializers for additional iterati

(
    body: Callable,
    lower_bound: int | Tile,
    upper_bound: int | Tile,
    step: int | Tile = 1,
    init_values: Sequence[Tile] = (),
    el_type=Int32,
    *,
    unsigned: bool = False,
    loc=None,
    ip=None,
)

@cuda_tile_op
def for_loop(
    body: Callable,
    lower_bound: int | Tile,
    upper_bound: int | Tile,
    step: int | Tile = 1,
    init_values: Sequence[Tile] = (),
    el_type=Int32,
    *,
    unsigned: bool = False,
    loc=None,
    ip=None,
) -> Tuple[Tile, ...]:
    """
    Constructs a for loop with the provided body. The body is a function taking
    as argument the iteration variables and building the operations within the
    body (including continue and break).

    By default, only the induction variable is created. If initializers for
    additional iteration variables are provided in `init_values`, additional
    iteration variables will be passed to the body and returned from the
    operation.

    By default, the induction variable element type is Int32, which can be
    overriden by setting `el_type`.

    By default, signed comparison is used for loop termination. Set `unsigned=True`
    to use unsigned integer comparison.
    """

    index_type = el_type.mlir_type

    def check_scalar(x: int | Tile, name: str) -> Tile:
        nonlocal index_type
        if isinstance(x, int):
            return constant(x, el_type)
        elif isinstance(x, Tile) and x.element_type == index_type and len(x.shape) == 0:
            return x
        else:
            raise TypeError(
                f"For loop {name} must be an integer or a scalar {el_type} tile value, got {x}"
            )

    lower_bound = check_scalar(lower_bound, "lower bound")
    upper_bound = check_scalar(upper_bound, "upper bound")
    step = check_scalar(step, "step")

    iter_arg_types = tuple(x.tile_type for x in init_values)
    _for_op = _cuda_tile.ForOp(
        resultValues=iter_arg_types,
        lowerBound=lower_bound,
        upperBound=upper_bound,
        step=step,
        initValues=init_values,
        unsignedCmp=unsigned,
        loc=loc,
        ip=ip,
    )