Csl stencil bufferize

class StencilTypeConversion(TypeConversionPattern):
    """
    Converts from tensorised stencil.field types to memref by extracting the element type which is a tensor
    and converting it to memref.

    For instance:
        `!stencil.field<[-1,1023]x[-1,511]xtensor<512xf32>` to `memref<512xf32>`
    """

    @attr_type_rewrite_pattern
    def convert_type(
        self, typ: stencil.FieldType[TensorType[Attribute]]
    ) -> memref.MemRefType:
        # todo should this convert to `memref` or `stencil.field<..xmemref<..>>`?
        return tensor_to_memref_type(typ.get_element_type())

@dataclass(frozen=True)
class ApplyOpBufferize(RewritePattern):
    """
    Bufferizes csl_stencil.apply, rewriting args and block args, changing them from tensor to memref types.
    For each converted arg, creates a `bufferization.to_memref` before the apply op.
    For each converted block arg, creates a `bufferization.to_tensor` at the start of the block.
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: csl_stencil.ApplyOp, rewriter: PatternRewriter, /):
        if isa(op.accumulator.type, memref.MemRefType):
            return

        # convert args
        buf_args: list[SSAValue] = []
        to_memrefs: list[Operation] = [buf_iter_arg := to_buffer_op(op.accumulator)]
        # in case of subsequent apply ops accessing this accumulator, replace uses with `bufferization.to_memref`
        rewriter.replace_uses_with_if(
            op.accumulator,
            buf_iter_arg.memref,
            lambda use: use.operation != buf_iter_arg,
        )
        for arg in [*op.args_rchunk, *op.args_dexchng]:
            if isa(arg.type, TensorType[Attribute]):
                to_memrefs.append(new_arg := to_buffer_op(arg))
                buf_args.append(new_arg.memref)
            else:
                buf_args.append(arg)

        # create new op
        buf_apply_op = csl_stencil.ApplyOp(
            operands=[
                op.field,
                buf_iter_arg.memref,
                op.args_rchunk,
                op.args_dexchng,
                op.dest,
            ],
            result_types=op.res.types or [[]],
            regions=[
                self._get_empty_bufferized_region(op.receive_chunk.block.args),
                self._get_empty_bufferized_region(op.done_exchange.block.args),
            ],
            properties=op.properties,
            attributes=op.attributes,
        )

        # insert to_tensor ops and create arg mappings for block inlining
        chunk_region_arg_mapping: Sequence[SSAValue] = []
        for idx, (old_arg, arg) in enumerate(
            zip(op.receive_chunk.block.args, buf_apply_op.receive_chunk.block.args)
        ):
            # arg0 has special meaning and does not need a `to_tensor` op
            if isinstance(old_arg.type, TensorType) and idx != 0:
                rewriter.insert(
                    # ensure iter_arg is writable
                    t := to_tensor_op(arg, writable=idx == 2),
                    InsertPoint.at_end(buf_apply_op.receive_chunk.block),
                )
                chunk_region_arg_mapping.append(t.tensor)
            else:
                chunk_region_arg_mapping.append(arg)

        done_exchange_arg_mapping: Sequence[SSAValue] = []
        for idx, (old_arg, arg) in enumerate(
            zip(op.done_exchange.block.args, buf_apply_op.done_exchange.block.args)
        ):
            if isinstance(old_arg.type, TensorType):
                rewriter.insert(
                    # ensure iter_arg is writable
                    t := to_tensor_op(arg, writable=idx == 1),
                    InsertPoint.at_end(buf_apply_op.done_exchange.block),
                )
                done_exchange_arg_mapping.append(t.tensor)
            else:
                done_exchange_arg_mapping.append(arg)

        typ = op.receive_chunk.block.args[0].type
        assert isa(typ, TensorType[Attribute])
        chunk_type = TensorType(typ.get_element_type(), typ.get_shape()[1:])

        # inline blocks from old into new regions
        rewriter.inline_block(
            op.receive_chunk.block,
            InsertPoint.at_end(buf_apply_op.receive_chunk.block),
            chunk_region_arg_mapping,
        )

        rewriter.inline_block(
            op.done_exchange.block,
            InsertPoint.at_end(buf_apply_op.done_exchange.block),
            done_exchange_arg_mapping,
        )

        self._inject_iter_arg_into_linalg_outs(
            buf_apply_op, rewriter, chunk_type, chunk_region_arg_mapping[2]
        )

        # insert new op
        rewriter.replace(op, new_ops=[*to_memrefs, buf_apply_op])

    @staticmethod
    def _get_empty_bufferized_region(args: Sequence[BlockArgument]) -> Region:
        """Helper function to create a new region with bufferized arg types."""
        return Region(
            Block(
                arg_types=[
                    (
                        tensor_to_memref_type(arg.type)
                        if AnyTensorTypeConstr.verifies(arg.type)
                        else arg.type
                    )
                    for arg in args
                ]
            )
        )

    @staticmethod
    def _inject_iter_arg_into_linalg_outs(
        op: csl_stencil.ApplyOp,
        rewriter: PatternRewriter,
        chunk_type: TensorType[Attribute],
        iter_arg: SSAValue,
    ):
        """
        Finds a linalg op with `chunk_type` shape in `outs` and injects
        an extracted slice of `iter_arg`. This is a work-around for the
        way bufferization works, causing it to use `iter_arg` as an accumulator
        and avoiding having an extra alloc + memref.copy.
        """
        linalg_op: linalg.abstract_ops.NamedOperation | None = None
        for curr_op in op.receive_chunk.block.ops:
            if (
                isinstance(curr_op, linalg.abstract_ops.NamedOperation)
                and len(curr_op.outputs) > 0
                and curr_op.outputs.types[0] == chunk_type
            ):
                linalg_op = curr_op
                break

        if linalg_op is None:
            return

        rewriter.replace(
            linalg_op,
            [
                extract_slice_op := tensor.ExtractSliceOp(
                    iter_arg,
                    (op.receive_chunk.block.args[1],),
                    chunk_type.get_shape(),
                    (1,),
                    chunk_type,
                ),
                type(linalg_op).build(
                    operands=[linalg_op.inputs, extract_slice_op.results],
                    result_types=linalg_op.result_types,
                    properties=linalg_op.properties,
                    attributes=linalg_op.attributes,
                    regions=[linalg_op.detach_region(r) for r in linalg_op.regions],
                ),
            ],
        )

    @staticmethod
    def _build_extract_slice(
        op: csl_stencil.ApplyOp, to_tensor: bufferization.ToTensorOp, offset: SSAValue
    ) -> tensor.ExtractSliceOp:
        """
        Helper function to create an early tensor.extract_slice in the apply.recv_chunk_cb region needed for bufferization.
        """

        # this is the unbufferized `tensor<(neighbours)x(ZDim)x(type)>` value
        typ = op.receive_chunk.block.args[0].type
        assert isa(typ, TensorType[Attribute])

        return tensor.ExtractSliceOp(
            to_tensor.tensor,
            (offset,),
            typ.get_shape()[1:],
            (1,),
            TensorType(typ.get_element_type(), typ.get_shape()[1:]),
        )

@dataclass(frozen=True)
class AccessOpBufferize(RewritePattern):
    """
    Bufferizes AccessOp.

    The type conversion pass creates the scenario that some `csl_stencil.access` ops are equal input and output types,
    for instance, `(memref<512xf32>) -> memref<512xf32>`. This only happens for ops accessing own data. In this case,
    the access op has no effect and can safely be folded away.
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: csl_stencil.AccessOp, rewriter: PatternRewriter, /):
        if not isa(op.result.type, TensorType[Attribute]):
            return
        r_type = tensor_to_memref_type(op.result.type)

        # accesses to own data that (after bufferization) have the same input and output type can be safely folded away
        if op.op.type == r_type and all(o == 0 for o in op.offset):
            rewriter.replace(op, to_tensor_op(op.op))
            return

        # accesses to buffers passed in additional args can read directly from memref underlying `to_tensor`
        source = (
            op.op.op.memref
            if isinstance(op.op, OpResult)
            and isinstance(op.op.op, bufferization.ToTensorOp)
            else op.op
        )

        rewriter.replace(
            op,
            [
                access := csl_stencil.AccessOp(
                    source,
                    op.offset,
                    r_type,
                    op.offset_mapping,
                ),
                to_tensor_op(access.result),
            ],
        )

