Stencil tensorize z dimension

class StencilTypeConversion(TypeConversionPattern):
    @attr_type_rewrite_pattern
    def convert_type(self, typ: FieldType[Attribute]) -> FieldType[Attribute]:
        return stencil_field_to_tensor(typ)

@dataclass(frozen=True)
class AccessOpTensorize(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: AccessOp, rewriter: PatternRewriter, /):
        if not is_tensorized(op.operands[0].type) or len(op.offset) != 3:
            return
        xy_offsets, z_offset = (
            tuple(o for o in op.offset)[:-1],
            tuple(o for o in op.offset)[-1],
        )
        a = AccessOp(op.temp, xy_offsets)
        # this conditional controls if ExtractSliceOps for x/y accesses should be generated
        # if xy_offsets[0] != 0 or xy_offsets[1] != 0:
        #     rewriter.replace(op, a)
        #     return
        assert isa(op.temp.type, TempType[Attribute])
        element_t = op.temp.type.get_element_type()
        assert is_tensor(element_t)
        extract = ExtractSliceOp.from_static_parameters(
            a, [z_offset], element_t.get_shape()
        )
        rewriter.insert(a, InsertPoint.before(op))
        rewriter.replace(op, extract)

class ArithOpTensorize(RewritePattern):
    """
    Tensorises arith binary ops.
    If both operands are tensor types, rebuilds the op with matching result type.
    If one operand is scalar and an `arith.constant`, create a tensor constant directly.
    If one operand is scalar and not an `arith.constant`, create an empty tensor and fill it with the scalar value.
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(
        self, op: FloatingPointLikeBinaryOperation, rewriter: PatternRewriter, /
    ):
        type_constructor = type(op)
        if is_tensor(op.result.type):
            return
        if is_tensor(op.lhs.type) and is_tensor(op.rhs.type):
            rewriter.replace(
                op,
                type_constructor(op.lhs, op.rhs, flags=None, result_type=op.lhs.type),
            )
        elif isa(op.lhs.type, TensorType[AnyFloat]) and is_scalar(op.rhs.type):
            new_rhs = ArithOpTensorize._rewrite_scalar_operand(
                op.rhs, op.lhs.type, op, rewriter
            )
            rewriter.replace(
                op,
                type_constructor(op.lhs, new_rhs, flags=None, result_type=op.lhs.type),
            )
        elif is_scalar(op.lhs.type) and isa(op.rhs.type, TensorType[AnyFloat]):
            new_lhs = ArithOpTensorize._rewrite_scalar_operand(
                op.lhs, op.rhs.type, op, rewriter
            )
            rewriter.replace(
                op,
                type_constructor(new_lhs, op.rhs, flags=None, result_type=op.rhs.type),
            )

    @staticmethod
    def _rewrite_scalar_operand(
        scalar_op: SSAValue,
        dest_typ: TensorType[AnyFloat],
        op: FloatingPointLikeBinaryOperation,
        rewriter: PatternRewriter,
    ) -> SSAValue:
        """
        Rewrites a scalar operand into a tensor.
        If it is a constant, create a corresponding tensor constant.
        If it is not a constant, create an empty tensor and `linalg.fill` it with the scalar value.
        """
        if isinstance(scalar_op, OpResult) and isinstance(scalar_op.op, ConstantOp):
            float_attr = scalar_op.op.value
            assert isinstance(float_attr, FloatAttr)
            scalar_value = float_attr.value.data
            tens_const = ConstantOp(
                DenseIntOrFPElementsAttr.from_list(dest_typ, [scalar_value])
            )
            rewriter.insert(tens_const, InsertPoint.before(scalar_op.op))
            return tens_const.result
        emptyop = EmptyOp((), dest_typ)
        fillop = FillOp((scalar_op,), (emptyop.tensor,), (dest_typ,))
        rewriter.insert(emptyop, InsertPoint.before(op))
        rewriter.insert(fillop, InsertPoint.before(op))
        return fillop.res[0]

@dataclass(frozen=True)
class ApplyOpTensorize(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: ApplyOp, rewriter: PatternRewriter, /):
        if all(is_tensorized(arg.type) for arg in op.args):
            access_patterns = dict[Operand, AccessPattern](
                zip(op.region.block.args, op.get_accesses())
            )
            for access_op in op.region.walk():
                if isinstance(access_op, AccessOp):
                    z_shift = -access_patterns[access_op.temp].halo_in_axis(2)[0]
                    access_op.offset = IndexAttr.from_indices(
                        *[idx.data for idx in access_op.offset.array.data[:-1]],
                        access_op.offset.array.data[-1].data + z_shift,
                    )

            body = Block(arg_types=op.operand_types)
            rewriter.inline_block(
                op.region.block, InsertPoint.at_start(body), body.args
            )

            rewriter.replace(
                op,
                ApplyOp(
                    op.args,
                    body,
                    [stencil_temp_to_tensor(r.type) for r in op.res],
                ),
            )

class FuncOpTensorize(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: FuncOp, rewriter: PatternRewriter, /):
        if not op.is_declaration:
            for arg in op.args:
                if isa(arg.type, FieldType[Attribute]):
                    op.replace_argument_type(
                        arg, stencil_field_to_tensor(arg.type), rewriter
                    )

class LoadOpTensorize(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: LoadOp, rewriter: PatternRewriter, /):
        assert isa(op.res.type, TempType[Attribute])
        bounds = op.res.type.bounds
        assert isinstance(bounds, StencilBoundsAttr)
        rewriter.replace(
            op,
            LoadOp(
                op.field,
                IndexAttr.from_indices(*[lb for lb in bounds.lb][:-1]),
                IndexAttr.from_indices(*[ub for ub in bounds.ub][:-1]),
            ),
        )

class DmpSwapOpTensorize(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: dmp.SwapOp, rewriter: PatternRewriter, /):
        if (
            is_tensorized(op.input_stencil.type)
            and op.swapped_values
            and not is_tensorized(op.swapped_values.type)
        ):
            rewriter.replace(
                op,
                dmp.SwapOp.get(op.input_stencil, op.strategy, ArrayAttr(op.swaps.data)),
            )

class StoreOpTensorize(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: StoreOp, rewriter: PatternRewriter, /):
        if (
            is_tensorized(op.field.type)
            and isinstance(op.field.type, ShapedType)
            and len(op.bounds.lb) != len(op.field.type.get_shape())
        ):
            rewriter.replace(
                op,
                StoreOp(
                    op.temp,
                    op.field,
                    StencilBoundsAttr(
                        zip(list(op.bounds.lb.array)[:-1], list(op.bounds.ub)[:-1])
                    ),
                ),
            )

class AccessOpUpdateShape(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: AccessOp, rewriter: PatternRewriter, /):
        if typ := get_required_result_type(op):
            if needs_update_shape(op.res.type, typ):
                rewriter.replace(
                    op,
                    AccessOp.build(
                        operands=[op.temp], attributes=op.attributes, result_types=[typ]
                    ),
                )

class CslStencilAccessOpUpdateShape(RewritePattern):
    """
    Updates the result type of a tensorized `csl_stencil.access` op
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: csl_stencil.AccessOp, rewriter: PatternRewriter, /):
        if typ := get_required_result_type(op):
            if needs_update_shape(op.result.type, typ) and (
                isa(op.op.type, TempType[TensorType[Attribute]])
                or isa(op.op.type, TensorType[Attribute])
            ):
                rewriter.replace(
                    op,
                    csl_stencil.AccessOp(
                        op.op,
                        op.offset,
                        (
                            op.op.type.get_element_type()
                            if isa(op.op.type, TempType[TensorType[Attribute]])
                            else typ
                        ),
                        op.offset_mapping,
                    ),
                )

