TheRock/rocm-libraries/projects/hipblaslt/tensilelite/Tensile/Components/MAC_BF16_HPA.py

################################################################################
#
# Copyright (C) 2022-2026 Advanced Micro Devices, Inc. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
################################################################################
from rocisa.code import Module
from rocisa.container import vgpr
from rocisa.instruction import SSetPrior, VDot2F32BF16, VDot2CF32BF16

from rocisa.enum import DataTypeEnum
from ..Common.DataType import DataType
from ..Component import Component, MAC

class FMA_BF16_HPA(MAC):
    asmCaps = {"v_fma_f32": True}
    kernel = {"ProblemType": {"MacDataTypeA": DataType(DataTypeEnum.BFloat16),
                              "MacDataTypeB": DataType(DataTypeEnum.BFloat16),
                              "HighPrecisionAccumulate": True},
              "UseDotInstruction": False,
              }

    def __call__(self, writer, m, innerUnroll):
        kernel = writer.states.kernel
        module = Module("FMA_BF16_HPA")
        module.addComment(self.commentHeader())
        priority = Component.Priority.find(writer)

        vars = {}
        vars["m"] = m
        vars["ThreadTile0"] = kernel["ThreadTile0"]
        vars["ThreadTile0Half"] = kernel["ThreadTile0"] // 2

        for iui in range(0, innerUnroll):
            vars["iui"] = iui
            for blockA in range(kernel["ThreadTileA"]//2-1, -1, -1):
                vars["blockA"] = blockA
                dst = "v[vgprValuA_X{m}_I{iui}+{blockA}*2+1]".format_map(vars)
                src = "v[vgprValuA_X{m}_I{iui}+{blockA}]".format_map(vars)
                module.addInst("v_and_b32", dst, "0xffff0000", src, "")
                dst = "v[vgprValuA_X{m}_I{iui}+{blockA}*2]".format_map(vars)
                module.addInst("v_lshlrev_b32", dst, 16, src, "")
            for blockB in range(kernel["ThreadTileB"]//2-1, -1, -1):
                vars["blockB"] = blockB
                dst = "v[vgprValuB_X{m}_I{iui}+{blockB}*2+1]".format_map(vars)
                src = "v[vgprValuB_X{m}_I{iui}+{blockB}]".format_map(vars)
                module.addInst("v_and_b32", dst, "0xffff0000", src, "")
                dst = "v[vgprValuB_X{m}_I{iui}+{blockB}*2]".format_map(vars)
                module.addInst("v_and_b32", dst, 16, src, "")

        for block1 in range(0, kernel["ThreadTile1"]//2):
            vars["block1"] = block1
            for block0 in range(0, kernel["ThreadTile0"]//2):
                vars["block0"] = block0
                if kernel["ProblemType"]["HighPrecisionAccumulate"]:
                    # we treat HighPrecisionAccumulate as expanded packed math
                    for iui in range(0, innerUnroll):
                        vars["iui"] = iui

                        vars["blockA"] = block0 if writer.tPB["tile01Idx"] else block1
                        vars["blockB"] = block1 if writer.tPB["tile01Idx"] else block0

                        aStr0 = "v[vgprValuA_X{m}_I{iui}+{blockA}*2+0]".format_map(vars)
                        aStr1 = "v[vgprValuA_X{m}_I{iui}+{blockA}*2+1]".format_map(vars)
                        bStr0 = "v[vgprValuB_X{m}_I{iui}+{blockB}*2+0]".format_map(vars)
                        bStr1 = "v[vgprValuB_X{m}_I{iui}+{blockB}*2+1]".format_map(vars)

                        cidx = block0*2 + block1*kernel["ThreadTile0"]*2 + 0
                        cStr = "v[vgprValuC+{block0}*2+{block1}*{ThreadTile0}*2+0*2+0]".format_map(vars)
                        module.addInst("v_fma_f32", cStr, aStr0, bStr0, cStr, "ValuC[%u]" % cidx)

                        module.add(priority(writer, 1, "Raise priority while processing macs"))

                        aStr = aStr1 if writer.tPB["tile01Idx"] else aStr0
                        bStr = bStr0 if writer.tPB["tile01Idx"] else bStr1
                        cidx = block0*2 + block1*kernel["ThreadTile0"]*2 + 1
                        cStr = "v[vgprValuC+{block0}*2+{block1}*{ThreadTile0}*2+0*2+1]".format_map(vars) # *2 b/c of fp32
                        module.addInst("v_fma_f32", cStr, aStr, bStr, cStr, "ValuC[%u]" % cidx)

                        aStr = aStr0 if writer.tPB["tile01Idx"] else aStr1
                        bStr = bStr1 if writer.tPB["tile01Idx"] else bStr0
                        cidx = block0*2 + block1*kernel["ThreadTile0"]*2 + kernel["ThreadTile0"] + 0
                        cStr = "v[vgprValuC+{block0}*2+{block1}*{ThreadTile0}*2+{ThreadTile0Half}*2+0]".format_map(vars)
                        module.addInst("v_fma_f32", cStr, aStr, bStr, cStr, "ValuC[%u]" % cidx)

                        cidx = block0*2 + block1*kernel["ThreadTile0"]*2 + kernel["ThreadTile0"] + 1
                        cStr = "v[vgprValuC+{block0}*2+{block1}*{ThreadTile0}*2+{ThreadTile0Half}*2+1]".format_map(vars)
                        module.addInst("v_fma_f32", cStr, aStr1, bStr1, cStr, "ValuC[%u]" % cidx)
                        """
                        ignore this, not quite correct for mixed precision
                        D.f[31:16] = S0.f[31:16] * S1.f[31:16] + S2.f[31:16]
                        D.f[15:00] = S0.f[15:00] * S1.f[15:00] + S2.f[15:00]
                        C[0] = A[0]*B[0]+D[0]
                        C[1] = A[1]*B[1]+D[1]
                        """
                        #module.add(self.getBomb(-13))

        module.add(priority(writer, 0, "Reset priority after macs"))
        return module

#dot2
class FMA_BF16_HPA_DOT2(MAC):
    asmCaps = lambda caps: caps['v_dot2_f32_bf16'] or caps['v_dot2c_f32_bf16']
    #archCaps = {}
    kernel = {"ProblemType": {"MacDataTypeA": DataType(DataTypeEnum.BFloat16),
                              "MacDataTypeB": DataType(DataTypeEnum.BFloat16),
                              "HighPrecisionAccumulate": True},
              "UseDotInstruction": True,
             }

    def __call__(self, writer, tPA, tPB, m, innerUnroll):
        kernel = writer.states.kernel

        module = Module("FMA_BF16_HPA_DOT2")
        module.addComment(self.commentHeader())

        vars = {}

        if writer.states.asmCaps["v_dot2_f32_bf16"]:
            instruction = VDot2F32BF16
        else:
            instruction = VDot2CF32BF16

        vars["m"] = m
        vars["kernel"] = kernel

        vars["ThreadTile0"] = kernel["ThreadTile0"]
        vars["ThreadTile1"] = kernel["ThreadTile1"]

        for block1 in range(0, kernel["ThreadTile1"]):
            for block0 in range(0, kernel["ThreadTile0"]):
                for iui in range(0, innerUnroll):
                    vars["block0"] = block0
                    vars["block1"] = block1
                    vars["blockA"] = block0 if tPA["tileIdx"] == 0 else block1
                    vars["blockB"] = block1 if tPB["tileIdx"] != 0 else block0
                    vars["iui"] = iui

                    vars["cIdxExpr"] = "%d+%d" % (vars["block0"], vars["block1"]*vars["ThreadTile0"])
                    cidx = eval(vars["cIdxExpr"])
                    cStr = "ValuC+{cIdxExpr}".format_map(vars)
                    aStr = "ValuA_X{m}_I{iui}+{blockA}".format_map(vars)
                    bStr = "ValuB_X{m}_I{iui}+{blockB}".format_map(vars)
                    if instruction == VDot2F32BF16:
                        module.add(instruction(dst=vgpr(cStr), src0=vgpr(aStr), src1=vgpr(bStr), src2=vgpr(cStr), \
                                            comment="ValuC[%u] iui=%u" % (cidx, vars["iui"])))
                    else:
                        module.add(instruction(dst=vgpr(cStr), src0=vgpr(aStr), src1=vgpr(bStr), \
                                            comment="ValuC[%u] iui=%u" % (cidx, vars["iui"])))

                    if (block1 == 0) and (block0 == 0) and (iui == 0):
                        module.add(SSetPrior(prior=1, comment="Raise priority while processing macs"))

        module.add(SSetPrior(prior=0, comment="Reset priority after macs"))

        return module