@dataclass(frozen=True)
class YieldOpBufferize(RewritePattern):
    """Bufferizes YieldOp."""

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: csl_stencil.YieldOp, rewriter: PatternRewriter, /):
        to_memrefs: list[Operation] = []
        args: list[SSAValue] = []
        for arg in op.arguments:
            if isa(arg.type, TensorType[Attribute]):
                to_memrefs.append(new_arg := to_buffer_op(arg))
                args.append(new_arg.memref)
            else:
                args.append(arg)

        if len(to_memrefs) == 0:
            return

        rewriter.replace(op, [*to_memrefs, csl_stencil.YieldOp(*args)])

@dataclass(frozen=True)
class FuncOpBufferize(RewritePattern):
    """
    Replace the function_type and let a separate type conversion pass handle the block args.
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: func.FuncOp, rewriter: PatternRewriter, /):
        function_type = FunctionType.from_lists(
            [
                (
                    tensor_to_memref_type(t.get_element_type())
                    if isa(t, stencil.FieldType[TensorType[Attribute]])
                    else t
                )
                for t in op.function_type.inputs
            ],
            [
                (
                    tensor_to_memref_type(t.get_element_type())
                    if isa(t, stencil.FieldType[TensorType[Attribute]])
                    else t
                )
                for t in op.function_type.outputs
            ],
        )
        if function_type == op.function_type:
            return
        rewriter.replace(
            op,
            func.FuncOp.build(
                operands=op.operands,
                result_types=op.result_types,
                regions=[op.detach_region(op.body)],
                properties={**op.properties, "function_type": function_type},
                attributes=op.attributes.copy(),
            ),
        )

@dataclass(frozen=True)
class ArithConstBufferize(RewritePattern):
    """
    Bufferize arith tensor constants to prevent mlir bufferize from promoting them to globals.
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: arith.ConstantOp, rewriter: PatternRewriter, /):
        if not isa(op.result.type, TensorType[Attribute]):
            return
        assert isa(op.value, DenseIntOrFPElementsAttr)
        assert isa(op.value.type, TensorType[AnyDenseElement])
        typ = DenseIntOrFPElementsAttr(
            tensor_to_memref_type(op.value.type), op.value.data
        )
        rewriter.replace(
            op,
            [
                c := arith.ConstantOp(typ),
                to_tensor_op(c.result),
            ],
        )

