chore: import upstream snapshot with attribution

python/tvm/s_tir/tensor_intrin/__init__.py

@@ -0,0 +1,25 @@
+# isort: skip_file
+# 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=unused-import
+"""Intrinsics for tensorization."""
+
+from tvm.runtime import enabled
+from . import cuda
+
+if enabled("llvm"):
+    from . import arm_cpu, x86, rocm, hexagon, riscv_cpu

python/tvm/s_tir/tensor_intrin/arm_cpu.py

@@ -0,0 +1,777 @@
+# 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())

python/tvm/s_tir/tensor_intrin/dot_product_common.py

@@ -0,0 +1,73 @@
+# 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
+"""Dot product related intrinsics."""
+
+from tvm.script import tirx as T
+
+from .. import TensorIntrin
+
+
+def get_dp4a_intrin(dtype_a, dtype_b, dtype_c):
+    if dtype_c == "uint32":
+        assert dtype_a == dtype_b == "uint8"
+    vec_type_a = "int8x4" if dtype_a == "int8" else "uint8x4"
+    vec_type_b = "int8x4" if dtype_b == "int8" else "uint8x4"
+
+    @T.prim_func(s_tir=True)
+    def dp4a_desc(
+        A: T.Buffer((4,), dtype_a, offset_factor=1, align=4, scope="shared"),
+        B: T.Buffer((4,), dtype_b, offset_factor=1, align=4, scope="shared"),
+        C: T.Buffer((1,), dtype_c, offset_factor=1, align=4, scope="local"),
+    ) -> None:
+        with T.sblock("root"):
+            T.reads(C[0], A[0:4], B[0:4])
+            T.writes(C[0])
+            for i in range(0, 4):
+                with T.sblock("update"):
+                    vi = T.axis.remap("R", [i])
+                    C[0] = C[0] + T.cast(A[vi], dtype_c) * T.cast(B[vi], dtype_c)
+
+    @T.prim_func(s_tir=True)
+    def dp4a_impl(
+        A: T.Buffer((4,), dtype_a, offset_factor=1, align=4, scope="shared"),
+        B: T.Buffer((4,), dtype_b, offset_factor=1, align=4, scope="shared"),
+        C: T.Buffer((1,), dtype_c, offset_factor=1, align=4, scope="local"),
+    ) -> None:
+        with T.sblock("root"):
+            T.reads(C[0], A[0:4], B[0:4])
+            T.writes(C[0])
+
+            C[0] += T.call_pure_extern(
+                "__dp4a",
+                A.vload([0], vec_type_a),
+                B.vload([0], vec_type_b),
+                T.uint32(0) if dtype_c == "uint32" else T.int32(0),
+                dtype=dtype_c,
+            )
+
+    return dp4a_desc, dp4a_impl
+
+
+DP4A_S8S8S32_INTRIN = "dp4a_s8s8s32"
+TensorIntrin.register(DP4A_S8S8S32_INTRIN, *get_dp4a_intrin("int8", "int8", "int32"))
+DP4A_U8S8S32_INTRIN = "dp4a_u8s8s32"
+TensorIntrin.register(DP4A_U8S8S32_INTRIN, *get_dp4a_intrin("uint8", "int8", "int32"))
+DP4A_S8U8S32_INTRIN = "dp4a_s8u8s32"
+TensorIntrin.register(DP4A_S8U8S32_INTRIN, *get_dp4a_intrin("int8", "uint8", "int32"))
+DP4A_U8U8U32_INTRIN = "dp4a_u8u8u32"
+TensorIntrin.register(DP4A_U8U8U32_INTRIN, *get_dp4a_intrin("uint8", "uint8", "uint32"))

python/tvm/s_tir/tensor_intrin/hexagon.py

