TheRock/rocm-libraries/shared/tensile/Tensile/Contractions.py

################################################################################
#
# Copyright (C) 2019-2024 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 .DataType import DataType
from . import Hardware
from . import Properties
from .SolutionStructs import Solution as OriginalSolution
from .Utils import state, state_key_ordering

from . import Common

@state_key_ordering
class FreeIndex:
    StateKeys = ['isA', 'i', 'c', 'd']

    def __init__(self, isA, i=None, c=None, d=None):
        self.isA = isA
        self.i = i # index into A or B (depending on isA)
        self.c = c
        self.d = d

@state_key_ordering
class BatchIndex:
    StateKeys = ['a', 'b', 'c', 'd']
    def __init__(self, a=None, b=None, c=None, d=None):
        self.a = a
        self.b = b
        self.c = c
        self.d = d

@state_key_ordering
class BoundIndex:
    StateKeys = ['a', 'b', 'aMirror', 'bMirror']
    def __init__(self, a=None, b=None, aMirror=False, bMirror=False):
        self.a = a
        self.b = b
        self.aMirror = aMirror
        self.bMirror = bMirror


class ProblemType:
    StateKeys = ['operationIdentifier', 'aType', 'bType', 'cType', 'dType',
                 'useBeta', 'highPrecisionAccumulate', 'useInitialStridesAB', 'useInitialStridesCD', 'stridedBatched', 'f32XdlMathOp', 'stochasticRounding']
    @classmethod
    def FromOriginalState(cls, d):
        indices = [None]*d['TotalIndices']
        freeIndices  = []
        batchIndices = []
        boundIndices = []

        for i in d['IndicesSummation']:
            bi = BoundIndex(aMirror=('MirrorDimsA' in d and i in d['MirrorDimsA']),
                            bMirror=('MirrorDimsB' in d and i in d['MirrorDimsB']))
            indices[i] = bi
            boundIndices.append(bi)

        for i in range(0,d["NumIndicesC"]):
            if i in d['IndicesBatch']:
                bi = BatchIndex(c=i, d=i)
                indices[i] = bi
                batchIndices.append(bi)
            else:
                assert i in d['IndicesFree']
                if i in d['IndexAssignmentsA']:
                    fi = FreeIndex(isA=True, i=d["IndexAssignmentsA"].index(i), c=i, d=i)
                elif i in d['IndexAssignmentsB']:
                    fi = FreeIndex(isA=False, i=d["IndexAssignmentsB"].index(i), c=i, d=i)
                else:
                    raise RuntimeError("free index %u not in ia or ib"%i)
                indices[i] = fi
                freeIndices.append(fi)

        for ia, ic in enumerate(d['IndexAssignmentsA']):
            indices[ic].a = ia

        for ib, ic in enumerate(d['IndexAssignmentsB']):
            indices[ic].b = ib

        for idx in indices:
            assert idx is not None
            idxState = state(idx)
            for (key, value) in list(idxState.items()):
                assert value is not None

        rv = cls()
        rv.indices = indices
        rv.freeIndices = freeIndices
        rv.batchIndices = batchIndices
        rv.boundIndices = boundIndices
        rv.aDims = len(d['IndexAssignmentsA'])
        rv.bDims = len(d['IndexAssignmentsB'])
        rv.cDims = d['NumIndicesC']
        rv.dDims = rv.cDims

        rv.aConjugate = d['ComplexConjugateA']
        rv.bConjugate = d['ComplexConjugateB']

        srcType = DataType(d['DataType'])
        dstType = DataType(d['DestDataType']) if 'DestDataType' in d else srcType
        computeType = DataType(d['ComputeDataType']) if 'ComputeDataType' in d else dstType

        rv.aType = srcType
        rv.bType = srcType
        # for hybrid 8bit float types, we need to split the type into a_type and b_type
        if srcType.isFloat8BFloat8():
            rv.aType = DataType("F8")
            rv.bType = DataType("B8")
        elif srcType.isBFloat8Float8():
            rv.aType = DataType("B8")
            rv.bType = DataType("F8")

        # We don't expect dstType and computeType to be hybrid types for now
        rv.cType = dstType
        rv.dType = dstType
        # we already checked the src/dst/compute types are supported and well-assigned in SolutionStruct
        rv.alphaType = computeType
        rv.betaType = computeType

        rv.highPrecisionAccumulate = False
        if 'HighPrecisionAccumulate' in d:
            rv.highPrecisionAccumulate = d['HighPrecisionAccumulate']

        rv.useInitialStridesAB = False
        if 'UseInitialStridesAB' in d:
            rv.useInitialStridesAB = d['UseInitialStridesAB']
        rv.useInitialStridesCD = False
        if 'UseInitialStridesCD' in d:
            rv.useInitialStridesCD = d['UseInitialStridesCD']

        rv.stridedBatched = True
        if 'StridedBatched' in d:
          rv.stridedBatched = d['StridedBatched']

        rv.setConstStrideA = []
        if 'SetConstStrideA' in d:
            rv.setConstStrideA = d['SetConstStrideA']
        rv.setConstStrideB = []
        if 'SetConstStrideB' in d:
            rv.setConstStrideB = d['SetConstStrideB']
        rv.zeroPadA=[]
        if 'ZeroPadA' in d:
            rv.zeroPadA = d['ZeroPadA']
        rv.zeroPadB=[]
        if 'ZeroPadB' in d:
            rv.zeroPadB = d['ZeroPadB']

        rv.mirrorDimsA = d['MirrorDimsA'] if 'MirrorDimsA' in d else []
        rv.mirrorDimsB = d['MirrorDimsB'] if 'MirrorDimsB' in d else []

        rv.useBeta = True
        if 'UseBeta' in d:
            rv.useBeta = d['UseBeta']

        rv.batched = d['Batched']

        rv.f32XdlMathOp = DataType(d['F32XdlMathOp']) if 'F32XdlMathOp' in d else DataType(0)

        rv.stochasticRounding = False
        if 'StochasticRounding' in d:
            rv.stochasticRounding = d['StochasticRounding']

        return rv

    def __init__(self, freeIndices=None, batchIndices=None, boundIndices=None, aDims=None, bDims=None, cDims=None, dDims=None):
        self.convolution = None
        self.freeIndices  = freeIndices
        self.batchIndices = batchIndices
        self.boundIndices = boundIndices
        self.aDims = aDims
        self.bDims = bDims
        self.cDims = cDims
        self.dDims = dDims

    @property
    def indexNames(self):
        aNames = ['_'] * self.aDims
        bNames = ['_'] * self.bDims
        cNames = ['_'] * self.cDims

        allIndexNames = 'ijklmnopqrstuvwxyz'
        index = 0

        dNames = list([allIndexNames[index+i] for i in range(self.cDims)])
        index += len(dNames)

        sumNames = list([allIndexNames[index+i] for i in range(len(self.boundIndices))])
        index += len(sumNames)

        for free in self.freeIndices:
            if free.isA:
                aNames[free.i ] = dNames[free.d]
            else:
                bNames[free.i ] = dNames[free.d]
            cNames[free.c] = dNames[free.d]

        for batch in self.batchIndices:
            name = dNames[batch.d]
            aNames[batch.a] = name
            bNames[batch.b] = name
            cNames[batch.c] = name

        for i, bound in enumerate(self.boundIndices):
            name = sumNames[i]
            aNames[bound.a] = name.upper() if bound.aMirror else name
            bNames[bound.b] = name.upper() if bound.bMirror else name

        aNames = ''.join(aNames)
        bNames = ''.join(bNames)
        cNames = ''.join(cNames)
        dNames = ''.join(dNames)
        sumNames = ''.join(sumNames)

        return (aNames, bNames, cNames, dNames, sumNames)

    @property
    def operationIdentifier(self):
        (aNames, bNames, cNames, dNames, sumNames) = self.indexNames

        aOp = 'C' if self.aConjugate else ''
        bOp = 'C' if self.bConjugate else ''

        return '_'.join(['Contraction', sumNames,
                         'A' + aNames + aOp,
                         'B' + bNames + bOp,
                         'C' + cNames,
                         'D' + dNames])

    def placeholderStr(self, includeBatch=False, includeOperation=False, includeType=False):
        ret = ""
        if includeOperation:
            ret = self.operationIdentifier
            if not self.useBeta:
                ret += "_Beta0"
            ret += "_StridedBatched{}".format(int(self.stridedBatched))
        if includeType:
            ret += "_Type_{}{}".format(DataType(self.aType).toChar(), DataType(self.cType).toChar())
            if self.highPrecisionAccumulate:
                ret += "_HPA"

        return ret

    def predicates(self, includeBatch=False, includeOperation=False, includeType=False):
        predicates = []

        #if includeBatch and not self.batched:
        #    predicates.append(ProblemPredicate("BatchSizeEqual", index=0, value=1))

        if includeOperation:
            predicates.append(ProblemPredicate("OperationIdentifierEqual", value=self.operationIdentifier))
            if not self.useBeta:
                predicates.append(ProblemPredicate("BetaZero"))
            predicates.append(ProblemPredicate("StridedBatched", value=self.stridedBatched))

        if includeType:
            predicates.append(ProblemPredicate("TypesEqual", value=(self.aType, self.bType, self.cType, self.dType)))
            predicates.append(ProblemPredicate("HighPrecisionAccumulate", value=self.highPrecisionAccumulate))
            predicates.append(ProblemPredicate("F32XdlMathOp", value=self.f32XdlMathOp))
            predicates.append(ProblemPredicate("StochasticRounding", value=self.stochasticRounding))


        return predicates