class ExtractSliceOpUpdateShape(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: ExtractSliceOp, rewriter: PatternRewriter, /):
        if typ := get_required_result_type(op):
            if needs_update_shape(op.result.type, typ):
                rewriter.replace(
                    op,
                    ExtractSliceOp.from_static_parameters(
                        op.source, op.static_offsets.get_values(), typ.get_shape()
                    ),
                )

class ArithOpUpdateShape(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(
        self, op: FloatingPointLikeBinaryOperation, rewriter: PatternRewriter, /
    ):
        arithBinaryOpUpdateShape(op, rewriter)

class VarithOpUpdateShape(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: varith.VarithOp, rewriter: PatternRewriter, /):
        type_constructor = type(op)
        if typ := get_required_result_type(op):
            if needs_update_shape(op.result_types[0], typ):
                rewriter.replace(
                    op, type_constructor.build(operands=[op.args], result_types=[typ])
                )

class EmptyOpUpdateShape(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: EmptyOp, rewriter: PatternRewriter, /):
        if typ := get_required_result_type(op):
            if needs_update_shape(op.results[0].type, typ):
                rewriter.replace(op, EmptyOp((), typ))

class FillOpUpdateShape(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: FillOp, rewriter: PatternRewriter, /):
        if typ := get_required_result_type(op):
            if needs_update_shape(op.results[0].type, typ):
                rewriter.replace(
                    op, FillOp(op.inputs, op.outputs, [typ] * len(op.outputs))
                )

class ConstOpUpdateShape(RewritePattern):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: ConstantOp, rewriter: PatternRewriter, /):
        if is_tensor(op.result.type):
            if typ := get_required_result_type(op):
                if needs_update_shape(op.result.type, typ):
                    assert isinstance(op.value, DenseIntOrFPElementsAttr)
                    rewriter.replace(
                        op, ConstantOp(DenseIntOrFPElementsAttr(typ, op.value.data))
                    )