@@ -0,0 +1,225 @@
+# 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
+"""Intrinsics for Hexagon tensorization."""
+
+from tvm.script import tirx as T
+
+from .. import TensorIntrin
+
+
+def generate_dma_load_intrin(
+    size: int,
+    dtype: str,
+):
+    """Generator of dma_load intrins"""
+
+    @T.prim_func(s_tir=True)
+    def sync_dma_load_desc(a: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (size), dtype, offset_factor=1, scope="global")
+        C = T.match_buffer(c, (size), dtype, offset_factor=1, scope="global.vtcm")
+        with T.sblock("root"):
+            T.reads(A[0:size])
+            T.writes(C[0:size])
+            for i in T.serial(size):
+                with T.sblock("load"):
+                    vii = T.axis.remap("S", [i])
+                    C[vii] = A[vii]
+
+    @T.prim_func(s_tir=True)
+    def sync_dma_load_impl(a: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (size), dtype, offset_factor=1, scope="global")
+        C = T.match_buffer(c, (size), dtype, offset_factor=1, scope="global.vtcm")
+        with T.sblock("root"):
+            T.reads(A[0:size])
+            T.writes(C[0:size])
+            T.evaluate(
+                T.tvm_call_packed(
+                    "device_api.hexagon.dma_copy_dltensor",
+                    T.tvm_stack_make_array(
+                        T.address_of(C[0], dtype="handle"),
+                        T.tvm_stack_make_shape(size, dtype="handle"),
+                        0,
+                        1,
+                        C.dtype,
+                        0,
+                        dtype="handle",
+                    ),
+                    T.tvm_stack_make_array(
+                        T.address_of(A[0], dtype="handle"),
+                        T.tvm_stack_make_shape(size, dtype="handle"),
+                        0,
+                        1,
+                        A.dtype,
+                        0,
+                        dtype="handle",
+                    ),
+                    T.cast(size, dtype="int"),
+                    False,  # Do not use experimental bypass mode.
+                    dtype="int32",
+                )
+            )
+
+    return sync_dma_load_desc, sync_dma_load_impl
+
+
+def generate_dot_product_32x4_u8u8i32(mem_scope="global"):
+    @T.prim_func(s_tir=True)
+    def dot_product_32x4_u8u8i32_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (4,), "uint8", offset_factor=1, scope=mem_scope)
+        B = T.match_buffer(b, (32, 4), "uint8", offset_factor=1, scope=mem_scope)
+        C = T.match_buffer(c, (32,), "int32", offset_factor=1, scope=mem_scope)
+        with T.sblock("root"):
+            T.reads(C[0:32], A[0:4], B[0:32, 0:4])
+            T.writes(C[0:32])
+            for i in T.serial(0, 32):
+                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 dot_product_32x4_u8u8i32_vrmpy(a: T.handle, b: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (4,), "uint8", offset_factor=1, scope=mem_scope)
+        B = T.match_buffer(b, (32, 4), "uint8", offset_factor=1, scope=mem_scope)
+        C = T.match_buffer(c, (32,), "int32", offset_factor=1, scope=mem_scope)
+        with T.sblock("root"):
+            T.reads(C[0:32], A[0:4], B[0:32, 0:4])
+            T.writes(C[0:32])
+
+            A_u8x4 = A.vload([0], "uint8x4")
+            A_i32 = T.reinterpret(A_u8x4, dtype="int32")
+
+            B_i8x128 = B.vload([0, 0], dtype="uint8x128")
+            B_i32x32 = T.reinterpret(B_i8x128, dtype="int32x32")
+
+            C[T.ramp(T.int32(0), 1, 32)] = T.call_llvm_pure_intrin(
+                T.llvm_lookup_intrinsic_id("llvm.hexagon.V6.vrmpyub.acc.128B"),
+                C[T.ramp(T.int32(0), 1, 32)],
+                B_i32x32,
+                A_i32,
+                dtype="int32x32",
+            )
+
+    return dot_product_32x4_u8u8i32_desc, dot_product_32x4_u8u8i32_vrmpy
+
+
+def generate_dot_product_32x4_u8i8i32(mem_scope="global"):
+    @T.prim_func(s_tir=True)
+    def dot_product_32x4_u8i8i32_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (4,), "uint8", offset_factor=1, scope=mem_scope)
+        B = T.match_buffer(b, (32, 4), "int8", offset_factor=1, scope=mem_scope)
+        C = T.match_buffer(c, (32,), "int32", offset_factor=1, scope=mem_scope)
+        with T.sblock("root"):
+            T.reads(C[0:32], A[0:4], B[0:32, 0:4])
+            T.writes(C[0:32])
+            for i in T.serial(0, 32):
+                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 dot_product_32x4_u8i8i32_vrmpy(a: T.handle, b: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (4,), "uint8", offset_factor=1, scope=mem_scope)
+        B = T.match_buffer(b, (32, 4), "int8", offset_factor=1, scope=mem_scope)
+        C = T.match_buffer(c, (32,), "int32", offset_factor=1, scope=mem_scope)
+        with T.sblock("root"):
+            T.reads(C[0:32], A[0:4], B[0:32, 0:4])
+            T.writes(C[0:32])
+
+            A_u8x4 = A.vload([0], "uint8x4")
+            A_i32 = T.reinterpret(A_u8x4, dtype="int32")
+
+            B_i8x128 = B.vload([0, 0], dtype="int8x128")
+            B_i32x32 = T.reinterpret(B_i8x128, dtype="int32x32")
+
+            C[T.ramp(T.int32(0), 1, 32)] = T.call_llvm_pure_intrin(
+                T.llvm_lookup_intrinsic_id("llvm.hexagon.V6.vrmpybusv.acc.128B"),
+                C[T.ramp(T.int32(0), 1, 32)],
+                T.broadcast(A_i32, 32),
+                B_i32x32,
+                dtype="int32x32",
+            )
+
+    return dot_product_32x4_u8i8i32_desc, dot_product_32x4_u8i8i32_vrmpy
+
+
+def generate_dot_product_32x2_i16i16i32(mem_scope="global"):
+    @T.prim_func(s_tir=True)
+    def dot_product_32x2_i16i16i32_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (2,), "int16", offset_factor=1, scope=mem_scope)
+        B = T.match_buffer(b, (32, 2), "int16", offset_factor=1, scope=mem_scope)
+        C = T.match_buffer(c, (32,), "int32", offset_factor=1, scope=mem_scope)
+        with T.sblock("root"):
+            T.reads(C[0:32], A[0:2], B[0:32, 0:2])
+            T.writes(C[0:32])
+            for i in T.serial(0, 32):
+                for k in T.serial(0, 2):
+                    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 dot_product_32x2_i16i16i32_vdmpy(a: T.handle, b: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (2,), "int16", offset_factor=1, scope=mem_scope)
+        B = T.match_buffer(b, (32, 2), "int16", offset_factor=1, scope=mem_scope)
+        C = T.match_buffer(c, (32,), "int32", offset_factor=1, scope=mem_scope)
+        with T.sblock("root"):
+            T.reads(C[0:32], A[0:2], B[0:32, 0:2])
+            T.writes(C[0:32])
+
+            A_i16x2 = A.vload([0], "int16x2")
+            A_i32 = T.reinterpret(A_i16x2, dtype="int32")
+
+            B_i16x64 = B.vload([0, 0], dtype="int16x64")
+            B_i32x32 = T.reinterpret(B_i16x64, dtype="int32x32")
+
+            C[T.ramp(T.int32(0), 1, 32)] = T.call_llvm_pure_intrin(
+                T.llvm_lookup_intrinsic_id("llvm.hexagon.V6.vdmpyhvsat.acc.128B"),
+                C[T.ramp(T.int32(0), 1, 32)],
+                T.Broadcast(A_i32, 32),
+                B_i32x32,
+                dtype="int32x32",
+            )
+
+    return dot_product_32x2_i16i16i32_desc, dot_product_32x2_i16i16i32_vdmpy
+
+
+VRMPY_u8u8i32_INTRIN = "dot_32x4_u8u8i32_vrmpy"
+
+TensorIntrin.register(VRMPY_u8u8i32_INTRIN, *generate_dot_product_32x4_u8u8i32())
+
+VRMPY_u8i8i32_INTRIN = "dot_32x4_u8i8i32_vrmpy"
+
+TensorIntrin.register(VRMPY_u8i8i32_INTRIN, *generate_dot_product_32x4_u8i8i32())
+
+VDMPY_i16i16i32_INTRIN = "dot_product_32x2_i16i16i32_vdmpy"
+
+TensorIntrin.register(VDMPY_i16i16i32_INTRIN, *generate_dot_product_32x2_i16i16i32())
+
+VRMPY_u8u8i32_VTCM_INTRIN = "dot_32x4_u8u8i32_vtcm_vrmpy"
+TensorIntrin.register(VRMPY_u8u8i32_VTCM_INTRIN, *generate_dot_product_32x4_u8u8i32("global.vtcm"))
+
+VRMPY_u8i8i32_VTCM_INTRIN = "dot_32x4_u8i8i32_vtcm_vrmpy"
+TensorIntrin.register(VRMPY_u8i8i32_VTCM_INTRIN, *generate_dot_product_32x4_u8i8i32("global.vtcm"))
+
+DMA_READ_128_u8 = "dma_read_128_u8"
+TensorIntrin.register(DMA_READ_128_u8, *generate_dma_load_intrin(128, "uint8"))
+
+DMA_READ_128_i8 = "dma_read_128_i8"
+TensorIntrin.register(DMA_READ_128_i8, *generate_dma_load_intrin(128, "int8"))

python/tvm/s_tir/tensor_intrin/metal.py

