Files
TheRock/rocm-libraries/projects/hipblaslt/tensilelite/Tensile/SolutionLibrary.py
T
derek.holloway 7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
Flatten subprojects into a single base monorepo
2026-09-02 16:34:47 -07:00

717 lines
27 KiB
Python

################################################################################
#
# 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 typing import Dict
from . import Properties
from . import Hardware
from . import Contractions
from Tensile.Common import state, IsaInfo
from Tensile.Common.Architectures import gfxToIsa, supportsChipIdPredicate
from Tensile.SolutionStructs.Naming import getSolutionNameMin, getKernelNameMin
class SingleSolutionLibrary:
Tag = "Single"
def __init__(self, solution):
self.solution = solution
@property
def tag(self):
return self.__class__.Tag
def state(self):
return {"type": self.tag, "index": self.solution.index}
def remapSolutionIndices(self, indexMap):
pass
class IndexSolutionLibrary(SingleSolutionLibrary):
Tag = "Index"
def state(self):
return self.solution.index
class PlaceholderLibrary:
Tag = 'Placeholder'
def __init__(self, name):
self.filenamePrefix = name
@property
def tag(self):
return self.__class__.Tag
def state(self):
return {'type': self.tag, 'value': self.filenamePrefix}
def remapSolutionIndices(self, indexMap):
pass
def merge(self, other):
otherName = getattr(other, 'filenamePrefix', None)
if otherName != self.filenamePrefix:
raise RuntimeError(
"[PlaceholderLibrary.merge] COLLISION: two source YAMLs "
"converge at the same dispatch slot but reference different "
"per-file libraries; one .dat will be orphaned.\n"
f" self: {self.filenamePrefix}\n"
f" other: {otherName!r}\n"
" Likely cause: chip-id placeholder-suffix gate is out of "
"sync with HardwarePredicate.FromHardware, or sibling YAMLs "
"declare divergent DeviceNames on a chip-id-unaware arch."
)
class MatchingLibrary:
Tag = "Matching"
StateKeys = [("type", "tag"), "properties", "table", "distance"]
@classmethod
def FromOriginalState(cls, d, solutions):
indices = d["indexOrder"]
distance = d["distance"]
origTable = d["table"]
propertyKeys = {
2: Properties.Property("FreeSizeA", index=0),
3: Properties.Property("FreeSizeB", index=0),
1: Properties.Property("BoundSize", index=0)
}
if distance == "Equality" or distance == "GridBased":
propertyKeys[0] = Properties.Property("BatchSize", index=0)
if distance == "Range":
propertyKeys = {
0: Properties.Property("FreeSizeA", index=0),
1: Properties.Property("FreeSizeA", index=0),
2: Properties.Property("FreeSizeB", index=0),
3: Properties.Property("FreeSizeB", index=0),
4: Properties.Property("BatchSize", index=0),
5: Properties.Property("BatchSize", index=0),
6: Properties.Property("BoundSize", index=0),
7: Properties.Property("BoundSize", index=0)
}
properties = list([propertyKeys[i] for i in indices if i in propertyKeys])
keyOrder = [i for i, j in enumerate(indices) if j in propertyKeys]
table = []
for row in origTable:
try:
index = row[1][0]
value = IndexSolutionLibrary(solutions[index])
key = list([row[0][i] for i in keyOrder])
if distance == "GridBased":
entry = {"key": key, "index": value}
else:
entry = {"key": key, "index": value, "speed": row[1][1]}
table.append(entry)
except KeyError:
pass
table.sort(key=lambda r: r["key"])
return cls(properties, table, distance)
@property
def tag(self):
return self.__class__.Tag
def merge(self, other):
assert self.__class__ == other.__class__ \
and self.properties == other.properties \
and self.distance == other.distance
self.table += other.table
self.table.sort(key=lambda r: r["key"])
def remapSolutionIndices(self, indexMap):
pass
def __init__(self, properties, table, distance):
self.properties = properties
self.table = table
self.distance = distance
class FreeSizeLibrary:
Tag = "FreeSize"
StateKeys = [("type", "tag"), "table"]
@classmethod
def FromOriginalState(cls, d, solutions):
origTable = d["table"]
table = []
try:
indexStart = origTable[0]
indexOffset = origTable[1]
for index in range(indexStart, indexStart + indexOffset):
value = IndexSolutionLibrary(solutions[index])
entry = {"index": value}
table.append(entry)
except KeyError:
pass
return cls(table)
@property
def tag(self):
return self.__class__.Tag
def merge(self, other):
assert self.__class__ == other.__class__
self.table += other.table
def remapSolutionIndices(self, indexMap):
pass
def __init__(self, table):
self.table = table
class PredictionLibrary:
Tag = "Prediction"
StateKeys = [("type", "tag"), "table"]
@classmethod
def FromOriginalState(cls, d, solutions):
origTable = d["table"]
table = []
try:
indexStart = origTable[0]
indexOffset = origTable[1]
for index in range(indexStart, indexStart + indexOffset):
value = IndexSolutionLibrary(solutions[index])
table.append(value)
except KeyError:
pass
return cls(table)
@property
def tag(self):
return self.__class__.Tag
def merge(self, other):
assert self.__class__ == other.__class__
self.table += other.table
def remapSolutionIndices(self, indexMap):
pass
def __init__(self, table):
self.table = table
class MLPClassificationLibrary:
Tag = "MLPClassification"
StateKeys = [("type", "tag"), "table", "mlp", "problemFeatures"]
@classmethod
def FromOriginalState(cls, d, solutions):
origTable = d["table"]
table = []
try:
indexStart = origTable[0]
indexOffset = origTable[1]
for index in range(indexStart, indexStart + indexOffset):
value = IndexSolutionLibrary(solutions[index])
table.append(value)
except KeyError:
pass
mlp = d["mlp"]
problem_features = d["problemFeatures"]
return cls(table, mlp, problem_features)
@property
def tag(self):
return self.__class__.Tag
def merge(self, other):
raise RuntimeError(
"MLPClassificationLibrary does not support merging."
)
def remapSolutionIndices(self, indexMap):
pass
def __init__(self, table, mlp, problem_features):
self.table = table
self.mlp = mlp
self.problemFeatures = problem_features
class ProblemMapLibrary:
Tag = "ProblemMap"
StateKeys = [("type", "tag"), ("property", "mappingProperty"), ("map", "mapping")]
def __init__(self, mappingProperty=None, mapping=None):
self.mappingProperty = mappingProperty
self.mapping = mapping
@property
def tag(self):
return self.__class__.Tag
def merge(self, other):
assert self.__class__ == other.__class__ and self.tag == other.tag and self.mappingProperty == other.mappingProperty
for key, value in list(other.mapping.items()):
if key in self.mapping:
self.mapping[key].merge(value)
else:
self.mapping[key] = value
def remapSolutionIndices(self, indexMap):
for key, value in list(self.mapping.items()):
value.remapSolutionIndices(indexMap)
class PredicateLibrary:
StateKeys = [("type", "tag"), "rows"]
def __init__(self, tag=None, rows=None):
self.tag = tag
if rows is None: rows = []
self.rows = rows
def merge(self, other):
assert self.__class__ == other.__class__ and self.tag == other.tag
rowPreds = [r["predicate"] for r in self.rows]
for row in other.rows:
if row["predicate"] in rowPreds:
myRownum = rowPreds.index(row["predicate"])
self.rows[myRownum]["library"].merge(row["library"])
else:
self.rows.append(row)
if self.rows[0]["library"].tag == "Placeholder":
# Sort to ensure pure gemm can be search first
self.rows.sort(key=lambda x: len(x["library"].filenamePrefix))
else:
# Sort to ensure consistent fallback logic.
self.rows.sort(key=lambda x: x["predicate"])
def remapSolutionIndices(self, indexMap):
for row in self.rows:
row["library"].remapSolutionIndices(indexMap)
class MasterSolutionLibrary:
StateKeys = ["solutions", "library"]
@staticmethod
def hardware(d, library, placeholderName, lazyLibrary, logicFile=None):
"""Build the Hardware-level PredicateLibrary row and update the
placeholder filename.
Lifted out of ``FromOriginalState`` so the placeholder-suffix gating
invariant (the ``_ID<chipid>`` suffix must only be appended when the
runtime ``HardwarePredicate`` also discriminates on chip-id, see
``Hardware.HardwarePredicate.FromHardware`` and
``supportsChipIdPredicate``) can be exercised behaviorally rather than
by source-string inspection.
"""
devicePart = d["ArchitectureName"]
cuCount = d["CUCount"]
pciChipId = d.get("DeviceNames", None)
newLib = PredicateLibrary(tag="Hardware")
if devicePart == "fallback":
pred = Hardware.HardwarePredicate("TruePred")
else:
pred = Hardware.HardwarePredicate.FromHardware(
gfxToIsa(devicePart), cuCount, pciChipId, logicFile=logicFile
)
newLib.rows.append({"predicate": pred, "library": library})
if lazyLibrary:
if cuCount: placeholderName += "_CU" + str(cuCount)
# Only append the chip-id suffix on architectures whose runtime
# HardwarePredicate also includes the chip-id discriminator
# (see Hardware.HardwarePredicate.FromHardware). Otherwise the
# filename diverges between sibling YAMLs while their predicates
# remain equal, producing PredicateLibrary.merge collisions whose
# PlaceholderLibrary children silently drop one leaf.
if pciChipId and supportsChipIdPredicate(devicePart):
# Convert device names list to a sanitized string for filename
# e.g., ['Device 75a0', 'Device 75b0'] -> 'ID75a0-75b0'
if isinstance(pciChipId, list):
chipIdStr = '-'.join([str(d).replace('Device ', '').strip() for d in pciChipId])
else:
chipIdStr = str(pciChipId).replace('Device ', '').strip()
placeholderName += "_ID" + chipIdStr
placeholderName += "_" + str(devicePart)
return newLib, placeholderName
@classmethod
def FixSolutionIndices(cls, solutions):
# fix missing and duplicate solution indices.
try:
maxSolutionIdx = max([s.index for s in solutions if s.index is not None])
except ValueError:
maxSolutionIdx = -1
solutionsSoFar = set()
for solution in solutions:
if solution.index is None or solution.index in solutionsSoFar:
maxSolutionIdx += 1
solution.index = maxSolutionIdx
else:
solutionsSoFar.add(solution.index)
@classmethod
def FromOriginalState(cls,
origData,
origSolutions,
splitGSU: bool,
printSolutionRejectionReason: bool,
printIndexAssignmentInfo: bool,
assembler,
isaInfoMap: Dict[str, IsaInfo],
lazyLibraryLoading: bool,
logicFile=None,
solutionClass=Contractions.Solution,
libraryOrder=None,
placeholderName='TensileLibrary'):
# functions for creating each "level" of the library
def hardware(d, problemType, solutions, library, placeholderName):
return cls.hardware(d, library, placeholderName, lazyLibrary, logicFile=logicFile)
def operationIdentifier(d, problemType, solutions, library, placeholderName):
operationID = problemType.operationIdentifier
prop = Properties.Property("OperationIdentifier")
mapping = {operationID: library}
newLib = ProblemMapLibrary(prop, mapping)
if lazyLibrary:
placeholderName += "_" + operationID
return newLib, placeholderName
def performanceMetric(d, problemType, solutions, library, placeholderName):
if d.get("PerfMetric", "DeviceEfficiency") != "DeviceEfficiency":
predicate = Properties.Predicate(tag=d["PerfMetric"])
else:
predicate = Properties.Predicate(tag="TruePred")
newLib = PredicateLibrary(tag="Problem")
newLib.rows.append({"predicate": predicate, "library": library})
if lazyLibrary and predicate.tag != "TruePred":
placeholderName += "_" + predicate.tag
return newLib, placeholderName
def predicates(d, problemType, solutions, library, placeholderName):
predicates = problemType.predicates(includeBatch=True, includeType=True)
predicate = Contractions.ProblemPredicate.And(predicates)
newLib = PredicateLibrary(tag="Problem")
newLib.rows.append({"predicate": predicate, "library": library})
if lazyLibrary:
placeholderName += problemType.placeholderStr(includeBatch=True, includeType=True)
return newLib, placeholderName
def placeholder(d, problemType, solutions, library, placeholderName):
newLib = PlaceholderLibrary(placeholderName)
return newLib, placeholderName
def selection(d, problemType, solutions, library, placeholderName):
if d["LibraryType"] == "Matching":
if d["Library"]["distance"] == "Equality":
predicate = Properties.Predicate(tag="EqualityMatching")
elif d["Library"]["distance"] == "Range":
predicate = Properties.Predicate(tag="RangeMatching")
else:
predicate = Properties.Predicate(tag="GridBasedMatching") # to make different from TruePred
matchingLib = MatchingLibrary.FromOriginalState(d["Library"], solutions)
library = PredicateLibrary(tag="Problem")
library.rows.append({"predicate": predicate, "library": matchingLib})
elif d["LibraryType"] == "FreeSize":
predicate = Properties.Predicate(tag="FreeSizeMatching")
freesizeLib = FreeSizeLibrary.FromOriginalState(d["Library"], solutions)
library = PredicateLibrary(tag="Problem")
library.rows.append({"predicate": predicate, "library": freesizeLib})
elif d["LibraryType"] == "Prediction":
predicate = Properties.Predicate(tag="PredictionMatching") # So that FreeSize / Prediction can co-exist
predictionLib = PredictionLibrary.FromOriginalState(d["Library"], solutions)
library = PredicateLibrary(tag="Problem")
library.rows.append({"predicate": predicate, "library": predictionLib})
elif d["LibraryType"] == "MLPClassification":
predicate = Properties.Predicate(tag="TruePred")
regressionLib = MLPClassificationLibrary.FromOriginalState(d["Library"], solutions)
library = PredicateLibrary(tag="Problem")
library.rows.append({"predicate": predicate, "library": regressionLib})
else:
assert 0 and "Unrecognized LibraryType."
if lazyLibraryLoading:
computeInputTypeStr = str(problemType.computeInputTypeA)
if problemType.computeInputTypeA != problemType.computeInputTypeB:
computeInputTypeStr = str(problemType.computeInputTypeA) + str(problemType.computeInputTypeB)
placeholderName += '_' + str(problemType.aType) + str(problemType.bType)
placeholderName += '_' + str(problemType.cType) + computeInputTypeStr
if problemType.activationType != 'none':
if str(problemType.activationType).upper() == 'ALL':
placeholderName += "_A"
elif str(problemType.activationType).upper() == 'HIPBLASLT_ALL':
placeholderName += "_HA"
else:
placeholderName += "_%s"%str(problemType.activationType).upper()
if problemType.mxBlockA:
placeholderName += ('_MXA' + str(problemType.mxTypeA) + 'B' + str(problemType.mxBlockA))
if problemType.mxBlockB:
placeholderName += ('_MXB' + str(problemType.mxTypeB) + 'B' + str(problemType.mxBlockB))
if problemType.swizzleTensorA:
placeholderName += '_STA'
if problemType.swizzleTensorB:
placeholderName += '_STB'
if problemType.useBias:
placeholderName += '_Bias'
if problemType.useE:
placeholderName += '_Grad' if problemType.useGradient else '_Aux'
if problemType.groupedGemm:
placeholderName += "_GG"
else:
placeholderName += "" if problemType.stridedBatched else "_GB" # legacy
if problemType.useScaleAB == "Scalar":
placeholderName += '_SAB'
elif problemType.useScaleAB == "Vector":
placeholderName += '_SABV'
if problemType.useScaleCD:
placeholderName += '_SCD'
if problemType.useScaleAlphaVec:
placeholderName += '_SAV'
if problemType.sparse:
placeholderName += '_SPB' if problemType.sparse == 2 else '_SPA'
placeholderName += "ML" + str(problemType.metadataLayout)
if not problemType.f32XdlMathOp.isSingle() and problemType.computeInputTypeA.isSingle() and problemType.computeInputTypeB.isSingle():
placeholderName += '_M' + str(problemType.f32XdlMathOp)
if problemType.supportDeviceUserArguments:
placeholderName += '_UA'
if getattr(problemType, 'useGateResidual', False):
placeholderName += '_GateRes'
return library, placeholderName
# end library creation functions
if libraryOrder is None:
if lazyLibraryLoading:
libraryOrder = [
hardware, operationIdentifier, performanceMetric, predicates,
placeholder, selection
]
else:
libraryOrder = [
hardware, operationIdentifier, performanceMetric, predicates,
selection
]
#assert libraryOrder[-1] == selection
lazyLibrary = None
if placeholder in libraryOrder:
placeholderIndex = libraryOrder.index(placeholder) + 1
lazyLibrary, placeholderName = \
MasterSolutionLibrary.FromOriginalState(origData,
origSolutions,
splitGSU,
printSolutionRejectionReason,
printIndexAssignmentInfo,
assembler,
isaInfoMap,
lazyLibraryLoading,
logicFile,
solutionClass,
libraryOrder[placeholderIndex:],
placeholderName)
libraryOrder = libraryOrder[0:placeholderIndex]
origSolutions = []
problemType = Contractions.ProblemType.FromOriginalState(origData["ProblemType"])
allSolutions = [solutionClass.FromSolutionStruct(
s,
splitGSU,
printSolutionRejectionReason,
printIndexAssignmentInfo,
assembler,
isaInfoMap
) for s in origSolutions]
cls.FixSolutionIndices(allSolutions)
# library is constructed in reverse order i.e. bottom-up
library = None
placeholderLibrary = None
for libName in reversed(libraryOrder):
library, placeholderName = libName(origData, problemType, allSolutions, library,
placeholderName)
if libName == placeholder:
placeholderLibrary = library
solutions = {s.index: s for s in allSolutions}
rv = cls(solutions, library)
if lazyLibrary and placeholderLibrary:
rv.lazyLibraries[placeholderName] = lazyLibrary
placeholderLibrary.filenamePrefix = placeholderName
return rv, placeholderName
@classmethod
def BenchmarkingLibrary(
cls,
solutions,
assembler,
splitGSU: bool,
printSolutionRejectionReason: bool,
printIndexAssignmentInfo: bool,
isaInfoMap
):
solutionObjs = list([Contractions.Solution.FromOriginalState(
s._state,
splitGSU,
printSolutionRejectionReason,
printIndexAssignmentInfo,
assembler,
isaInfoMap)
for s in solutions])
cls.FixSolutionIndices(solutionObjs)
predRows = list([{
"predicate": s.problemPredicate,
"library": SingleSolutionLibrary(s)
} for s in solutionObjs])
library = PredicateLibrary(tag="Problem", rows=predRows)
solutionMap = {s.index: s for s in solutionObjs}
return cls(solutionMap, library)
def __init__(self, solutions, library, version=None):
self.lazyLibraries = {}
self.solutions = solutions
self.library = library
self.version = version
def state(self):
rv = {
"solutions": state(self.solutions.values()),
"library": state(self.library)
}
if self.version is not None:
rv["version"] = self.version
return rv
def applyNaming(self, splitGSU: bool):
for s in list(self.solutions.values()):
s.name = getSolutionNameMin(s.originalSolution.getKernels()[0], splitGSU)
s.kernelName = getKernelNameMin(s.originalSolution.getKernels()[0], splitGSU)
def remapSolutionIndicesStartingFrom(self, curIndex):
reIndexMap = {}
solutionCopy = self.solutions
self.solutions = dict()
for k, s in solutionCopy.items():
reIndexMap[s.index] = curIndex
s.index = curIndex
self.solutions[curIndex] = s
curIndex += 1
self.library.remapSolutionIndices(reIndexMap)
def merge(self, other, startIndex=0):
assert self.__class__ == other.__class__
curIndex = max(startIndex, max(self.solutions.keys()) + 1 if self.solutions else 0)
if self.lazyLibraries:
curIndex = max(
curIndex,
max(max(lib.solutions.keys()) for _, lib in self.lazyLibraries.items()) + 1)
#Merge separate library files
for name, lib in other.lazyLibraries.items():
if name in self.lazyLibraries.keys():
curIndex = self.lazyLibraries[name].merge(lib, curIndex)
else:
reIndexMap = {}
newSolutions = {}
for k, s in lib.solutions.items():
reIndexMap[s.index] = curIndex
s.index = curIndex
newSolutions[curIndex] = s
curIndex += 1
lib.solutions = newSolutions
lib.library.remapSolutionIndices(reIndexMap)
self.lazyLibraries[name] = lib
reIndexMap = {}
for k, s in other.solutions.items():
reIndexMap[s.index] = curIndex
s.index = curIndex
self.solutions[curIndex] = s
curIndex += 1
other.library.remapSolutionIndices(reIndexMap)
self.library.merge(other.library)
return curIndex #Next unused index
@property
def cpp_base_class(self):
return "SolutionLibrary<ContractionProblemGemm, ContractionSolution>"
@property
def cpp_class(self):
return "MasterSolutionLibrary<ContractionProblemGemm, ContractionSolution>"