@dataclass(frozen=True)
class BackpropagateStencilShapes(ModulePass):
    """
    Greedily back-propagates the result types of tensorized ops.
    Use after creating/modifying tensorization.
    """

    name = "backpropagate-stencil-shapes"

    def apply(self, ctx: Context, op: builtin.ModuleOp) -> None:
        backpropagate_stencil_shapes = PatternRewriteWalker(
            GreedyRewritePatternApplier(
                [
                    CslStencilAccessOpUpdateShape(),
                    ExtractSliceOpUpdateShape(),
                    EmptyOpUpdateShape(),
                    FillOpUpdateShape(),
                    ArithOpUpdateShape(),
                    VarithOpUpdateShape(),
                    ConstOpUpdateShape(),
                ]
            ),
            walk_reverse=True,
            apply_recursively=False,
        )
        backpropagate_stencil_shapes.rewrite_module(op)

@dataclass(frozen=True)
class StencilTensorizeZDimension(ModulePass):
    name = "stencil-tensorize-z-dimension"

    def apply(self, ctx: Context, op: ModuleOp) -> None:
        module_pass = PatternRewriteWalker(
            GreedyRewritePatternApplier(
                [
                    FuncOpTensorize(),
                    StencilTypeConversion(),  # this needs to come after FuncOpTensorize()
                    LoadOpTensorize(),
                    ApplyOpTensorize(),
                    StoreOpTensorize(),
                    DmpSwapOpTensorize(),
                    # AccessOpTensorize(),   # this doesn't work here, using second pass
                ]
            ),
            walk_reverse=False,
            apply_recursively=False,
        )
        module_pass.rewrite_module(op)
        stencil_pass = PatternRewriteWalker(
            GreedyRewritePatternApplier(
                [
                    AccessOpTensorize(),
                    ArithOpTensorize(),
                ]
            ),
            walk_reverse=False,
            apply_recursively=False,
        )
        stencil_pass.rewrite_module(op)
        BackpropagateStencilShapes().apply(ctx=ctx, op=op)