def extractDimPredicate(cls, key, value, predicateName):
    """
    Extract the predicate for AssertStrideEqual*
    Value is a dictionary
    """
    predicates = []
    for pos,val in value.items():
        if val != -1:
            predicates.append(cls(predicateName, index=pos, value=val))
    if len(predicates) == 1:
        return predicates[0]
    elif len(predicates) > 1:
        return cls.And(predicates)

class ProblemPredicate(Properties.Predicate):
    @classmethod
    def FromOriginalKeyPair(cls, pair):
        (key, value) = pair
        if key == "AssertStrideAEqual":
            return extractDimPredicate(cls, key, value, "StrideAEqual")
        if key == "AssertStrideBEqual":
            return extractDimPredicate(cls, key, value, "StrideBEqual")
        if key == "AssertStrideCEqual":
            return extractDimPredicate(cls, key, value, "StrideCEqual")
        if key == "AssertStrideDEqual":
            return extractDimPredicate(cls, key, value, "StrideDEqual")

        if key == "AssertSizeEqual":
            return extractDimPredicate(cls, key, value, "SizeEqual")
        if key == "AssertSizeGreaterThan":
            return extractDimPredicate(cls, key, value, "SizeGreaterThan")
        if key == "AssertSizeLessThan":
            return extractDimPredicate(cls, key, value, "SizeLessThan")
        if key == "AssertSizeMultiple":
            return extractDimPredicate(cls, key, value, "SizeMultiple")

        # Arithmetic intensity assertions
        if key == "AssertAIGreaterThanEqual":
            return cls("AIGreaterThanEqual", value=value) if value > 0 else None
        if key == "AssertAILessThanEqual":
            return cls("AILessThanEqual", value=value) if value > 0 else None

        # Alpha and beta value assertions
        if key == "AssertAlphaValue":
            return cls("AlphaValue", value=str(value)) if value != False else None
        if key == "AssertBetaValue":
            return cls("BetaValue", value=str(value)) if value != False else None

        if key == "AssertCEqualsD":
            return cls("CEqualsD") if value != False else None

        # TODO - remove this when logic files have been updated
        if key == 'AssertMinApproxSize':
            return None

        if key.endswith('Multiple'):
            if value == 1:
                return None

            if key == "AssertFree0ElementMultiple":
                tag = "Free0SizeMultiple"
                index = 0
            elif key == "AssertFree1ElementMultiple":
                tag = "Free1SizeMultiple"
                index = 0
            elif key == "AssertSummationElementMultiple":
                tag = "BoundSizeMultiple"
                index = -1
            else:
                raise RuntimeError("Unknown Multiple Value: {}".format(key))

            return cls(tag, index=index, value=value)

        if key.startswith('Assert'):
            raise RuntimeError("Unknown assertion key: {}".format(key))

        if key == "Fp16AltImpl":
            return cls("Fp16AltImpl") if value != False else None

        if key == "Fp16AltImplRound":
            return cls("Fp16AltImplRound") if value != False else None

    @classmethod
    def CompoundPredicates(cls, state, problemType):
        rv = []

        if not problemType.aType.isInt8x4():
            # calculate the minimum supported free dimension size
            TLUA = state['ProblemType']['TLUA']
            TLUB = state['ProblemType']['TLUB']
            minFree0 = max(state['GlobalLoadVectorWidthA'], 1) if TLUA else 1
            minFree1 = max(state['GlobalLoadVectorWidthB'], 1) if TLUB else 1
            rv += [cls('LeadingFree0SizesGreaterOrEqual', value=minFree0)]
            rv += [cls('LeadingFree1SizesGreaterOrEqual', value=minFree1)]

        if len(state["PackedC0IndicesX"]) > 1:
          rv += [cls("CDStridesEqual")]

        if "KernelLanguage" in state:
            rv += [cls("KernelLanguageCompatible", value=state["KernelLanguage"])]

        if ('GlobalSplitU' in state) and (state['GlobalSplitU'] > 1):
            if ('_GlobalAccumulation' not in state) or (state['_GlobalAccumulation'] != 'MultipleBuffer'):
                rv += [cls("DeterministicMode", value = False)]

        if ('StreamK' in state) and (state['StreamK'] > 0) and ('StreamKAtomic' in state) and (state['StreamKAtomic'] == 1):
            # StreamKAtomic = 1 uses atomic for partial tiles
            rv += [cls("DeterministicMode", value = False)]

        # debugging: mark this set to allow the problem always runnable with PK
        if 'PersistentKernel' in state and state['PersistentKernel']:
            rv += [cls("PersistentKernelCheck")]

        if ("MatrixInstruction" in state and state["MatrixInstruction"]) or \
           ("EnableMatrixInstruction" in state and state["EnableMatrixInstruction"] is True):
            rv += [cls("ArithmeticUnitCompatible", value="MFMA")]
        else:
            rv += [cls("ArithmeticUnitCompatible", value="VALU")]

        # if bufferload is performed, we output some predication info for host side,
        # to prevent from some extremely large problems from launching and causing bufferload offset limit < 2^32
        # those cases will not satisfy the assertion thus won't use the kernel.
        # See Common.py for more details, we will need four values:
        # TODO - haven't been fully tested for FP16 and BF, need to verify the false-positive
        if 'BufferLoad' in state and state['BufferLoad'] == True:
            TLUA = state['ProblemType']['TLUA']
            TLUB = state['ProblemType']['TLUB']
            MayShiftA = TLUA and state['AssertFree0ElementMultiple'] < state['GlobalLoadVectorWidthA']
            MayShiftB = TLUB and state['AssertFree1ElementMultiple'] < state['GlobalLoadVectorWidthB']
            subrv={}
            subrv['ShiftPtrElemB'] = state['GlobalLoadVectorWidthB'] if MayShiftB else 0
            subrv['ShiftPtrElemA'] = state['GlobalLoadVectorWidthA'] if MayShiftA else 0
            subrv['DUorMT1'] = state['DepthU'] if TLUB else state['MacroTile1']
            subrv['DUorMT0'] = state['DepthU'] if TLUA else state['MacroTile0']
            # value is also a dict for better readability, client side need to handel the serialization
            rv += [cls('BufferLoadOffsetLimitCheck', value=subrv)]

        # similar check is applied for bufferstore,
        # for bufferstore offset, test if the bot-right offset < 2^32,
        # it should be StrideA*MT1, so we need to output MT1 and use the StrideA of problem in host-side for predication
        if 'BufferStore' in state and state['BufferStore'] == True:
            val = state['MacroTile1']
            if (1 in state["AssertSizeLessThan"].keys()):
                # use smaller value if array size assert is enabled
                val = min(val, state["AssertSizeLessThan"][1] - 1)
            rv += [cls('BufferStoreOffsetLimitCheck', value=val)]

        if '_GlobalAccumulation' in state and state['_GlobalAccumulation'] != None and not state["StreamK"]:
            value = state['MinKForGSU'] * state['GlobalSplitU']
            rv += [cls('GlobalSplitUCheckMinK', value=value)]

        return rv

    @classmethod
    def FromOriginalState(cls, d, problemType, morePreds=[]):
        problemTypePreds = problemType.predicates(True, True, True)
        compoundPreds = cls.CompoundPredicates(d, problemType)
        extraPreds = problemTypePreds + compoundPreds + morePreds

        predicates = [p for p in map(cls.FromOriginalKeyPair, d.items()) if p is not None] + extraPreds
        return cls.And(predicates)

