# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name,missing-function-docstring,unused-import
# ruff: noqa: E501, F401
"""Intrinsics for ARM tensorization."""

from tvm import tirx
from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder.tirx import prim_func as build_prim_func
from tvm.target.codegen import llvm_version_major

from .. import TensorIntrin
from .dot_product_common import (
    DP4A_S8S8S32_INTRIN,
    DP4A_S8U8S32_INTRIN,
    DP4A_U8S8S32_INTRIN,
    DP4A_U8U8U32_INTRIN,
)

# TODO(masahi): Parametrize the TVMScript description of dot product by
# shape and dtype, and share the common description with x86.


@T.prim_func(s_tir=True)
def neon_4x4_i8i8i32_desc(
    A: T.Buffer((4,), "int8", offset_factor=1),
    B: T.Buffer((4, 4), "int8", offset_factor=1),
    C: T.Buffer((4,), "int32", offset_factor=1),
) -> None:
    with T.sblock("root"):
        T.reads(C[0:4], A[0:4], B[0:4, 0:4])
        T.writes(C[0:4])
        for i in T.serial(0, 4):
            for k in T.serial(0, 4):
                with T.sblock("update"):
                    vi, vk = T.axis.remap("SR", [i, k])
                    C[vi] = C[vi] + T.cast(A[vk], "int32") * T.cast(B[vi, vk], "int32")


@T.prim_func(s_tir=True)
def neon_4x4_i8i8i32_impl(
    A: T.Buffer((4,), "int8", offset_factor=1),
    B: T.Buffer((4, 4), "int8", offset_factor=1),
    C: T.Buffer((4,), "int32", offset_factor=1),
) -> None:
    with T.sblock("root"):
        T.reads(C[0:4], A[0:4], B[0:4, 0:4])
        T.writes(C[0:4])

        A_int8 = A.vload([0], "int8x4")
        re_int32 = T.reinterpret(A_int8, dtype="int32")
        vec_ai32 = T.broadcast(re_int32, 2)
        vec_a = T.reinterpret(vec_ai32, dtype="int8x8")

        vec_b = B.vload([0, 0], dtype="int8x16")

        # TODO(masahi): Remove duplication when inlined function call is supported
        vec_b_low = T.vectorlow(vec_b, dtype="int8x8")

        multiply_low = T.call_llvm_pure_intrin(
            T.llvm_lookup_intrinsic_id("llvm.aarch64.neon.smull.v8i16"),
            vec_a,
            vec_b_low,
            dtype="int16x8",
        )

        pairwise_reduction_low = T.call_llvm_pure_intrin(
            T.llvm_lookup_intrinsic_id("llvm.aarch64.neon.saddlp.v4i32.v8i16"),
            multiply_low,
            dtype="int32x4",
        )

        vec_b_high = T.vectorhigh(vec_b, dtype="int8x8")

        multiply_high = T.call_llvm_pure_intrin(
            T.llvm_lookup_intrinsic_id("llvm.aarch64.neon.smull.v8i16"),
            vec_a,
            vec_b_high,
            dtype="int16x8",
        )

        pairwise_reduction_high = T.call_llvm_pure_intrin(
            T.llvm_lookup_intrinsic_id("llvm.aarch64.neon.saddlp.v4i32.v8i16"),
            multiply_high,
            dtype="int32x4",
        )

        C[T.ramp(T.int32(0), 1, 4)] += T.call_llvm_pure_intrin(
            T.llvm_lookup_intrinsic_id("llvm.aarch64.neon.addp.v4i32"),
            pairwise_reduction_low,
            pairwise_reduction_high,
            dtype="int32x4",
        )


