Convert stencil to ll mlir

@dataclass
class CastOpToMemRef(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: CastOp, rewriter: PatternRewriter, /):
        assert isa(op.result.type, FieldType[Attribute])
        assert isinstance(op.result.type.bounds, StencilBoundsAttr)

        result_type = StencilToMemRefType(op.result.type)

        cast = memref.CastOp.get(op.field, result_type)

        rewriter.replace(op, cast)

@dataclass
class ReturnOpToMemRef(RewritePattern):
    return_target: dict[ApplyOp, list[SSAValue | None]]

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: ReturnOp, rewriter: PatternRewriter, /):
        unroll_factor = op.unroll_factor
        n_res = len(op.arg) // unroll_factor

        store_list: list[Operation] = []
        apply = op.parent_op()
        assert isinstance(apply, ApplyOp)
        if op.unroll is not None:
            apply.attributes["__unroll__"] = op.unroll

        n_dims = apply.get_rank()

        for j in range(n_res):
            if len(apply.res) > 0:
                target = self.return_target[apply][j]
                # Insert pending subview if it was created but not yet inserted
                # (StencilStoreToSubview defers insertion to avoid dead code elimination)
                if (
                    target is not None
                    and isinstance(target.owner, memref.SubviewOp)
                    and target.owner.parent_block() is None
                ):
                    subview_op = target.owner
                    field = subview_op.source
                    if isinstance(field.owner, Operation):
                        rewriter.insert(subview_op, InsertPoint.after(field.owner))
                    else:
                        rewriter.insert(subview_op, InsertPoint.at_start(field.owner))
            else:
                target = apply.dest[j]
                rewriter.insert(
                    subview := field_subview(target), InsertPoint.before(apply)
                )
                target = subview

            unroll = op.unroll
            if unroll is None:
                unroll = IndexAttr.from_indices(*([1] * n_dims))

            for k, offset in enumerate(product(*(range(u) for u in unroll))):
                arg = op.arg[j * unroll_factor + k]
                index_ops: list[Operation] = list(
                    IndexOp(
                        properties={
                            "dim": builtin.IntegerAttr.from_index_int_value(i),
                            "offset": IndexAttr.from_indices(*([0] * n_dims)),
                        },
                        result_types=[builtin.IndexType()],
                    )
                    for i in range(n_dims)
                )
                store_list += index_ops

                for i in range(n_dims):
                    if offset[i] != 0:
                        constant_op = arith.ConstantOp.from_int_and_width(
                            offset[i], builtin.IndexType()
                        )
                        add_op = arith.AddiOp(index_ops[i], constant_op)
                        index_ops[i] = add_op
                        store_list.append(constant_op)
                        store_list.append(add_op)

                if isinstance(arg.type, ResultType):
                    result_owner = _find_result_store(arg)
                    for owner in result_owner:
                        if owner.arg:
                            if target is not None:
                                store = memref.StoreOp.get(owner.arg, target, index_ops)
                            else:
                                store = list[Operation]()
                            rewriter.replace(
                                owner,
                                store,
                                new_results=[owner.arg],
                            )
                        else:
                            dummy = UnrealizedConversionCastOp.get([], [arg.type.elem])
                            rewriter.replace(owner, dummy)

                else:
                    if target is not None:
                        store = memref.StoreOp.get(arg, target, index_ops)
                        store_list.append(store)

        rewriter.insert(store_list)
        rewriter.erase(op)

class LoadOpToMemRef(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: LoadOp, rewriter: PatternRewriter, /):
        for use in op.field.uses:
            if isa(use.operation, StoreOp):
                raise VerifyException(
                    "Cannot lower directly if loading and storing the same field! Try running `stencil-bufferize` before."
                )
        field = op.field.type
        assert isa(field, FieldType[Attribute])
        assert isa(field.bounds, StencilBoundsAttr)
        temp = op.res.type
        assert isa(temp, TempType[Attribute])
        assert isa(temp.bounds, StencilBoundsAttr)

        assert_subset(field, temp)

        offsets = [i for i in -field.bounds.lb]
        sizes = [i for i in temp.get_shape()]
        strides = [1] * len(sizes)

        subview = memref.SubviewOp.from_static_parameters(
            op.field, StencilToMemRefType(field), offsets, sizes, strides
        )

        rewriter.replace(op, subview)
        name = None
        if subview.source.name_hint:
            name = subview.source.name_hint + "_loadview"
        subview.result.name_hint = name