class TaskPredicate(Properties.Predicate):
    @classmethod
    def FromOriginalKeyPair(cls, pair):
        (key, value) = pair
        if key == "_WorkspaceSizePerElemC" and value > 0:
            return cls("WorkspaceCheck")
        return None

    @classmethod
    def FromOriginalState(cls, d, problemType, morePreds=[]):
        predicates = [p for p in map(cls.FromOriginalKeyPair, d.items()) if p is not None]
        return cls.And(predicates)

class SizeMapping:
    StateKeys = ['workGroup',
                 'macroTile',
                 'threadTile',
                 'depthU',
                 'staggerU',
                 'globalSplitU',
                 'staggerStrideShift',
                 'workGroupMapping',
                 'packSummationDims',
                 'packBatchDims',
                 'magicDivAlg',
                 'streamK',
                 'streamKAtomic',
                 'persistentKernel',
                 'persistentKernelAlongBatch',
                 'sourceKernel',
                 'globalAccumulation',
                 'workspaceSizePerElemC',
                 'preloadKernargs'
                 ]

    @classmethod
    def FromOriginalState(cls, d):
        globalAccum = 0
        if d['_GlobalAccumulation'] == 'SingleBuffer':
            globalAccum = 1
        if d['_GlobalAccumulation'] == 'MultipleBuffer':
            globalAccum = 2
        if d['_GlobalAccumulation'] == 'PartialsBuffer':
            globalAccum = 3

        assert d["PreloadKernelArguments"] in (0, 1), d["PreloadKernelArguments"]
        return cls(workGroup             = d['WorkGroup'],
                   macroTile             = cls.ReadOriginalMacroTile(d),
                   threadTile            = d['ThreadTile'],
                   workGroupMapping      = d['WorkGroupMapping'],
                   staggerU              = d['StaggerU'] if 'StaggerU' in d else 0,
                   depthU                = d['DepthU'],
                   globalSplitU          = d['GlobalSplitU'],
                   staggerStrideShift    = d['_staggerStrideShift'] if '_staggerStrideShift' in d else 0,
                   packSummationDims     = d['PackSummationDims'] if 'PackSummationDims' in d else 0,
                   packBatchDims         = d['PackBatchDims'] if 'PackBatchDims' in d else 0,
                   streamK               = d['StreamK'] if 'StreamK' in d else 0,
                   streamKAtomic         = d['StreamKAtomic'] if 'StreamKAtomic' in d else 0,
                   persistentKernel      = d['PersistentKernel'] if 'PersistentKernel' in d else 0,
                   persistentKernelAlongBatch   = d['PersistentKernelAlongBatch'] if 'PersistentKernelAlongBatch' in d else False,
                   magicDivAlg           = d.get('MagicDivAlg', 1),
                   sourceKernel          = d['KernelLanguage'] == 'Source',
                   globalAccumulation    = globalAccum,
                   workspaceSizePerElemC = d['_WorkspaceSizePerElemC'],
                   preloadKernargs       = d["PreloadKernelArguments"]
                   )

    @classmethod
    def ReadOriginalMacroTile(cls, d):
        rv = [1,1,1]
        rv[0] = d['MacroTile0']
        rv[1] = d['MacroTile1']
        return rv

    def __init__(self, **kwargs):
        for (key, value) in list(kwargs.items()):
            setattr(self, key, value)