def get_dotprod_intrin(in_dtype, out_dtype):
    if in_dtype == "uint8":
        instr = "udot.v4u32.v16u8"
    else:  # if in_dtype == "int8"
        instr = "sdot.v4i32.v16i8"

    in_dtype_x4 = f"{in_dtype}x4"
    out_dtype_x4 = f"{out_dtype}x4"
    in_dtype_x16 = f"{in_dtype}x16"

    @T.prim_func(s_tir=True)
    def dot_prod_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
        A = T.match_buffer(a, (4,), dtype=in_dtype, offset_factor=1)
        B = T.match_buffer(b, (4, 4), dtype=in_dtype, offset_factor=1)
        C = T.match_buffer(c, (4,), dtype=out_dtype, offset_factor=1)
        with T.sblock("root"):
            T.reads(C[0:4], A[0:4], B[0:4, 0:4])
            T.writes(C[0:4])
            for i in T.serial(0, 4):
                for k in T.serial(0, 4):
                    with T.sblock("update"):
                        vi, vk = T.axis.remap("SR", [i, k])
                        C[vi] = C[vi] + T.cast(A[vk], dtype=out_dtype) * T.cast(
                            B[vi, vk], dtype=out_dtype
                        )

    @T.prim_func(s_tir=True)
    def dot_prod_impl(a: T.handle, b: T.handle, c: T.handle) -> None:
        A = T.match_buffer(a, (4,), dtype=in_dtype, offset_factor=1)
        B = T.match_buffer(b, (4, 4), dtype=in_dtype, offset_factor=1)
        C = T.match_buffer(c, (4,), dtype=out_dtype, offset_factor=1)
        with T.sblock("root"):
            T.reads(C[0:4], A[0:4], B[0:4, 0:4])
            T.writes(C[0:4])

            A_i8x4 = A.vload([0], in_dtype_x4)
            A_i32 = T.reinterpret(A_i8x4, dtype=out_dtype)
            vec_ai32 = T.broadcast(A_i32, 4)
            vec_a = T.reinterpret(vec_ai32, dtype=in_dtype_x16)

            vec_b = B.vload([0, 0], dtype=in_dtype_x16)

            vec_c = C.vload([0], dtype=out_dtype_x4)

            C[T.ramp(T.int32(0), 1, 4)] = T.call_llvm_pure_intrin(
                T.llvm_lookup_intrinsic_id(f"llvm.aarch64.neon.{instr}"),
                vec_c,
                vec_a,
                vec_b,
                dtype=out_dtype_x4,
            )

    return dot_prod_desc, dot_prod_impl


def _create_ptrue_mask(dtype):
    """
    Creates a mask that enables all lanes of a scalable vector.
    """
    return T.broadcast(T.bool(True), tirx.get_vscale_expr(dtype))


def _create_active_lane_mask(tensor, relative_offsets, vertical_limit):
    """
    Get the active lane mask intrinsic call for predicated accesses.

    Parameters
    ----------
    tensor : tvm.tirx.Buffer
        The tensor the buffer access will be performed on.
    relative_offsets : Tuple[Expr, Expr]
        The vertical and horizontal offsets into the accumulator tile.
    vertical_limit : Expr
        An absolute offset specifying the limit at which rows should be stored.

    Returns
    -------
    Expr
        The active lane mask intrinsic.
    """
    vertical_offset, horizontal_offset = relative_offsets
    stride = tensor.strides[0]

    # The base is the offset of the first value we wish to store
    base = T.int32(tensor.offset_of([vertical_offset, horizontal_offset])[0])

    # The limit is the maximum offset in the current row of 'base' that we wish to allow values
    # to be stored. Calculating this limit is a bit tricky since we can only request offsets of
    # elements in the tensorized tile of the output tensor. One way to calculate this is to find
    # the offset of the first value in the row of the output tensor that 'base' is in and add
    # 'stride' to it.
    limit = (
        base
        - T.int32(horizontal_offset)
        - T.int32(tensor.offset_of([0, 0])[0] % stride)
        + T.int32(stride)
    )
    limit = T.Min(limit, T.Cast("int32", vertical_limit) * stride)

    return T.get_active_lane_mask(
        "uint1xvscalex4",
        T.Cast("int32", base),
        T.Cast("int32", limit),
    )