@dataclass
class BufferOpToMemRef(RewritePattern):
    return_targets: dict[ApplyOp, list[SSAValue | None]]

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: BufferOp, rewriter: PatternRewriter, /):
        # The current lowering simply allocate at block entry and deallocates at block
        # exit.
        # One could be smarter, e.g., aim precisly at where the first and last need is
        # But first, this requires more code with stencil.combine in the mix.
        # And second, do we want to be that smart? We use it iterated either way, so
        # probably always hoping to hoist (de)allocations out of the loop anyway?
        temp_t = op.temp.type
        assert isa(temp_t, TempType[Attribute])
        temp_bounds = temp_t.bounds
        assert isa(temp_bounds, StencilBoundsAttr)

        block = op.parent_block()
        assert block is not None
        first_op = block.first_op
        assert first_op is not None
        last_op = block.last_op
        assert last_op is not None

        shape = temp_t.get_shape()
        strides = [prod(shape[i + 1 :]) for i in range(len(shape))]
        offset = -sum(o * s for o, s in zip(temp_bounds.lb, strides, strict=True))

        layout = memref.StridedLayoutAttr(strides, offset)
        alloc = memref.AllocOp.get(
            temp_t.get_element_type(), shape=temp_t.get_shape(), layout=layout
        )
        alloc_type = alloc.memref.type
        assert isa(alloc_type, MemRefType)

        rewriter.insert(alloc, InsertPoint.before(first_op))

        update_return_target(self.return_targets, op.temp, alloc.memref)

        dealloc = memref.DeallocOp.get(alloc.memref)

        if not op.res.uses:
            rewriter.insert(dealloc, InsertPoint.after(op))
            rewriter.erase(op)
            return

        rewriter.insert(dealloc, InsertPoint.before(last_op))
        rewriter.replace(op, [], [alloc.memref])

class AllocOpToMemRef(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: AllocOp, rewriter: PatternRewriter, /):
        alloc = memref.AllocOp(
            [], [], StencilToMemRefType(cast(StencilType[Attribute], op.field.type))
        )
        rewriter.replace(op, alloc)

@dataclass
class ApplyOpFieldSubviews(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: ApplyOp, rewriter: PatternRewriter, /):
        args = [
            field_subview(arg) if isinstance(arg.type, FieldType) else arg
            for arg in op.args
        ]
        if args == list(op.args):
            return

        new_apply = ApplyOp.create(
            operands=[SSAValue.get(arg) for arg in args] + list(op.dest),
            result_types=[r.type for r in op.res],
            regions=[op.detach_region(0)],
            attributes=op.attributes,
            properties=op.properties,
        )
        rewriter.replace(
            op, [*(arg for arg in args if isinstance(arg, Operation)), new_apply]
        )

@dataclass
class ApplyOpToParallel(RewritePattern):
    return_targets: dict[ApplyOp, list[SSAValue | None]]

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: ApplyOp, rewriter: PatternRewriter, /):
        if len(op.res) > 0:
            res_type = op.res[0].type
            assert isa(res_type, TempType[Attribute])
            assert isinstance(res_type.bounds, StencilBoundsAttr)
            lb = res_type.bounds.lb
            ub = res_type.bounds.ub
        else:
            assert op.bounds is not None
            lb = op.bounds.lb
            ub = op.bounds.ub

        # Get this apply's ReturnOp
        unroll = op.attributes.get("__unroll__", None)
        if unroll is not None:
            assert isinstance(unroll, IndexAttr)
            unroll = list(u for u in unroll)

        rank = op.get_rank()
        body = prepare_apply_body(op)
        if unroll is None:
            unroll = [1] * rank
        else:
            unroll = [i for i in unroll]

        # Then create the corresponding scf.parallel
        boilerplate_ops = [
            *(
                lowerBounds := [
                    arith.ConstantOp.from_int_and_width(x, builtin.IndexType())
                    for x in lb
                ]
            ),
            *(
                steps := [
                    arith.ConstantOp.from_int_and_width(x, builtin.IndexType())
                    for x in unroll
                ]
            ),
            *(
                upperBounds := [
                    arith.ConstantOp.from_int_and_width(x, builtin.IndexType())
                    for x in ub
                ]
            ),
        ]

        # Generate an outer parallel loop as well as two inner sequential
        # loops. The inner sequential loops ensure that the computational
        # kernel itself is not slowed down by the OpenMP runtime.
        tiled_steps = steps
        p = scf.ParallelOp(
            lower_bounds=lowerBounds,
            upper_bounds=upperBounds,
            steps=tiled_steps,
            body=Region(body),
        )
        for index in body.walk():
            if isinstance(index, IndexOp):
                offset = list(index.offset)
                ops: list[Operation] = []
                res: list[SSAValue] = [body.args[index.dim.value.data]]
                if offset[index.dim.value.data] != 0:
                    ops = [
                        cst := arith.ConstantOp.from_int_and_width(
                            offset[index.dim.value.data], builtin.IndexType()
                        ),
                        add := arith.AddiOp(body.args[index.dim.value.data], cst),
                    ]
                    res = [add.result]
                rewriter.replace(index, ops, res)

        # Get the maybe updated results
        new_results = self.return_targets[op] if op in self.return_targets else []
        # Replace with the loop and necessary constants.
        assert isa(boilerplate_ops, list[Operation])
        rewriter.insert([*boilerplate_ops, p])
        rewriter.replace(op, [], new_results)

