apache--tvm/python/tvm/s_tir/tensor_intrin/riscv_cpu.py

# 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.")