class Solution:
    StateKeys = ['name',
                'problemType',
                'hardwarePredicate',
                'problemPredicate',
                'taskPredicate',
                'sizeMapping',
                'debugKernel',
                'libraryLogicIndex',
                'index',
                'ideals',
                'linearModel']
    HiddenKeys = ['originalSolution']

    @classmethod
    def FromSolutionStruct(cls, solution):
        return cls.FromOriginalState(solution._state)

    @classmethod
    def FromOriginalState(cls, d, deviceInfo=None):
        rv = cls()


        if 'SolutionNameMin' in d:
            rv.name = d['SolutionNameMin']

        rv.problemType = ProblemType.FromOriginalState(d['ProblemType'])

        rv.problemPredicate = ProblemPredicate.FromOriginalState(d, rv.problemType)
        rv.taskPredicate = TaskPredicate.FromOriginalState(d, rv.problemType)

        if 'DebugKernel' in d:
            rv.debugKernel = d['DebugKernel']

        if 'SolutionIndex' in d:
            rv.index = d['SolutionIndex']

        info = cls.ReadOriginalInfo(d)
        rv.libraryLogicIndex = int(info.get("SolutionIndex", -1))

        rv.sizeMapping = SizeMapping.FromOriginalState(d)
        if 'Ideals' in d:
            rv.ideals = d['Ideals']
        else:
            rv.ideals = {}

        if 'LinearModel' in d:
            rv.linearModel = d['LinearModel']
        else:
            rv.linearModel = {}

        if 'ISA' not in d:
            if d['KernelLanguage'] == 'Assembly':
                d['ISA'] = Common.gfxArch(deviceInfo[1])
            else:
                d['ISA'] = [0,0,0]

        rv.originalSolution = OriginalSolution(d)
        # hacky, can just construct Convolution yet again?
        rv.problemType.convolution = rv.originalSolution["ProblemType"].convolution

        return rv

    @classmethod
    def ReadOriginalInfo(cls, d):
        return dict([(key, str(value)) for (key, value) in list(d.items()) if key != 'ProblemType'])

    def __init__(self, **kwargs):
        self.name = None
        self.problemType = None
        self.hardwarePredicate = Hardware.HardwarePredicate('TruePred')
        self.problemPredicate = ProblemPredicate('TruePred')
        self.taskPredicate = TaskPredicate('TruePred')
        self.sizeMapping = None
        self.debugKernel = False
        self.libraryLogicIndex = {}
        self.index = None
        self.ideals = {}

        for key, value in kwargs:
            if key not in Solution.StateKeys and key not in Solution.HiddenKeys:
                raise KeyError("{0} is not a property of Solution.".format(key))

            setattr(self, key, value)