@dataclass
class AccessOpToMemRef(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: AccessOp, rewriter: PatternRewriter, /):
        temp = op.temp.type
        assert StencilTypeConstr.verifies(temp)
        assert isinstance(temp.bounds, StencilBoundsAttr)

        memref_offset = op.offset

        mapping = (
            op.offset_mapping
            if op.offset_mapping is not None
            else range(len(memref_offset))
        )

        args = [
            IndexOp(
                properties={
                    "dim": builtin.IntegerAttr.from_index_int_value(i),
                    "offset": IndexAttr.from_indices(
                        *([0] * op.get_apply().get_rank())
                    ),
                },
                result_types=[builtin.IndexType()],
            )
            for i in mapping
        ]

        off_const_ops: list[Operation] = []
        memref_load_args: list[BlockArgument | OpResult] = []

        # This will apply an offset to the index if one is required
        # (e.g the offset is not zero), otherwise will use the index value directly
        for arg, x in zip(args, memref_offset):
            if x != 0:
                constant_op = arith.ConstantOp.from_int_and_width(
                    x, builtin.IndexType()
                )
                add_op = arith.AddiOp(arg, constant_op)
                memref_load_args.append(add_op.results[0])
                off_const_ops += [constant_op, add_op]
            else:
                memref_load_args.append(arg.idx)

        load = memref.LoadOp(
            operands=[op.temp, memref_load_args], result_types=[temp.element_type]
        )

        rewriter.insert(args)
        rewriter.replace(op, [*off_const_ops, load], [load.res])

@dataclass
class StencilStoreToSubview(RewritePattern):
    return_targets: dict[ApplyOp, list[SSAValue | None]]

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: StoreOp, rewriter: PatternRewriter, /):
        for use in op.field.uses:
            if isa(use.operation, LoadOp):
                raise VerifyException(
                    "Cannot lower directly if loading and storing the same field! "
                    "Try running `stencil-bufferize` before."
                )
            if isa(use.operation, StoreOp) and use.operation is not op:
                raise VerifyException(
                    "Cannot lower directly if storing to the same field multiple "
                    "times! Try running `stencil-bufferize` before."
                )
        field = op.field
        assert isa(field.type, FieldType[Attribute])
        assert isa(field.type.bounds, StencilBoundsAttr)
        temp = op.temp
        assert isa(temp.type, TempType[Attribute])
        offsets = [i for i in -field.type.bounds.lb]
        sizes = [i for i in temp.type.get_shape()]
        subview = memref.SubviewOp.from_static_parameters(
            field,
            StencilToMemRefType(field.type),
            offsets,
            sizes,
            [1] * len(sizes),
        )
        name = None
        if subview.source.name_hint:
            name = subview.source.name_hint + "_storeview"
        subview.result.name_hint = name

        # Don't insert the subview here - defer insertion to ReturnOpToMemRef
        # to avoid the subview being eliminated as dead code before it gets users.
        # The subview will be inserted when ReturnOpToMemRef processes the
        # corresponding return operation and needs this target.

        rewriter.erase(op)

        update_return_target(self.return_targets, field, subview.result)