@dataclass(frozen=True)
class InjectApplyOutsIntoLinalgOuts(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: csl_stencil.ApplyOp, rewriter: PatternRewriter, /):
        # require bufferized apply (with op.dest specified)
        # zero-output apply ops may be used for communicate-only, to which this pattern does not apply
        if not op.dest:
            return

        yld = op.done_exchange.block.last_op
        assert isinstance(yld, csl_stencil.YieldOp)
        new_dest: list[SSAValue] = []
        new_yield_args: list[SSAValue] = []
        additional_args: list[SSAValue] = []
        to_remove: list[Operation] = []

        for arg, yld_arg in zip(op.dest, yld.arguments, strict=True):
            if (
                not isinstance(yld_arg, OpResult)
                or not isinstance(yld_arg.op, bufferization.ToBufferOp)
                or not isinstance(yld_arg.op.tensor, OpResult)
                or not isinstance(
                    linalg_op := yld_arg.op.tensor.op,
                    linalg.abstract_ops.NamedOperation | linalg.ops.GenericOp,
                )
                or not isa(arg_t := arg.type, MemRefType)
                or not isa(yld_arg.type, MemRefType)
            ):
                new_dest.append(arg)
                new_yield_args.append(yld_arg)
                continue
            additional_args.append(arg)
            if yld_arg.has_one_use():
                to_remove.append(yld_arg.op)

            arg = op.done_exchange.block.insert_arg(
                arg.type, len(op.done_exchange.block.args)
            )
            arg_to_tensor = to_tensor_op(arg, writable=True)

            # offset of core data, assuming symmetric ghost cells in each direction
            offsets = tuple(
                (src - dst) // 2  # symmetric offset
                for src, dst in zip(
                    arg_t.get_shape(), yld_arg.type.get_shape(), strict=True
                )
            )

            extract_slice_op = tensor.ExtractSliceOp(
                arg_to_tensor,
                offsets,
                yld_arg.type.get_shape(),
                (1,),
                yld_arg.op.tensor.type,
            )
            rewriter.insert(
                [arg_to_tensor, extract_slice_op], InsertPoint.before(linalg_op)
            )
            rewriter.replace(
                linalg_op,
                type(linalg_op).build(
                    operands=[linalg_op.inputs, [extract_slice_op.result]],
                    result_types=linalg_op.result_types,
                    regions=[linalg_op.detach_region(r) for r in linalg_op.regions],
                    properties=linalg_op.properties,
                    attributes=linalg_op.attributes,
                ),
            )
        if additional_args:
            rewriter.replace(yld, csl_stencil.YieldOp(*new_yield_args))
            for r in to_remove:
                rewriter.erase(r)
            rewriter.replace(
                op,
                csl_stencil.ApplyOp(
                    operands=[
                        op.field,
                        op.accumulator,
                        [*op.args_rchunk],
                        [*op.args_dexchng, *additional_args],
                        [*op.dest],
                    ],
                    result_types=op.res.types or [[]],
                    regions=[op.detach_region(r) for r in op.regions],
                    properties=op.properties,
                    attributes=op.attributes,
                ),
            )