def get_sme_transpose_interleave_2svlx2svl_fp32_intrin(cols, rows):
    """
    Transpose a matrix of size 2SVL x 2SVL (where 'SVL' is the Scalable Vector Length) using
    the Scalable Matrix Extension (SME).

    This is completed by loading rows of the input matrix into the accumulator tile,
    then storing the columns. The SME accumulator tile is divided into a series of sub-tiles
    which must be loaded to / stored from independently.

    Example
    -------
    An example case for float32. In this instance the accumulator tile is divided into 4
    sub-tiles of size SVLxSVL numbered 0-3. We start by loading rows of A, each SVL in length,
    into each of the sub-tiles. In the diagram below, each load for a sub-tile is sequenced by
    a, b, ... till the tile is full.

    The columns of each sub-tile are then stored into A_t. Note that to perform a transpose,
    the contents of sub-tile 1 and 2 are stored in opposite locations - see the diagram
    below.

    ::

        A:                                  Accumulator tile:                     A_t:
                    2SVL                                2SVL                               2SVL
             +----------------+                 +-----------------+                +-------------------+
             | --0a--  --1a-- |                 |                 |                | |  |     |  |     |
             | --0b--  --1b-- |                 |    0       1    |                | 0a 0b .. 2a 2b .. |
             |   ...     ...  | ld1w.horiz      |                 | st1w.vert      | |  |     |  |     |
        2SVL | --2a--  --3a-- |   ====>    2SVL |                 |   ====>   2SVL | |  |     |  |     |
             | --2a--  --3b-- |                 |    2       3    |                | 1a 1b .. 3a 3b .. |
             |   ...     ...  |                 |                 |                | |  |     |  |     |
             +----------------+                 +-----------------+                +-------------------+

    Returns
    -------
    intrin : TensorIntrin
        The SME TensorIntrin that can be used in tensorizing a schedule.

    """
    SVF = tirx.get_vscale_expr("float32")
    SVF2 = 2 * SVF

    @T.prim_func(s_tir=True)
    def desc(a: T.handle, a_t: T.handle) -> None:
        A = T.match_buffer(a, (SVF2, SVF2), dtype="float32", offset_factor=1)
        A_t = T.match_buffer(a_t, (SVF2, SVF2), dtype="float32", offset_factor=1)
        with T.sblock("root"):
            T.reads(A[0:SVF2, 0:SVF2])
            T.writes(A_t[0:SVF2, 0:SVF2])
            for k, m in T.grid(SVF2, SVF2):
                with T.sblock("transpose"):
                    v_m, v_k = T.axis.remap("SS", [m, k])
                    A_t[v_k, v_m] = A[v_m, v_k]

    def impl():
        sub_tile_count = 4

        with IRBuilder() as ib:
            with build_prim_func():
                a = T.arg("a", T.handle())
                a_t = T.arg("a_t", T.handle())

                A = T.match_buffer(
                    a, (SVF2, SVF2), "float32", offset_factor=1, strides=[T.int32(), 1]
                )
                A_t = T.match_buffer(
                    a_t,
                    (SVF2, SVF2),
                    "float32",
                    offset_factor=1,
                    strides=[T.int32(), 1],
                )

                with T.sblock("root"):
                    T.reads(A[0:SVF2, 0:SVF2])
                    T.writes(A_t[0:SVF2, 0:SVF2])

                    # Load rows of the input matrix
                    with T.serial(0, SVF) as slice_idx:
                        for sub_tile_idx in range(0, sub_tile_count):
                            row_offset = SVF if sub_tile_idx >= (sub_tile_count // 2) else 0
                            col_offset = SVF if sub_tile_idx % 2 else 0
                            offset = (slice_idx + row_offset) * A.strides[0] + col_offset

                            input_ptr = A.access_ptr("r", offset=offset)
                            sub_tile = T.int32(sub_tile_idx)
                            predicate = _create_active_lane_mask(
                                A, (row_offset + slice_idx, col_offset), cols
                            )
                            T.evaluate(
                                T.call_llvm_intrin(
                                    "void",
                                    "llvm.aarch64.sme.ld1w.horiz",
                                    predicate,
                                    input_ptr,
                                    sub_tile,
                                    slice_idx,
                                )
                            )

                    # Store columns to the output matrix
                    with T.serial(0, SVF) as slice_idx:
                        for sub_tile_idx in range(0, sub_tile_count):
                            col_offset = SVF if sub_tile_idx >= (sub_tile_count // 2) else 0
                            row_offset = SVF if sub_tile_idx % 2 else 0
                            offset = (slice_idx + row_offset) * A_t.strides[0] + col_offset

                            output_ptr = A_t.access_ptr("w", offset=offset)
                            sub_tile = T.int32(sub_tile_idx)
                            predicate = _create_active_lane_mask(
                                A_t, (row_offset + slice_idx, col_offset), rows
                            )
                            T.evaluate(
                                T.call_llvm_intrin(
                                    "void",
                                    "llvm.aarch64.sme.st1w.vert",
                                    predicate,
                                    output_ptr,
                                    sub_tile,
                                    slice_idx,
                                )
                            )

        return ib.get()

    return desc, impl()


def get_sme_transpose_interleave_block2_2svl_fp16_intrin():
    # pylint: disable=line-too-long
    """
    Transpose and block pack a matrix of size 2SVL x 1SVL (where 'SVL' is the Scalable Vector
    Length for the fp16 datatype) using the Scalable Matrix Extension (SME).

    Rows of the fp16 input matrix are loaded into the accumulator tile and columns are stored
    as fp32 SVL length vectors to the output matrix. When loading, the accumulator tile is
    interpreted to be of shape 2 * 8 * vscale x 8 * vscale. When  storing, we interpret the
    accumulator tile to be of shape 2 * 4 * vscale x 2 * 4 * vscale.

    Example
    -------
    In the fp16 instance, the accumulator tile consists of two sub-tiles numbered 0-1. Rows
    of A are loaded onto the accumulator tile by interleaving rows in the first half (0, SVL//2]
    of the tile and rows in the second half (SVL//2, SVL]. Columns of fp32 values are stored
    into the output buffer. The fp32 store is used to group pairs of consecutive values together,
    resulting in the arrangement displayed below::

        A:                                Accumulator tile:
             +----------------+            +----------------+
             |-------0a-------|            |-------0a-------|
             |-------0b-------|            |-------0x-------|
             |      ...       |            |-------0b-------|            A_t:
             |-------0x-------|            |-------0y-------|             +------------------------------------------------+
             |-------0y-------|            |      ...       |             |0a.0 0a.1 0b.0 0b.1    | 1a.0 1a.1 1b.0 1b.1    |
             |      ...       | ld1h.horiz |                |  st1w.vert  |0x.0 0x.1 0y.0 0y.1    | 1x.0 1x.1 1y.0 1y.1    |
             |================|   ====>    |================|    ====>    |0a.2 0a.3 0b.2 0b.3 ...| 1a.2 1a.3 1b.2 1b.3 ...|
             |-------1a-------|            |-------1a-------|             |0x.2 0x.3 0y.2 0y.3    | 1x.2 1x.3 1y.2 1y.3    |
             |-------1b-------|            |-------1x-------|             |...  ...  ...  ...     | ...  ...  ...  ...     |
             |      ...       |            |-------1b-------|             +------------------------------------------------+
             |-------1x-------|            |-------1y-------|
             |-------1y-------|            |      ...       |
             |      ...       |            |                |
             +----------------+            +----------------+

    In the A_t output matrix in the diagram above, .x is used to denote the offset into the
    labelled row.

    Returns
    -------
    intrin : TensorIntrin
        The SME TensorIntrin that can be used in tensorizing a schedule.

    """
    # pylint: enable=line-too-long
    SVF = tirx.get_vscale_expr("float16")
    SVF2 = 2 * SVF

    @T.prim_func(s_tir=True)
    def desc(a: T.handle, a_t: T.handle) -> None:
        A = T.match_buffer(a, (SVF2, SVF), dtype="float16", offset_factor=1)
        A_t = T.match_buffer(a_t, (SVF, SVF2), dtype="float16", offset_factor=1)
        with T.sblock("root"):
            T.reads(A[0:SVF2, 0:SVF])
            T.writes(A_t[0:SVF, 0:SVF2])
            for k, m in T.grid(SVF, SVF2):
                with T.sblock("transpose"):
                    v_m, v_k = T.axis.remap("SS", [m, k])
                    A_t[v_k, v_m] = A[v_m, v_k]

    def impl():
        with IRBuilder() as ib:
            with build_prim_func():
                a = T.arg("a", T.handle())
                a_t = T.arg("a_t", T.handle())

                A = T.match_buffer(
                    a, (SVF2, SVF), "float16", offset_factor=1, strides=[T.int32(), 1]
                )
                A_t = T.match_buffer(
                    a_t, (SVF, SVF2), "float16", offset_factor=1, strides=[T.int32(), 1]
                )

                ptrue_fp16 = _create_ptrue_mask("float16")
                ptrue_fp32 = _create_ptrue_mask("float32")

                with T.sblock("root"):
                    T.reads(A[0:SVF2, 0:SVF])
                    T.writes(A_t[0:SVF, 0:SVF2])

                    # Load rows of the input matrix
                    with T.serial(SVF // 2) as slice_idx:
                        for sub_tile_idx in range(2):
                            offset = slice_idx * A.strides[0] + (SVF * A.strides[0] * sub_tile_idx)
                            input_ptr = A.access_ptr("r", offset=offset)
                            T.evaluate(
                                T.call_llvm_intrin(
                                    "void",
                                    "llvm.aarch64.sme.ld1h.horiz",
                                    ptrue_fp16,
                                    input_ptr,
                                    sub_tile_idx,
                                    slice_idx * 2,
                                )
                            )
                            input_ptr = A.access_ptr("r", offset=offset + (SVF // 2) * A.strides[0])
                            T.evaluate(
                                T.call_llvm_intrin(
                                    "void",
                                    "llvm.aarch64.sme.ld1h.horiz",
                                    ptrue_fp16,
                                    input_ptr,
                                    sub_tile_idx,
                                    slice_idx * 2 + 1,
                                )
                            )

                    # Store columns to the output matrix
                    with T.serial(SVF // 2) as slice_idx:
                        for sub_tile_idx in range(2):
                            offset = slice_idx * 2 * A_t.strides[0] + (SVF * sub_tile_idx)
                            output_ptr = A_t.access_ptr("w", offset=offset)
                            T.evaluate(
                                T.call_llvm_intrin(
                                    "void",
                                    "llvm.aarch64.sme.st1w.vert",
                                    ptrue_fp32,
                                    output_ptr,
                                    sub_tile_idx,
                                    slice_idx,
                                )
                            )
                            output_ptr = A_t.access_ptr("w", offset=offset + A_t.strides[0])
                            T.evaluate(
                                T.call_llvm_intrin(
                                    "void",
                                    "llvm.aarch64.sme.st1w.vert",
                                    ptrue_fp32,
                                    output_ptr,
                                    sub_tile_idx + 2,
                                    slice_idx,
                                )
                            )

        return ib.get()

    return desc, impl()


def get_transpose_interleave_intrin_name(in_dtype, out_dtype, extent_cols, extent_rows):
    if in_dtype == "float32" and out_dtype == "float32":
        sme_transpose_interleave_intrin_name = (
            ARM_SME_2SVLx2SVL_FP32_TRANSPOSE_INTERLEAVE + f"_{extent_cols}_{extent_rows}"
        )
        tirx.TensorIntrin.register(
            sme_transpose_interleave_intrin_name,
            *get_sme_transpose_interleave_2svlx2svl_fp32_intrin(extent_cols, extent_rows),
            override=True,
        )
        return sme_transpose_interleave_intrin_name
    elif in_dtype == "float16" and out_dtype == "float32":
        return ARM_SME_BLOCK2_2SVLx1SVL_FP16_TRANSPOSE_INTERLEAVE
    else:
        raise ValueError("Input/output data type combination not supported.")


def get_sme_gemm_interleaved_mopa_2svlx2svl_intrin(M, K, in_dtype):
    """
    Compute a GEMM of size 2SVL x 2SVL (where 'SVL' is the Scalable Vector Length using
    outer product operations from the Scalable Matrix Extension (SME).

    The inputs A and B are expected to be of size K x 2SVL and produce a result C of
    size 2SVL x 2SVL.

    The SME accumulator tile is divided into sub-tiles, each of which is utilized to
    calculate the outer-product using columns / rows of A and B respectively. For each
    sub-tile, elements in the first column of input matrix A (accessed sequentially due
    to being transpose-interleaved) and first row of input matrix B are used to calculate
    an outer-product. This is then accumulated with the result of performing an
    outer-product on the second column and row of A and B respectively. This process is
    repeated K times. Finally, the results of the accumulation are stored.

    Note: The input tensor 'A' must be transpose-interleaved.

    Example
    -------

    Diagram showing outer-product performed on each of the accumulator sub-tiles
    for the fp32 datatype:

    ::

                           SVL           SVL
                    +----------------------------+
                    |       l     |       h      | K
                K   +----------------------------+
             +---+  +----------------------------+
             |   |  |  0:            1:          |-+
             |   |  |  mopa(l, l)    mopa(l, h)  | |-+
           l |   |  |                            | | |
             |   |  |                            | | |
             |---|  |                            | | |
             |   |  |  2:            3:          | | |
           h |   |  |  mopa(h, l)    mopa(h, h)  | | |
             |   |  |                            | | |
             |   |  |                            | | |
             +---+  +----------------------------+ | |
                      +----------------------------+ |
                         +---------------------------+
                                        (accumulate K times)

    Pseudo code computing 2SVL x 2SVL GEMM for fp32 inputs:

    .. code-block:: c

        // Number of fp32 elements in a scalable vector
        int SVF = SVL / 32;

        // Reset the accumulator tile
        sme.zero();

        // Calculate outer products and accumulate
        for (k = 0; k < K; k++) {
            float32xSVF A_row_0 = A[k][0];
            float32xSVF A_row_1 = A[k][SVF];
            float32xSVF B_row_0 = B[k][0];
            float32xSVF B_row_1 = B[k][SVF];

            float32xSVFxSVF sub_tile_0 += sme.mopa(A_row_0, B_row_0);
            float32xSVFxSVF sub_tile_1 += sme.mopa(A_row_0, B_row_1);
            float32xSVFxSVF sub_tile_2 += sme.mopa(A_row_1, B_row_0);
            float32xSVFxSVF sub_tile_3 += sme.mopa(A_row_1, B_row_1);
        }

        // Store the results of accumulation
        for (i = 0; i < SVF; i++) {
            C[i][0] = sme.horiz(sub_tile_0[i]);
            C[i][0] = sme.horiz(sub_tile_0[i + SVF]);
            C[i + SVF][0] = sme.horiz(sub_tile_0[i]);
            C[i + SVF][0] = sme.horiz(sub_tile_0[i + SVF]);
        }

    Notes:

    - Recall that A has been transposed beforehand such that each column is now accessed
      by row.
    - 'sme.zero' resets the accumulator tile to contain all zero's.
    - 'sme.mopa' is the outer product and accumulate intrinsic.
    - 'sme.horiz' stores rows of an accumulator sub-tile to memory.

    Returns
    -------
    intrin : TensorIntrin
        The SME TensorIntrin that can be used in tensorizing a schedule.

    """
    SVF = tirx.get_vscale_expr("float32")
    SVF2 = 2 * SVF
    fmopa_intrin = (
        "llvm.aarch64.sme.mopa" if in_dtype == "float32" else "llvm.aarch64.sme.mopa.wide"
    )

    @T.prim_func(s_tir=True)
    def desc(a: T.handle, b: T.handle, c: T.handle):
        A = T.match_buffer(a, (K, SVF2), dtype=in_dtype, offset_factor=1)
        B = T.match_buffer(b, (K, SVF2), dtype=in_dtype, offset_factor=1)
        C = T.match_buffer(c, (SVF2, SVF2), dtype="float32", offset_factor=1)

        with T.sblock("root"):
            T.reads(C[0:SVF2, 0:SVF2], A[0:K, 0:SVF2], B[0:K, 0:SVF2])
            T.writes(C[0:SVF2, 0:SVF2])
            for m, n, k in T.grid(SVF2, SVF2, K):
                with T.sblock("gemm"):
                    v_m, v_n, v_k = T.axis.remap("SSR", [m, n, k])
                    C[v_m, v_n] += T.Cast("float32", A[v_k, v_m]) * T.Cast("float32", B[v_k, v_n])

    def impl():
        sub_tile_count = 4

        with IRBuilder() as ib:
            with build_prim_func():
                a = T.arg("a", T.handle())
                b = T.arg("b", T.handle())
                c = T.arg("c", T.handle())

                A = T.match_buffer(a, (K, SVF2), in_dtype, offset_factor=1, strides=[T.int32(), 1])
                B = T.match_buffer(b, (K, SVF2), in_dtype, offset_factor=1, strides=[T.int32(), 1])
                C = T.match_buffer(
                    c, (SVF2, SVF2), "float32", offset_factor=1, strides=[T.int32(), 1]
                )

                ptrue = _create_ptrue_mask(in_dtype)

                with T.sblock("root"):
                    T.reads(C[0:SVF2, 0:SVF2], A[0:K, 0:SVF2], B[0:K, 0:SVF2])
                    T.writes(C[0:SVF2, 0:SVF2])

                    # Iterate over the reduction axis applying outer product and accumulate
                    rows_per_iter = 1 if in_dtype == "float32" else 2
                    with T.serial(T.ceildiv(K, rows_per_iter)) as k:
                        k_row = k * rows_per_iter
                        in_dtype_svf = tirx.get_vscale_expr(in_dtype)

                        # Ideally we'd rely on predicating the loads and use the same predicate
                        # for the outer product operation. However, support for predicated
                        # buffers is not currently supported by multiple lowering passes such as
                        # "LowerMatchBuffer", therefore the predicate is passed directly to the
                        # outer product operation for now.
                        if in_dtype == "float32":
                            a_low = (
                                T.BufferLoad(A, [k_row, T.Ramp(0, 1, in_dtype_svf)]),
                                _create_active_lane_mask(A, (k_row, 0), K),
                            )
                            b_low = (
                                T.BufferLoad(B, [k_row, T.Ramp(0, 1, in_dtype_svf)]),
                                _create_active_lane_mask(B, (k_row, 0), K),
                            )
                            a_high = (
                                T.BufferLoad(A, [k_row, T.Ramp(in_dtype_svf, 1, in_dtype_svf)]),
                                _create_active_lane_mask(A, (k_row, in_dtype_svf), K),
                            )
                            b_high = (
                                T.BufferLoad(B, [k_row, T.Ramp(in_dtype_svf, 1, in_dtype_svf)]),
                                _create_active_lane_mask(B, (k_row, in_dtype_svf), K),
                            )
                        else:
                            a_low = (T.BufferLoad(A, [k_row, T.Ramp(0, 1, in_dtype_svf)]), ptrue)
                            b_low = (T.BufferLoad(B, [k_row, T.Ramp(0, 1, in_dtype_svf)]), ptrue)
                            a_high = (
                                T.BufferLoad(A, [k_row + 1, T.Ramp(0, 1, in_dtype_svf)]),
                                ptrue,
                            )
                            b_high = (
                                T.BufferLoad(B, [k_row + 1, T.Ramp(0, 1, in_dtype_svf)]),
                                ptrue,
                            )

                        input_combinations = [
                            (a_low, b_low),
                            (a_low, b_high),
                            (a_high, b_low),
                            (a_high, b_high),
                        ]
                        for sub_tile_idx in range(0, sub_tile_count):
                            sub_tile = T.int32(sub_tile_idx)
                            input_1 = input_combinations[sub_tile_idx][0]
                            input_2 = input_combinations[sub_tile_idx][1]

                            T.evaluate(
                                T.call_llvm_intrin(
                                    "void",
                                    fmopa_intrin,
                                    sub_tile,
                                    input_1[1],
                                    input_2[1],
                                    input_1[0],
                                    input_2[0],
                                )
                            )

                    # Store the accumulated tile results
                    with T.serial(SVF) as slice_idx:
                        for sub_tile_idx in range(sub_tile_count):
                            vert_offset = SVF if sub_tile_idx >= (sub_tile_count // 2) else 0
                            horiz_offset = SVF if sub_tile_idx % 2 else 0
                            local_offset = (slice_idx + vert_offset) * C.strides[0] + horiz_offset
                            output_ptr = C.access_ptr("w", offset=local_offset, extent=SVF)

                            T.evaluate(
                                T.call_llvm_intrin(
                                    "void",
                                    "llvm.aarch64.sme.st1w.horiz",
                                    _create_active_lane_mask(
                                        C, (vert_offset + slice_idx, horiz_offset), M
                                    ),
                                    output_ptr,
                                    T.int32(sub_tile_idx),
                                    T.int32(slice_idx),
                                )
                            )

            return ib.get()

    return desc, impl()


def get_sme_init_intrin():
    """
    Reset the entire matrix tile storage to 0.
    """
    SVF2 = 2 * 4 * T.vscale()

    @T.prim_func(s_tir=True)
    def desc(c: T.handle) -> None:
        C = T.match_buffer(c, (SVF2, SVF2), "float32", offset_factor=1)
        with T.sblock("root"):
            T.reads()
            T.writes(C[0:SVF2, 0:SVF2])
            for m, n in T.grid(SVF2, SVF2):
                with T.sblock("init"):
                    v_m, v_n = T.axis.remap("SS", [m, n])
                    C[v_m, v_n] = T.float32(0)

    @T.prim_func(s_tir=True)
    def impl(c: T.handle) -> None:
        C = T.match_buffer(c, (SVF2, SVF2), "float32", offset_factor=1)
        with T.sblock("root"):
            T.reads()
            T.writes(C[0:SVF2, 0:SVF2])
            clear_all_tiles = T.int32(255)
            T.evaluate(T.call_llvm_intrin("void", "llvm.aarch64.sme.zero", clear_all_tiles))

    return desc, impl


ARM_DOT_4x4_i8_NEON_INTRIN = "dot_4x4_i8i8s32_neon"
ARM_DOT_4x4_i8_SDOT_INTRIN = "dot_4x4_i8i8s32_sdot"
ARM_DOT_4x4_u8_UDOT_INTRIN = "dot_4x4_u8u8u32_udot"
ARM_DOT_4x4_u8_HDOT_INTRIN = "dot_4x4_u8u8i32_hdot"

TensorIntrin.register(ARM_DOT_4x4_i8_NEON_INTRIN, neon_4x4_i8i8i32_desc, neon_4x4_i8i8i32_impl)
TensorIntrin.register(ARM_DOT_4x4_i8_SDOT_INTRIN, *get_dotprod_intrin("int8", "int32"))
TensorIntrin.register(ARM_DOT_4x4_u8_UDOT_INTRIN, *get_dotprod_intrin("uint8", "uint32"))
TensorIntrin.register(ARM_DOT_4x4_u8_HDOT_INTRIN, *get_dotprod_intrin("uint8", "int32"))

ARM_SME_INIT = "sme_init"
ARM_SME_2SVLx2SVL_FP32_TRANSPOSE_INTERLEAVE = "sme_2svlx2svl_fp32_transpose_interleave"
ARM_SME_BLOCK2_2SVLx1SVL_FP16_TRANSPOSE_INTERLEAVE = (
    "sme_block2_2svlx1svl_fp16_transpose_interleave"
)
ARM_SME_2SVLx2SVL_GEMM_INTERLEAVED_MOPA = "sme_2svlx2svl_gemm_interleaved_mopa"


# The following tensor intrinsics use LLVM intrinsics that are only available
# in versions of LLVM >= 15. Installations with older versions of LLVM will
# not be able to use them.
if llvm_version_major() >= 15:
    TensorIntrin.register(
        ARM_SME_BLOCK2_2SVLx1SVL_FP16_TRANSPOSE_INTERLEAVE,
        *get_sme_transpose_interleave_block2_2svl_fp16_intrin(),
    )
    TensorIntrin.register(ARM_SME_INIT, *get_sme_init_intrin())