@@ -0,0 +1,351 @@
+# 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-variable
+"""Intrinsics for tensorization on Apple GPU."""
+
+from typing import Literal
+
+from tvm.script import tirx as T
+from tvm.tirx import Buffer, Expr, PrimFunc, TensorIntrin
+
+######## simdgroup matrix intrinsics ########
+
+
+def get_simdgroup_index(buffer: Buffer, stride: Expr, col: int, row: int):
+    """Compute simdgroup index using elem_offset of the buffer"""
+
+    # NOTE: Need further check the usage between `col`` and `row`
+    # Currently, Metal only supports 8x8, which means the values of `col` and `row` are the same
+    frag_index_m = buffer.elem_offset // stride // col
+    frag_index_n = buffer.elem_offset % stride // row
+
+    num_fragments_per_row = stride // row
+    return frag_index_m * num_fragments_per_row + frag_index_n
+
+
+def get_make_filled_simdgroup_matrix_intrin(
+    dtype: str, col: int = 8, row: int = 8
+) -> tuple[PrimFunc, PrimFunc]:
+    @T.prim_func(s_tir=True)
+    def desc(a: T.handle) -> None:
+        A = T.match_buffer(a, (col, row), dtype, scope="metal.simdgroup", offset_factor=1)
+        with T.sblock("root"):
+            T.reads()
+            T.writes(A[0:col, 0:row])
+            for i, j in T.grid(col, row):
+                with T.sblock("init"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    A[vi, vj] = T.float32(0)
+
+    @T.prim_func(s_tir=True)
+    def impl(a: T.handle) -> None:
+        d0, d1 = T.int32(), T.int32()
+        A = T.match_buffer(
+            a, (col, row), dtype, scope="metal.simdgroup", strides=[d1, d0], offset_factor=1
+        )
+        with T.sblock("root"):
+            T.reads()
+            T.writes(A[0:col, 0:row])
+            T.metal.make_filled_simdgroup_matrix(
+                A.data,
+                index=get_simdgroup_index(A, d1, col, row),
+                value=T.float32(0),
+                col=col,
+                row=row,
+            )
+
+    return desc, impl
+
+
+def get_simdgroup_load_intrin(
+    dtype: str,
+    scope: Literal["global", "shared"],
+    col: int = 8,
+    row: int = 8,
+    transpose_matrix: bool = False,
+) -> tuple[PrimFunc, PrimFunc]:
+    align = col * row
+
+    @T.prim_func(s_tir=True)
+    def desc(a: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (col, row), dtype, align=align, scope=scope, offset_factor=1)
+        C = T.match_buffer(
+            c, (col, row), dtype, align=align, scope="metal.simdgroup", offset_factor=1
+        )
+        with T.sblock("root"):
+            T.reads(A[0:col, 0:row])
+            T.writes(C[0:col, 0:row])
+            for i, j in T.grid(col, row):
+                with T.sblock("load"):
+                    vii, vjj = T.axis.remap("SS", [i, j])
+                    if transpose_matrix:
+                        # C[vii, vjj] = A[vjj, vii]
+                        C[vjj, vii] = A[vii, vjj]
+                    else:
+                        C[vii, vjj] = A[vii, vjj]
+
+    @T.prim_func(s_tir=True)
+    def impl(a: T.handle, c: T.handle) -> None:
+        s0, s1, d0, d1 = T.int32(), T.int32(), T.int32(), T.int32()
+        A = T.match_buffer(
+            a,
+            (col, row),
+            dtype,
+            align=align,
+            scope=scope,
+            strides=[s1, s0],
+            offset_factor=1,
+        )
+        C = T.match_buffer(
+            c,
+            (col, row),
+            dtype,
+            align=align,
+            scope="metal.simdgroup",
+            strides=[d1, d0],
+            offset_factor=1,
+        )
+        with T.sblock("root"):
+            T.reads(A[0:col, 0:row])
+            T.writes(C[0:col, 0:row])
+            T.metal.simdgroup_load(
+                C.data,
+                index=get_simdgroup_index(C, d1, col, row),
+                ptr=A.access_ptr("r", ptr_type=dtype),
+                stride=s1,
+                col=col,
+                row=row,
+                transpose_matrix=transpose_matrix,
+            )
+
+    return desc, impl
+
+
+def get_simdgroup_store_intrin(
+    dtype: str,
+    scope: Literal["global", "shared"],
+    col: int = 8,
+    row: int = 8,
+    transpose_matrix: bool = False,
+) -> tuple[PrimFunc, PrimFunc]:
+    align = col * row
+
+    @T.prim_func(s_tir=True)
+    def desc(a: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(
+            a, (col, row), dtype, align=align, scope="metal.simdgroup", offset_factor=1
+        )
+        C = T.match_buffer(c, (col, row), dtype, align=align, scope=scope, offset_factor=1)
+        with T.sblock("root"):
+            T.reads(A[0:col, 0:row])
+            T.writes(C[0:col, 0:row])
+            for i, j in T.grid(col, row):
+                with T.sblock("store"):
+                    vii, vjj = T.axis.remap("SS", [i, j])
+                    if transpose_matrix:
+                        C[vjj, vii] = A[vii, vjj]
+                    else:
+                        C[vii, vjj] = A[vii, vjj]
+
+    @T.prim_func(s_tir=True)
+    def impl(a: T.handle, c: T.handle) -> None:
+        s0, s1, d0, d1 = T.int32(), T.int32(), T.int32(), T.int32()
+        A = T.match_buffer(
+            a,
+            (col, row),
+            dtype,
+            align=align,
+            scope="metal.simdgroup",
+            strides=[s1, s0],
+            offset_factor=1,
+        )
+        C = T.match_buffer(
+            c, (col, row), dtype, align=align, scope=scope, strides=[d1, d0], offset_factor=1
+        )
+        with T.sblock("root"):
+            T.reads(A[0:col, 0:row])
+            T.writes(C[0:col, 0:row])
+            T.metal.simdgroup_store(
+                A.data,
+                index=get_simdgroup_index(A, s1, col, row),
+                ptr=C.access_ptr("w", ptr_type=dtype),
+                stride=d1,
+                col=col,
+                row=row,
+                transpose_matrix=transpose_matrix,
+            )
+
+    return desc, impl
+
+
+def get_simdgroup_multiply_accumulate_intrin(
+    m_dim: int, n_dim: int, k_dim: int, dtype: str
+) -> tuple[PrimFunc, PrimFunc]:
+    @T.prim_func(s_tir=True)
+    def desc(a: T.handle, b: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (m_dim, k_dim), dtype, scope="metal.simdgroup", offset_factor=1)
+        B = T.match_buffer(b, (k_dim, n_dim), dtype, scope="metal.simdgroup", offset_factor=1)
+        C = T.match_buffer(c, (m_dim, n_dim), dtype, scope="metal.simdgroup", offset_factor=1)
+        with T.sblock("root"):
+            T.reads(C[0:m_dim, 0:n_dim], A[0:m_dim, 0:k_dim], B[0:k_dim, 0:n_dim])
+            T.writes(C[0:m_dim, 0:n_dim])
+            for i, j, k in T.grid(m_dim, n_dim, k_dim):
+                with T.sblock(""):
+                    vii, vjj, vkk = T.axis.remap("SSR", [i, j, k])
+                    C[vii, vjj] += A[vii, vkk] * B[vkk, vjj]
+
+    @T.prim_func(s_tir=True)
+    def impl(a: T.handle, b: T.handle, c: T.handle) -> None:
+        a0, a1, b0, b1, c0, c1 = T.int32(), T.int32(), T.int32(), T.int32(), T.int32(), T.int32()
+        A = T.match_buffer(
+            a, (m_dim, k_dim), dtype, scope="metal.simdgroup", strides=[a1, a0], offset_factor=1
+        )
+        B = T.match_buffer(
+            b, (k_dim, n_dim), dtype, scope="metal.simdgroup", strides=[b1, b0], offset_factor=1
+        )
+        C = T.match_buffer(
+            c, (m_dim, n_dim), dtype, scope="metal.simdgroup", strides=[c1, c0], offset_factor=1
+        )
+        with T.sblock("root"):
+            T.reads(C[0:m_dim, 0:n_dim], A[0:m_dim, 0:k_dim], B[0:k_dim, 0:n_dim])
+            T.writes(C[0:m_dim, 0:n_dim])
+            T.metal.simdgroup_multiply_accumulate(
+                C.data,
+                get_simdgroup_index(C, c1, m_dim, n_dim),
+                A.data,
+                get_simdgroup_index(A, a1, m_dim, k_dim),
+                B.data,
+                get_simdgroup_index(B, b1, k_dim, n_dim),
+                C.data,
+                get_simdgroup_index(C, c1, m_dim, n_dim),
+            )
+
+    return desc, impl
+
+
+# Make filled simdgroup matrix intrinsics
+
+SIMDGROUP_MAKE_FILLED_8x8x8_f16_INTRIN = "simdgroup_make_filled_8x8x8_f16"
+TensorIntrin.register(
+    SIMDGROUP_MAKE_FILLED_8x8x8_f16_INTRIN,
+    *get_make_filled_simdgroup_matrix_intrin("float16", 8, 8),
+)
+
+SIMDGROUP_FILLED_8x8x8_f32_INTRIN = "simdgroup_fill_8x8x8_f32"
+TensorIntrin.register(
+    SIMDGROUP_FILLED_8x8x8_f32_INTRIN, *get_make_filled_simdgroup_matrix_intrin("float32", 8, 8)
+)
+
+SIMDGROUP_FILLED_8x8x8_bf16_INTRIN = "simdgroup_fill_8x8x8_bf16"
+TensorIntrin.register(
+    SIMDGROUP_FILLED_8x8x8_bf16_INTRIN, *get_make_filled_simdgroup_matrix_intrin("bfloat16", 8, 8)
+)
+
+# Load intrinsics
+
+SIMDGROUP_LOAD_8x8x8_f16_SHARED_INTRIN = "simdgroup_load_8x8x8_f16_shared"
+TensorIntrin.register(
+    SIMDGROUP_LOAD_8x8x8_f16_SHARED_INTRIN,
+    *get_simdgroup_load_intrin("float16", "shared", 8, 8, False),
+)
+
+SIMDGROUP_LOAD_8x8x8_f16_SHARED_TRANS_INTRIN = "simdgroup_load_8x8x8_f16_shared_trans"
+TensorIntrin.register(
+    SIMDGROUP_LOAD_8x8x8_f16_SHARED_TRANS_INTRIN,
+    *get_simdgroup_load_intrin("float16", "shared", 8, 8, True),
+)
+
+# Store intrinsics
+
+SIMDGROUP_STORE_8x8x8_f16_GLOBAL_INTRIN = "simdgroup_store_8x8x8_f16_global"
+TensorIntrin.register(
+    SIMDGROUP_STORE_8x8x8_f16_GLOBAL_INTRIN,
+    *get_simdgroup_store_intrin("float16", "global", 8, 8, False),
+)
+
+SIMDGROUP_STORE_8x8x8_f16_SHARED_INTRIN = "simdgroup_store_8x8x8_f16_shared"
+TensorIntrin.register(
+    SIMDGROUP_STORE_8x8x8_f16_SHARED_INTRIN,
+    *get_simdgroup_store_intrin("float16", "shared", 8, 8, False),
+)
+# Multiply accumulate intrinsics
+
+SIMDGROUP_MULTI_ACC_8x8x8_f16_INTRIN = "simdgroup_multiply_accumulate_8x8x8_f16"
+TensorIntrin.register(
+    SIMDGROUP_MULTI_ACC_8x8x8_f16_INTRIN,
+    *get_simdgroup_multiply_accumulate_intrin(8, 8, 8, "float16"),
+)
+
+
+def get_simdgroup_intrin_group(
+    load_scope: Literal["shared"],
+    store_scope: Literal["global", "shared"],
+    dtype: str,
+    trans_a: bool = False,
+    trans_b: bool = False,
+) -> dict[str, str]:
+    """Get a group of intrinsics for tensorization on Apple GPU.
+
+    Parameters
+    ----------
+    load_scope : Literal["shared"]
+        The memory scope of the input buffer.
+
+    store_scope : Literal["global", "shared"]
+        The memory scope of the result buffer.
+
+    dtype : str
+        The data type of the input and output buffers.
+
+    trans_a : bool
+        Whether the input matrix A is transposed.
+
+    trans_b : bool
+        Whether the input matrix B is transposed.
+
+    Returns
+    -------
+    ret : Dict[str, str]
+        A group of tensor intrinsics.
+    """
+    assert load_scope in ["shared"]
+    assert store_scope in ["global", "shared"]
+    assert dtype in ["float16", "bfloat16", "float32"]
+
+    shape = "8x8x8"
+    dtype = "f16" if dtype == "float16" else "bf16" if dtype == "bfloat16" else "f32"
+    trans_a = "_trans" if trans_a else ""
+    trans_b = "_trans" if trans_b else ""
+
+    # e.g. simdgroup_load_8x8x8_f16_shared
+    load_a_intrin = f"simdgroup_load_{shape}_{dtype}_{load_scope}{trans_a}"
+    # e.g. simdgroup_load_8x8x8_f16_shared_trans
+    load_b_intrin = f"simdgroup_load_{shape}_{dtype}_{load_scope}{trans_b}"
+    # e.g. simdgroup_multiply_accumulate_8x8x8_f16
+    compute_intrin = f"simdgroup_multiply_accumulate_{shape}_{dtype}"
+    # e.g. simdgroup_make_filled_8x8x8_f16
+    init_intrin = f"simdgroup_make_filled_{shape}_{dtype}"
+    # e.g. simdgroup_store_8x8x8_f16_global
+    store_intrin = f"simdgroup_store_{shape}_{dtype}_{store_scope}"
+
+    return {
+        "init": init_intrin,
+        "load_a": load_a_intrin,
+        "load_b": load_b_intrin,
+        "compute": compute_intrin,
+        "store": store_intrin,
+    }

python/tvm/s_tir/tensor_intrin/riscv_cpu.py

@@ -0,0 +1,241 @@
+# 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,line-too-long
+# ruff: noqa: E501
+"""Intrinsics for RISCV tensorization"""
+
+import logging
+
+import tvm_ffi
+
+from tvm.runtime import DataType
+from tvm.script import tirx as T
+from tvm.target.codegen import Target, llvm_get_vector_width, target_has_features
+
+from .. import TensorIntrin
+
+logger = logging.getLogger(__name__)
+
+
+def get_max_elems(vlen: int, lmul: int, sew: int) -> int:
+    """Returns number of elements of a given data type (SEW)
+    that fits multiple (LMUL) of the vector registers (VLEN).
+
+    Args:
+        vlen (int): VLEN vector length in bits
+        lmul (int): LMUL vector lenght multiplier
+        sew (int): SEW standard (single) element width
+
+    Returns:
+        int: Number of elements
+    """
+    return (vlen // sew) * lmul
+
+
+def rvv_vec_dot_product_kernels(
+    n_elems: int,
+    n_lanes: int,
+    data_dtype: str,
+    weight_dtype: str,
+    out_dtype: str,
+    lmul: int,
+):
+    """Dot product of vector and matrix rows using RISC-V vector instructions.
+
+    These kernels takes two arrays A[ELEMS] and B[ELEMS][MACS] and computes
+    dot product of A[ELEMS] with each row of B[LANES], accumulating results
+    with C[LANES].
+
+    The pseudo code is as follows:
+
+    .. code-block:: c
+
+        void vec_dot_prod(A[ELEMS], B[LANES][ELEMS], C[LANES]){
+            for (j = 0; j < LANES; j++) {
+                for (k = 0; k < ELEMS; k++) {
+                    C[j] += A[k] * B[j][k]
+                }
+            }
+        }
+    """
+
+    @T.prim_func(s_tir=True)
+    def rvv_vec_dot_prod_desc(
+        A: T.Buffer((n_elems,), data_dtype, offset_factor=1),
+        B: T.Buffer((n_lanes, n_elems), weight_dtype, offset_factor=1),
+        C: T.Buffer((n_lanes,), out_dtype, offset_factor=1),
+    ) -> None:
+        with T.sblock("root"):
+            T.reads(C[0:n_lanes], A[0:n_elems], B[0:n_lanes, 0:n_elems])
+            T.writes(C[0:n_lanes])
+            for j in T.serial(0, n_lanes):
+                for k in T.serial(0, n_elems):
+                    with T.sblock("update"):
+                        vj, vk = T.axis.remap("SR", [j, k])
+                        C[vj] = C[vj] + T.cast(A[vk], out_dtype) * T.cast(B[vj, vk], out_dtype)
+
+    # LLVM only supports ELEN=32 or ELEN=64
+    # https://llvm.org/docs//RISCV/RISCVVectorExtension.html
+    d_dtype_lanes = (64 // DataType(data_dtype).bits) * lmul
+    w_dtype_lanes = (64 // DataType(weight_dtype).bits) * lmul
+    # reduction lanes narrows
+    o_dtype_lanes = (64 // DataType(out_dtype).bits) * lmul // n_lanes
+    # data type widening case
+    o_dtype_lanes = max(o_dtype_lanes, 2)
+
+    mask_args = () if data_dtype[0] in ("i", "u") else (T.uint64(7),)
+
+    wide_dtype = out_dtype
+    if DataType(out_dtype).bits > DataType(data_dtype).bits:
+        wide_dtype = "".join(c for c in data_dtype if not c.isdigit())
+        wide_dtype += str(DataType(data_dtype).bits * 2)
+
+    # fmt: off
+    @T.prim_func(s_tir=True)
+    def rvv_vec_dot_prod_impl(
+        A: T.Buffer((n_elems,), data_dtype, offset_factor=1),
+        B: T.Buffer((n_lanes, n_elems), weight_dtype, offset_factor=1),
+        C: T.Buffer((n_lanes,), out_dtype, offset_factor=1),
+    ) -> None:
+        with T.sblock("root"):
+            T.reads(C[0:n_lanes], A[0:n_elems], B[0:n_lanes, 0:n_elems])
+            T.writes(C[0:n_lanes])
+
+            vec_A = T.call_llvm_intrin(
+                f"{data_dtype}xvscalex{d_dtype_lanes}",
+                "llvm.riscv.vle",
+                T.broadcast(T.Cast(data_dtype, 0), T.vscale() * d_dtype_lanes),
+                T.tvm_access_ptr(T.type_annotation(data_dtype), A.data, 0, n_elems, 1),
+                T.int64(n_elems))
+
+            for i in range(n_lanes):
+                with T.sblock("update"):
+                    T.reads(B[i, 0:n_elems])
+                    T.writes(C[i])
+
+                    vec_B_row = T.call_llvm_intrin(
+                        f"{weight_dtype}xvscalex{w_dtype_lanes}",
+                        "llvm.riscv.vle",
+                        T.broadcast(T.Cast(data_dtype, 0), T.vscale() * w_dtype_lanes),
+                        T.tvm_access_ptr(T.type_annotation(weight_dtype), B.data, i * n_elems, n_elems, 1),
+                        T.int64(n_elems))
+
+                    product = T.call_llvm_intrin(
+                        f"{wide_dtype}xvscalex{w_dtype_lanes}",
+                        "llvm.riscv.vfmul" if out_dtype[0] == "f" else \
+                        "llvm.riscv.vwmulsu" if (data_dtype[0] != weight_dtype[0]) else \
+                        "llvm.riscv.vwmul",
+                        T.broadcast(T.Cast(wide_dtype, 0), T.vscale() * w_dtype_lanes),
+                        vec_B_row,
+                        vec_A,
+                        *mask_args,
+                        T.uint64(n_elems))
+
+                    ini_acc = T.call_llvm_intrin(
+                        f"{out_dtype}xvscalex{o_dtype_lanes}",
+                        "llvm.riscv.vle",
+                        T.broadcast(T.Cast(out_dtype, 0), T.vscale() * o_dtype_lanes),
+                        T.tvm_access_ptr(T.type_annotation(out_dtype), C.data, i, 1, 1),
+                        T.int64(1))
+
+                    red_sum = T.call_llvm_intrin(
+                        f"{out_dtype}xvscalex{o_dtype_lanes}",
+                        "llvm.riscv.vfredusum" if out_dtype[0] == "f" else \
+                        "llvm.riscv.vwredsum",
+                        T.broadcast(T.Cast(out_dtype, 0), T.vscale() * o_dtype_lanes),
+                        product,
+                        ini_acc,
+                        *mask_args,
+                        T.uint64(n_elems))
+
+                    C[i] = T.call_llvm_intrin(
+                        out_dtype,
+                        "llvm.riscv.vfmv.f.s" if out_dtype[0] == "f" else \
+                        "llvm.riscv.vmv.x.s",
+                        red_sum)
+    # fmt: on
+    return rvv_vec_dot_prod_desc, rvv_vec_dot_prod_impl
+
+
+@tvm_ffi.register_global_func("tirx.tensor_intrin.register_rvv_isa_intrinsics")
+def register_rvv_isa_intrinsics(target: Target, inventory_only=False) -> dict():
+    """Register RISCV V (vector) intrinsics
+    [x] Implementation follows version 1.0 vector specifications:
+        https://github.com/riscvarchive/riscv-v-spec/releases/tag/v1.0
+
+    Args:
+        target (Target): TVM target
+        inventory_only (bool): No registration inventory only
+
+    Returns:
+        dict(): A catalog with registered kernel names and properties
+    """
+    if not target_has_features("v", target):
+        raise RuntimeError("Current target does not support `v` extension.")
+
+    vlen = llvm_get_vector_width(target)
+    # get maximum reduction lanes (without grouping)
+    n_lanes = get_max_elems(vlen, lmul=1, sew=32)
+
+    kernels_inventory = {}
+
+    data_dtype = ["uint8", "int8", "float16", "float32"]
+    weight_dtype = ["int8", "int8", "float16", "float32"]
+    output_dtype = ["int32", "int32", "float16", "float32"]
+
+    for d_dtype, w_dtype, o_dtype in zip(data_dtype, weight_dtype, output_dtype):
+        # max elements to grouped registers
+        max_elems = get_max_elems(vlen, lmul=8, sew=DataType(d_dtype).bits)
+        # data widening halves available vector registers
+        if DataType(o_dtype).bits > DataType(d_dtype).bits:
+            max_elems //= 2
+        # compute optimal LMUL for full load
+        lmul = max_elems // (vlen // DataType(d_dtype).bits)
+
+        n_elems = max_elems
+        while n_elems >= 4:
+            dt = DataType(d_dtype)
+            wt = DataType(w_dtype)
+            ot = DataType(o_dtype)
+            kernel_name = "rvv_dot"
+            kernel_name += f"_{n_elems}{dt[0]}{dt.bits}"
+            kernel_name += f"_{n_lanes}x{n_elems}{wt[0]}{wt.bits}"
+            kernel_name += f"_{n_lanes}{ot[0]}{ot.bits}"
+            kernels_inventory[kernel_name] = n_elems
+
+            if not inventory_only:
+                logger.debug(f"Registering kernel {kernel_name}")
+                desc, impl = rvv_vec_dot_product_kernels(
+                    n_elems, n_lanes, d_dtype, w_dtype, o_dtype, lmul
+                )
+                TensorIntrin.register(kernel_name, desc, impl, override=True)
+
+            n_elems //= 2
+
+    return kernels_inventory
+
+
+def register_riscv_intrinsics(target: Target):
+    """Register RISCV intrinsics
+
+    Args:
+        target (Target): TVM target
+    """
+
+    # RISCV `v` 1.0 extension templates
+    _ = register_rvv_isa_intrinsics(target)
+    logger.debug("Finished registering riscv intrinsics.")

python/tvm/s_tir/tensor_intrin/rocm.py

@@ -0,0 +1,477 @@
+# 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
+"""Intrinsics for AMDGPU tensorization."""
+
+from tvm.runtime import convert
+from tvm.script import tirx as T
+from tvm.tirx.expr import Cast, IntImm
+
+from .. import TensorIntrin
+from .dot_product_common import get_dp4a_intrin
+
+lift = convert
+
+
+@T.prim_func(s_tir=True)
+def sdot4(
+    A: T.Buffer((4,), "int8", offset_factor=1, align=4, scope="shared"),
+    B: T.Buffer((4,), "int8", offset_factor=1, align=4, scope="shared"),
+    C: T.Buffer((1,), "int32", offset_factor=1, align=4, scope="local"),
+) -> None:
+    with T.sblock("root"):
+        T.reads(C[0], A[0:4], B[0:4])
+        T.writes(C[0])
+
+        C[0] += T.call_llvm_pure_intrin(
+            T.llvm_lookup_intrinsic_id("llvm.amdgcn.sdot4"),
+            T.reinterpret(A.vload([0], "int8x4"), dtype="int32"),
+            T.reinterpret(B.vload([0], "int8x4"), dtype="int32"),
+            T.int32(0),
+            T.bool(1),
+            dtype="int32",
+        )
+
+
+AMDGPU_SDOT4_INTRIN = "sdot4"
+
+dp4a_desc, _ = get_dp4a_intrin("int8", "int8", "int32")
+TensorIntrin.register(AMDGPU_SDOT4_INTRIN, dp4a_desc, sdot4)
+
+WARP_SIZE = 64
+M_DIM = 16
+N_DIM = 16
+
+
+def shared_16x4_to_local_64x1_layout_A(i, j):
+    thread_id = j * 16 + i
+    return thread_id, convert(0)
+
+
+def thread_id_shared_access_64x1_to_16x4_layout_A(thread_id, local_id):
+    i = thread_id % 16
+    j = thread_id // 16 + local_id
+    return i, j
+
+
+def shared_4x16_to_local_64x1_layout_B(i, j):
+    thread_id = i * 16 + j
+    return thread_id, convert(0)
+
+
+def thread_id_shared_access_64x1_to_4x16_layout_B(thread_id, local_id):
+    i = thread_id // 16
+    j = thread_id % 16 + local_id
+    return i, j
+
+
+def shared_16x16_to_local_64x4_layout_C(i, j):
+    thread_id = j + (i // 4) * 16
+    local = i % 4
+    return thread_id, local
+
+
+def thread_id_shared_access_64x4_to_16x16_layout_A(thread_id, local_id):
+    i = thread_id % 16
+    j = (thread_id // 16) * 4 + local_id
+    return i, j
+
+
+def shared_16x16_to_local_64x4_layout_A(i, j):
+    thread_id = i + 16 * (j // 4)
+    local = j % 4
+    return thread_id, local
+
+
+def thread_id_shared_access_64x4_to_16x16_layout_B(thread_id, local_id):
+    i = local_id + (thread_id // 16) * 4
+    j = thread_id % 16
+    return i, j
+
+
+def shared_16x16_to_local_64x4_layout_B(i, j):
+    thread_id = j + (i // 4) * 16
+    local = i % 4
+    return thread_id, local
+
+
+def thread_id_shared_access_64x4_to_16x16_layout_C(thread_id, local_id):
+    i = local_id + (thread_id // 16) * 4
+    j = thread_id % 16
+    return i, j
+
+
+def get_mma_fill_intrin(dtype, local_size):
+    zero = IntImm("int32", 0).astype(dtype)
+
+    # Assume M = N = 16
+    index_map = shared_16x16_to_local_64x4_layout_C
+
+    @T.prim_func(s_tir=True)
+    def mma_fill_desc(a: T.handle) -> None:
+        C_warp = T.match_buffer(a, [WARP_SIZE, local_size], dtype=dtype, scope="warp")
+
+        with T.sblock("root"):
+            T.reads()
+            T.writes(C_warp[0:WARP_SIZE, 0:local_size])
+            for i0, i1 in T.grid(M_DIM, N_DIM):
+                with T.sblock("C_warp"):
+                    i, j = T.axis.remap("SS", [i0, i1])
+                    thread_id, local_id = T.meta_var(index_map(i, j))
+                    T.reads()
+                    T.writes(C_warp[thread_id, local_id])
+                    C_warp[thread_id, local_id] = zero
+
+    @T.prim_func(s_tir=True)
+    def mma_fill_impl(a: T.handle) -> None:
+        C_warp = T.match_buffer(
+            a, [WARP_SIZE, local_size], dtype=dtype, scope="warp", offset_factor=1
+        )
+
+        with T.sblock("root"):
+            T.reads()
+            T.writes(C_warp[0:WARP_SIZE, 0:local_size])
+            tx = T.env_thread("threadIdx.x")
+            T.launch_thread(tx, WARP_SIZE)
+            for local_id in T.serial(0, local_size):
+                C_warp[tx, local_id] = zero
+
+    return mma_fill_desc, mma_fill_impl
+
+
+def get_mfma_load_intrin(
+    k_dim=4,
+    dtype="float32",
+    scope="shared",
+    is_b=False,
+    transposed=False,
+):
+    local_size = (M_DIM * k_dim) // WARP_SIZE if not is_b else (N_DIM * k_dim) // WARP_SIZE
+    memory_shape = (M_DIM, k_dim)
+    if is_b:
+        memory_shape = (N_DIM, k_dim) if transposed else (k_dim, N_DIM)
+
+    row_dim, col_dim = memory_shape
+
+    if k_dim == 4:
+        index_map = shared_16x4_to_local_64x1_layout_A
+        reverse_index_map = thread_id_shared_access_64x1_to_16x4_layout_A
+        if is_b:
+            index_map = (
+                shared_16x4_to_local_64x1_layout_A
+                if transposed
+                else shared_4x16_to_local_64x1_layout_B
+            )
+            reverse_index_map = (
+                thread_id_shared_access_64x1_to_16x4_layout_A
+                if transposed
+                else thread_id_shared_access_64x1_to_4x16_layout_B
+            )
+    elif k_dim == 16:
+        index_map = shared_16x16_to_local_64x4_layout_A
+        reverse_index_map = thread_id_shared_access_64x4_to_16x16_layout_A
+
+        if is_b:
+            index_map = (
+                shared_16x16_to_local_64x4_layout_A
+                if transposed
+                else shared_16x16_to_local_64x4_layout_B
+            )
+            reverse_index_map = (
+                thread_id_shared_access_64x4_to_16x16_layout_A
+                if transposed
+                else thread_id_shared_access_64x4_to_16x16_layout_B
+            )
+    else:
+        raise ValueError("k_dim must be 4 or 16 currently")
+
+    @T.prim_func(s_tir=True)
+    def mfma_load_desc(reg_handle: T.handle, memory_handle: T.handle) -> None:
+        memory = T.match_buffer(
+            memory_handle,
+            memory_shape,
+            dtype,
+            offset_factor=1,
+            scope=scope,
+        )
+        reg = T.match_buffer(
+            reg_handle, (WARP_SIZE, local_size), dtype, offset_factor=1, scope="warp"
+        )
+
+        with T.sblock("root"):
+            T.reads(memory[0:row_dim, 0:col_dim])
+            T.writes(reg[0:WARP_SIZE, 0:local_size])
+
+            for ax0, ax1 in T.grid(row_dim, col_dim):
+                with T.sblock("memory_reg"):
+                    v0, v1 = T.axis.remap("SS", [ax0, ax1])
+                    T.reads(memory[v0, v1])
+
+                    thread_id, local_id = T.meta_var(index_map(v0, v1))
+                    T.writes(reg[thread_id, local_id])
+                    reg[thread_id, local_id] = memory[v0, v1]
+
+    @T.prim_func(s_tir=True)
+    def mfma_load_impl(reg_handle: T.handle, memory_handle: T.handle) -> None:
+        s0 = T.int32()
+        s1 = T.int32()
+
+        memory = T.match_buffer(
+            memory_handle,
+            memory_shape,
+            dtype,
+            align=64,
+            offset_factor=1,
+            scope=scope,
+            strides=[s0, s1],
+        )
+        reg = T.match_buffer(
+            reg_handle, (WARP_SIZE, local_size), dtype, align=64, offset_factor=1, scope="warp"
+        )
+
+        with T.sblock("root"):
+            T.reads(memory[0:row_dim, 0:col_dim])
+            T.writes(reg[0:WARP_SIZE, 0:local_size])
+            tx = T.env_thread("threadIdx.x")
+            for local_id in T.serial(0, local_size):
+                row, col = T.meta_var(reverse_index_map(tx, local_id))
+                T.launch_thread(tx, WARP_SIZE)
+                reg[tx, local_id] = memory[row, col]
+
+    return mfma_load_desc, mfma_load_impl
+
+
+def get_mfma_intrin(k_dim, in_dtype="float32", out_dtype="float32", b_transposed=False):
+    local_size = (M_DIM * k_dim) // WARP_SIZE
+    local_size_out = (M_DIM * N_DIM) // WARP_SIZE
+    if k_dim == 4:
+        index_map_A = shared_16x4_to_local_64x1_layout_A
+        index_map_B = shared_4x16_to_local_64x1_layout_B
+        index_map_C = shared_16x16_to_local_64x4_layout_C
+    elif k_dim == 16:
+        index_map_A = shared_16x16_to_local_64x4_layout_A
+        index_map_B = shared_16x16_to_local_64x4_layout_B
+        index_map_C = shared_16x16_to_local_64x4_layout_C
+    else:
+        raise ValueError("k_dim must be 4 or 16 currently")
+
+    out_dtype_abbrv = {"float16": "f16", "float32": "f32", "int8": "i8", "int32": "i32"}[out_dtype]
+
+    in_dtype_abbrv = {"float16": "f16", "float32": "f32", "int8": "i8", "int32": "i32"}[in_dtype]
+
+    mfma_intrin = f"llvm.amdgcn.mfma.{out_dtype_abbrv}.{M_DIM}x{N_DIM}x{k_dim}{in_dtype_abbrv}"
+
+    def maybe_cast(v):
+        if out_dtype != in_dtype:
+            return Cast(out_dtype, v)
+        return v
+
+    def maybe_swap(i, j):
+        if b_transposed:
+            return j, i
+        return i, j
+
+    @T.prim_func(s_tir=True)
+    def mfma_sync_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
+        B = T.match_buffer(b, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
+        C = T.match_buffer(c, (WARP_SIZE, local_size_out), out_dtype, offset_factor=1, scope="warp")
+
+        with T.sblock("root"):
+            T.reads(
+                C[0:WARP_SIZE, 0:local_size_out],
+                A[0:WARP_SIZE, 0:local_size],
+                B[0:WARP_SIZE, 0:local_size],
+            )
+            T.writes(C[0:WARP_SIZE, 0:local_size_out])
+
+            for i, j, k in T.grid(M_DIM, N_DIM, k_dim):
+                with T.sblock("C"):
+                    vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                    b_row_ind, b_col_ind = T.meta_var(maybe_swap(vk, vj))
+
+                    thread_id_C, local_id_C = T.meta_var(index_map_C(vi, vj))
+                    thread_id_A, local_id_A = T.meta_var(index_map_A(vi, vk))
+                    thread_id_B, local_id_B = T.meta_var(index_map_B(b_row_ind, b_col_ind))
+
+                    T.reads(
+                        C[thread_id_C, local_id_C],
+                        A[thread_id_A, local_id_A],
+                        B[thread_id_B, local_id_B],
+                    )
+                    T.writes(C[thread_id_C, local_id_C])
+
+                    C[thread_id_C, local_id_C] += maybe_cast(
+                        A[thread_id_A, local_id_A]
+                    ) * maybe_cast(B[thread_id_B, local_id_B])
+
+    @T.prim_func(s_tir=True)
+    def mfma_sync_impl_float(a: T.handle, b: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
+        B = T.match_buffer(b, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
+        C = T.match_buffer(c, (WARP_SIZE, local_size_out), out_dtype, offset_factor=1, scope="warp")
+
+        with T.sblock("root"):
+            T.reads(
+                A[0:WARP_SIZE, 0:local_size],
+                B[0:WARP_SIZE, 0:local_size],
+                C[0:WARP_SIZE, 0:local_size_out],
+            )
+            T.writes(C[0:WARP_SIZE, 0:local_size_out])
+            tx = T.env_thread("threadIdx.x")
+            T.launch_thread(tx, WARP_SIZE)
+            C[tx, 0:local_size_out] = T.call_llvm_pure_intrin(
+                T.llvm_lookup_intrinsic_id(mfma_intrin),
+                A[tx, 0:local_size],
+                B[tx, 0:local_size],
+                C[tx, 0:local_size_out],
+                T.int32(0),
+                T.int32(0),
+                T.int32(0),
+                dtype=f"{out_dtype}x4",
+            )
+
+    @T.prim_func(s_tir=True)
+    def mfma_sync_impl_integer(a: T.handle, b: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
+        B = T.match_buffer(b, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
+        C = T.match_buffer(c, (WARP_SIZE, local_size_out), out_dtype, offset_factor=1, scope="warp")
+
+        with T.sblock("root"):
+            T.reads(
+                A[0:WARP_SIZE, 0:local_size],
+                B[0:WARP_SIZE, 0:local_size],
+                C[0:WARP_SIZE, 0:local_size_out],
+            )
+            T.writes(C[0:WARP_SIZE, 0:local_size_out])
+            tx = T.env_thread("threadIdx.x")
+            T.launch_thread(tx, WARP_SIZE)
+
+            C[tx, 0:local_size_out] = T.call_llvm_pure_intrin(
+                T.llvm_lookup_intrinsic_id(mfma_intrin),
+                T.call_intrin("int32", "tirx.reinterpret", A[tx, 0:local_size]),
+                T.call_intrin("int32", "tirx.reinterpret", A[tx, 0:local_size]),
+                C[tx, 0:local_size_out],
+                T.int32(0),
+                T.int32(0),
+                T.int32(0),
+                dtype=f"{out_dtype}x4",
+            )
+
+    return (
+        (mfma_sync_desc, mfma_sync_impl_integer)
+        if in_dtype == "int8"
+        else (mfma_sync_desc, mfma_sync_impl_float)
+    )
+
+
+def get_mfma_store_intrin(local_size=4, dtype="float32", scope="global"):
+    index_map = shared_16x16_to_local_64x4_layout_C
+
+    @T.prim_func(s_tir=True)
+    def mfma_store_desc(a: T.handle, c: T.handle) -> None:
+        C_warp = T.match_buffer(a, [WARP_SIZE, local_size], dtype=dtype, scope="warp")
+        C = T.match_buffer(c, [M_DIM, N_DIM], dtype=dtype, scope=scope)
+
+        with T.sblock("root"):
+            T.reads(C_warp[0:WARP_SIZE, 0:local_size])
+            T.writes(C[0:M_DIM, 0:N_DIM])
+            for i0, i1 in T.grid(M_DIM, N_DIM):
+                with T.sblock("C_warp"):
+                    v0, v1 = T.axis.remap("SS", [i0, i1])
+                    thread_id, local_id = T.meta_var(index_map(v0, v1))
+                    T.reads(C_warp[thread_id, local_id])
+                    T.writes(C[v0, v1])
+                    C[v0, v1] = C_warp[thread_id, local_id]
+
+    @T.prim_func(s_tir=True)
+    def mfma_store_impl(a: T.handle, c: T.handle) -> None:
+        s0 = T.int32()
+        s1 = T.int32()
+
+        C_warp = T.match_buffer(
+            a, [WARP_SIZE, local_size], dtype=dtype, scope="warp", offset_factor=1
+        )
+        C = T.match_buffer(
+            c, [M_DIM, N_DIM], dtype=dtype, scope=scope, offset_factor=1, strides=[s0, s1]
+        )
+
+        with T.sblock("root"):
+            T.reads(C_warp[0:WARP_SIZE, 0:local_size])
+            T.writes(C[0:M_DIM, 0:N_DIM])
+            tx = T.env_thread("threadIdx.x")
+            T.launch_thread(tx, WARP_SIZE)
+            for i in range(local_size):
+                C[((tx // 16) * 4) + i, (tx % 16)] = C_warp[tx, i]
+
+    return mfma_store_desc, mfma_store_impl
+
+
+ROCM_MFMA_fill_16x16_f32_INTRIN = "ROCM_mfma_fill_16x16_f32"
+TensorIntrin.register(ROCM_MFMA_fill_16x16_f32_INTRIN, *get_mma_fill_intrin("float32", 4))
+
+ROCM_MFMA_fill_16x16_i32_INTRIN = "ROCM_mfma_fill_16x16_i32"
+TensorIntrin.register(ROCM_MFMA_fill_16x16_i32_INTRIN, *get_mma_fill_intrin("int", 4))
+
+ROCM_MFMA_LOAD_16x16_A_SHARED_s8_INTRIN = "rocm_mfma_load_16x16_a_shared_s8"
+TensorIntrin.register(
+    ROCM_MFMA_LOAD_16x16_A_SHARED_s8_INTRIN, *get_mfma_load_intrin(16, "int8", "shared")
+)
+ROCM_MFMA_LOAD_16x16_B_SHARED_s8_INTRIN = "rocm_mfma_load_b_16x16_shared_s8"
+TensorIntrin.register(
+    ROCM_MFMA_LOAD_16x16_B_SHARED_s8_INTRIN, *get_mfma_load_intrin(16, "int8", "shared", is_b=True)
+)
+
+ROCM_MFMA_LOAD_16x16_A_SHARED_f16_INTRIN = "rocm_mfma_load_16x16_a_shared_f16"
+TensorIntrin.register(
+    ROCM_MFMA_LOAD_16x16_A_SHARED_f16_INTRIN, *get_mfma_load_intrin(16, "float16", "shared")
+)
+ROCM_MFMA_LOAD_16x16_B_SHARED_f16_INTRIN = "rocm_mfma_load_b_16x16_shared_f16"
+TensorIntrin.register(
+    ROCM_MFMA_LOAD_16x16_B_SHARED_f16_INTRIN,
+    *get_mfma_load_intrin(16, "float16", "shared", is_b=True),
+)
+
+ROCM_MFMA_LOAD_16x4_A_SHARED_f32_INTRIN = "rocm_mfma_load_16x4_a_shared_f32"
+TensorIntrin.register(
+    ROCM_MFMA_LOAD_16x4_A_SHARED_f32_INTRIN, *get_mfma_load_intrin(4, "float32", "shared")
+)
+ROCM_MFMA_LOAD_16x4_B_SHARED_f32_INTRIN = "rocm_mfma_load_b_16x4_shared_f32"
+TensorIntrin.register(
+    ROCM_MFMA_LOAD_16x4_B_SHARED_f32_INTRIN,
+    *get_mfma_load_intrin(4, "float32", "shared", is_b=True),
+)
+
+
+ROCM_MFMA_f32f32f32_INTRIN = "rocm_mfma_f32f32f32"
+TensorIntrin.register(ROCM_MFMA_f32f32f32_INTRIN, *get_mfma_intrin(4, "float32", "float32"))
+
+ROCM_MFMA_f16f16f32_INTRIN = "rocm_mfma_f16f16f32"
+TensorIntrin.register(ROCM_MFMA_f16f16f32_INTRIN, *get_mfma_intrin(16, "float16", "float32"))
+
+ROCM_MFMA_s8s8s32_INTRIN = "rocm_mfma_s8s8s32"
+TensorIntrin.register(ROCM_MFMA_s8s8s32_INTRIN, *get_mfma_intrin(16, "int8", "int32"))
+
+ROCM_MFMA_STORE_16x16_s32_INTRIN = "rocm_mfma_store_16x16_s32"
+TensorIntrin.register(
+    ROCM_MFMA_STORE_16x16_s32_INTRIN, *get_mfma_store_intrin(4, "int32", "global")
+)
+
+ROCM_MFMA_STORE_16x16_f32_INTRIN = "rocm_mfma_store_16x16_f32"
+TensorIntrin.register(
+    ROCM_MFMA_STORE_16x16_f32_INTRIN, *get_mfma_store_intrin(4, "float32", "global")
+)

python/tvm/s_tir/tensor_intrin/x86.py

@@ -0,0 +1,112 @@
+# 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
+"""Intrinsics for x86 tensorization."""
+
+from tvm.script import tirx as T
+
+from .. import TensorIntrin
+
+# Tensorized intrinsic description and VNNI-specific implementation.
+# Equivalent to the ones in topi/x86/tensor_intrin.py
+
+
+@T.prim_func(s_tir=True)
+def dot_product_16x4_u8i8i32_desc(
+    A: T.Buffer((4,), "uint8", offset_factor=1),
+    B: T.Buffer((16, 4), "int8", offset_factor=1),
+    C: T.Buffer((16,), "int32", offset_factor=1),
+) -> None:
+    with T.sblock("root"):
+        T.reads(C[0:16], A[0:4], B[0:16, 0:4])
+        T.writes(C[0:16])
+        for i in T.serial(0, 16):
+            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 dot_product_16x4_u8i8i32_vnni(
+    A: T.Buffer((4,), "uint8", offset_factor=1),
+    B: T.Buffer((16, 4), "int8", offset_factor=1),
+    C: T.Buffer((16,), "int32", offset_factor=1),
+) -> None:
+    with T.sblock("root"):
+        T.reads(C[0:16], A[0:4], B[0:16, 0:4])
+        T.writes(C[0:16])
+
+        A_u8x4 = A.vload([0], "uint8x4")
+        A_i32 = T.reinterpret(A_u8x4, dtype="int32")
+
+        B_i8x64 = B.vload([0, 0], dtype="int8x64")
+        B_i32x16 = T.reinterpret(B_i8x64, dtype="int32x16")
+        C_i32x16 = C.vload([0], dtype="int32x16")
+
+        C[T.ramp(T.int32(0), 1, 16)] = T.call_llvm_pure_intrin(
+            T.llvm_lookup_intrinsic_id("llvm.x86.avx512.vpdpbusd.512"),
+            C_i32x16,
+            T.broadcast(A_i32, 16),
+            B_i32x16,
+            dtype="int32x16",
+        )
+
+
+@T.prim_func(s_tir=True)
+def dot_product_16x4_u8i8i32_avx512(
+    A: T.Buffer((4,), "uint8", offset_factor=1),
+    B: T.Buffer((16, 4), "int8", offset_factor=1),
+    C: T.Buffer((16,), "int32", offset_factor=1),
+) -> None:
+    with T.sblock("root"):
+        T.reads(C[0:16], A[0:4], B[0:16, 0:4])
+        T.writes(C[0:16])
+
+        A_u8x4 = A.vload([0], "uint8x4")
+        A_i32 = T.reinterpret(A_u8x4, dtype="int32")
+        A_brdcst = T.broadcast(A_i32, 16)
+        A_u8x64 = T.reinterpret(A_brdcst, dtype="uint8x64")
+
+        B_i8x64 = B.vload([0, 0], dtype="int8x64")
+
+        Red = T.call_llvm_pure_intrin(
+            T.llvm_lookup_intrinsic_id("llvm.x86.avx512.pmaddubs.w.512"),
+            A_u8x64,
+            B_i8x64,
+            dtype="int16x32",
+        )
+
+        C[T.ramp(T.int32(0), 1, 16)] += T.call_llvm_pure_intrin(
+            T.llvm_lookup_intrinsic_id("llvm.x86.avx512.pmaddw.d.512"),
+            Red,
+            T.int16x32(1),
+            dtype="int32x16",
+        )
+
+
+VNNI_DOT_16x4_INTRIN = "dot_16x4_vnni"
+
+TensorIntrin.register(
+    VNNI_DOT_16x4_INTRIN, dot_product_16x4_u8i8i32_desc, dot_product_16x4_u8i8i32_vnni
+)
+
+AVX512_DOT_16x4_INTRIN = "dot_16x4_avx512"
+
+TensorIntrin.register(
+    AVX512_DOT_16x4_INTRIN, dot_product_16x4_u8i8i32_desc, dot_product_16x4_u8i8i32_avx512
+)