@dataclass(frozen=True)
class ReselectLinalgOutsFromInputs(RewritePattern):
    """
    Reselects the DPS `outs` of a linalg op if it is one of its inputs.
      * select `writable` tensor input with no later use, or else
      * select linalg op input with no later use
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(
        self,
        op: linalg.abstract_ops.NamedOperation | linalg.ops.GenericOp,
        rewriter: PatternRewriter,
        /,
    ):
        # only apply rewrite when re-selecting `outs` from `ins`
        if (
            op.outputs[0] not in op.inputs
            or self.is_writable(op.outputs[0])
            or len(op.outputs) != 1
        ):
            return

        # the new `outs` to re-select
        out: SSAValue | None = None

        for arg in op.inputs:
            # reselect outs that has no later use to avoid read-after-write conflicts
            if arg.has_one_use():
                # check for a `writable` input with no later uses and break immediately
                if self.is_writable(arg):
                    out = arg
                    break

                # check for a linalg op input with no later uses and keep looking
                if isinstance(arg, OpResult) and isinstance(
                    arg.op, linalg.abstract_ops.NamedOperation | linalg.ops.GenericOp
                ):
                    out = arg

        # replace the op with `out` as `output[0]`
        if out:
            rewriter.replace(
                op,
                type(op).build(
                    operands=[op.inputs, [out]],
                    result_types=op.result_types,
                    regions=[op.detach_region(r) for r in op.regions],
                    properties=op.properties,
                    attributes=op.attributes,
                ),
            )

    @staticmethod
    def is_writable(val: SSAValue) -> bool:
        """Returns if `val` is a `writable` tensor."""
        return (
            isinstance(val, OpResult)
            and isinstance(val.op, bufferization.ToTensorOp)
            and val.op.writable is not None
        )

@dataclass(frozen=True)
class CslStencilBufferize(ModulePass):
    """
    Bufferizes the csl_stencil dialect.

    Attempts to inject `csl_stencil.apply.recv_chunk_cb.accumulator` into linalg compute ops `outs` within that region
    for improved bufferization. Ideally be run after `--lift-arith-to-linalg`.

    In preparation for bufferization with minimal overhead, linalg ops `outs` are set as follows:
      if a linalg op's destination is one of its inputs
      1. prefer a `writable` input with no other uses
      2. prefer a linalg op input with no other uses
    """

    name = "csl-stencil-bufferize"

    def apply(self, ctx: Context, op: ModuleOp) -> None:
        module_pass = PatternRewriteWalker(
            GreedyRewritePatternApplier(
                [
                    StencilTypeConversion(),
                    ApplyOpBufferize(),
                    AccessOpBufferize(),
                    YieldOpBufferize(),
                    FuncOpBufferize(),
                    ArithConstBufferize(),
                ]
            )
        )
        module_pass.rewrite_module(op)
        PatternRewriteWalker(
            GreedyRewritePatternApplier(
                [
                    InjectApplyOutsIntoLinalgOuts(),
                    ReselectLinalgOutsFromInputs(),
                ]
            ),
            apply_recursively=False,
        ).rewrite_module(op)