fix(tensilelite): deduplicate kernels by assembly identity

projects/hipblaslt/tensilelite/tensilelite/BenchmarkProblems.py

@@ -46,7 +46,7 @@
 from .Common.TypeValidationErrors import ConfigTypeError
 from .SolutionStructs.Validators.MatrixInstruction import matrixInstructionToMIParameters, \
                                                                  validateMIParameters
-from .SolutionStructs.Naming import getKeyNoInternalArgs, getSolutionNameMin, getKernelNameMin
+from .SolutionStructs.Naming import getSolutionNameMin, getKernelNameMin
 
 from .BenchmarkStructs import BenchmarkProcess
 from .backends import BackendFactory
@@ -73,6 +73,7 @@
     "InternalSupportParams": "internalSupportParams",
     "CustomKernelWildcard": "customKernelWildcard",
 }
+_KERNEL_NAMING_VERSION = 2
 
 # Deterministic synthetic GFlops value emitted per (problem size, solution) cell in the
 # --cpu-only results CSV. Fixed (never random / never timestamped) so the file is
@@ -123,7 +124,9 @@ def _writeSyntheticResultsCSV(resultsFileName, problemSizes, gfxName, numSolutio
 
 def _cacheDataMatches(cacheData, benchmarkStep):
     """Check if cached data matches the current benchmark step parameters."""
-    return all(cacheData[f] == getattr(benchmarkStep, attr) for f, attr in _CACHE_FIELDS.items())
+    return cacheData.get("KernelNamingVersion") == _KERNEL_NAMING_VERSION and all(
+        cacheData[f] == getattr(benchmarkStep, attr) for f, attr in _CACHE_FIELDS.items()
+    )
 
 
 def cmdLineArchsFor(archNames, isaInfoMap) -> List[str]:
@@ -152,6 +155,7 @@ def _computeCacheKey(benchmarkStep, cmdLineArchs):
     hand the wrong machine code to silicon that cannot run it.
     """
     cacheFields = {f: getattr(benchmarkStep, attr) for f, attr in _CACHE_FIELDS.items()}
+    cacheFields["KernelNamingVersion"] = _KERNEL_NAMING_VERSION
     cacheFields["_Architectures"] = ",".join(sorted(cmdLineArchs))
     canonical = json.dumps(cacheFields, sort_keys=True, default=str)
     return hashlib.sha256(canonical.encode()).hexdigest()[:_CACHE_KEY_LEN]
@@ -494,18 +498,14 @@ def writeBenchmarkFiles(
 
     kernels = []
     kernelHelperObjs = []
-    kernelNames = set()
     kernelHelperNames = set()
 
     # get unique kernels and kernel helpers
     with timing_context("python_kernel_bench_setup"):
         for solution in tqdm(solutions, "Finding unique solutions"):
             solutionKernels = solution.getKernels()
             for kernel in solutionKernels:
-                kName = getKeyNoInternalArgs(kernel, debugConfig.splitGSU)
-                if kName not in kernelNames:
-                    kernels.append(kernel)
-                    kernelNames.add(kName)
+                kernels.append(kernel)
 
             solutionHelperKernels = initHelperKernelObjects(solution,
                                                             KernelHelperEnum.All,
@@ -761,6 +761,7 @@ def benchmark_runner(solutions, isCached=False, buildOnly=False):
                 with timing_context("python_write_cache"):
                     cachePath = os.path.join(cacheDir, "cache.yaml")
                     cacheData = {
+                        "KernelNamingVersion": _KERNEL_NAMING_VERSION,
                         "CodeObjectFiles": codeObjectFiles,
                         "LibraryFile": libraryFileRel,
                         "ConstantParams": benchmarkStep.constantParams,

projects/hipblaslt/tensilelite/tensilelite/KernelWriterAssembly.py

@@ -185,10 +185,6 @@ def _getCustomKernelSource(self, kernel, customKernelDirectory=None):
 
   def getSourceFileString(self, kernel) -> Tuple[int, str]:
     assert kernel["KernelLanguage"] == "Assembly"
-    # Skip if .o files will have already been built for this file
-    if kernel.duplicate:
-      self.language = "ASM"
-      return (0, "") # should this be an non zero number
 
     isCustom = isCustomKernelConfig(kernel)
     try:

projects/hipblaslt/tensilelite/tensilelite/SolutionStructs/Naming.py

@@ -30,6 +30,18 @@
 from .Problem import ProblemType
 
 
+def getKernelCompileKey(state, splitGSU: bool) -> str:
+  """Use the emitted symbol as the identity for codegen and metadata lookup.
+
+  The full solution name also encodes selection constraints and legacy tuning
+  inputs (such as ThreadTile on matrix-instruction kernels). Those can differ
+  while the emitted symbol and assembly are identical. Keep the code-changing
+  parameter rules in getKernelNameMin so grouping and artifact naming agree.
+  """
+  s = state._state if hasattr(state, '_state') else state
+  return getKernelNameMin(s, splitGSU)
+
+
 # Parameters that are "internal args" — runtime dispatch parameters that don't
 # affect the generated kernel assembly. Two kernels differing only in these
 # fields compile to identical code objects.
@@ -46,23 +58,10 @@
 )
 
 def getKeyNoInternalArgs(state, splitGSU: bool) -> str:
-  """Return a string key that identifies a kernel ignoring internal args.
-
-  Internal args (WorkGroupMapping, StaggerU, etc.) are runtime dispatch
-  parameters — they don't change the generated assembly. This function
-  produces a canonical key where those parameters are masked to "M" and
-  GroupedGemm is forced to False, so that kernels differing only in
-  internal args map to the same key. GroupedGemm masking is skipped when
-  SupportUserArgs is set, because the batch-offset codegen is gated on
-  GroupedGemm there and the assembly genuinely differs.
-
-  Used to:
-    - Deduplicate kernels before code generation (BenchmarkProblems.py,
-      Run.py:getUniqueKernels) — avoids compiling the same assembly twice.
-    - Identify invalid kernels after compilation and propagate failures
-      to all solutions sharing that kernel (Run.py:removeInvalidSolutionsAndKernels).
-    - Build kernel-to-solution mappings for post-processing
-      (Run.py:passPostKernelInfoToSolution).
+  """Identify solution records, including tuning inputs and library placement.
+
+  This key preserves distinct records during library collection. Code generation
+  uses getKernelCompileKey after collection to group their emitted symbols.
   """
   # Work on the raw dict to avoid Solution.__setitem__ invalidating _name cache
   s = state._state if hasattr(state, '_state') else state
@@ -87,14 +86,13 @@ def getKeyNoInternalArgs(state, splitGSU: bool) -> str:
   for k in _INTERNAL_ARGS:
     s[k] = "M"
 
-  # Compute string key (same as what str(deep_copied_solution) would produce)
-  key = _getName(s, getRequiredParametersFull(), splitGSU, False)
-
-  # Restore
-  pt["GroupedGemm"] = gg_backup
-  s["GlobalSplitU"] = gsu_backup
-  for k in _INTERNAL_ARGS:
-    s[k] = backups[k]
+  try:
+    key = _getName(s, getRequiredParametersFull(), splitGSU, False)
+  finally:
+    pt["GroupedGemm"] = gg_backup
+    s["GlobalSplitU"] = gsu_backup
+    for k in _INTERNAL_ARGS:
+      s[k] = backups[k]
 
   # Include codeObjectFile and DeviceNames in the key to prevent
   # over-deduplication across different code object files / devices.
@@ -170,101 +168,110 @@ def _getName(state, requiredParameters: frozenset, splitGSU: bool, ignoreInterna
   ggBackup = state["ProblemType"]["GroupedGemm"]
   wgmxccBackup = state["WorkGroupMappingXCC"]
 
-  # Include WGMXCC in kernel name as either n1 for auto or 1 for set value
-  # Fixed values produce different assembly code
-  # If the key is missing from name, kernels are dropped as duplicates when they should be kept
-  if "WorkGroupMappingXCC" in state and state["WorkGroupMappingXCC"] != -1:
-    state["WorkGroupMappingXCC"] = 1
-
-  if ignoreInternalArgs:
-    # GroupedGemm is masked so kernels differing only in GroupedGemm dedup to
-    # one key. When SupportUserArgs is set the batch-offset codegen depends on
-    # GroupedGemm, so grouped and non-grouped kernels are not identical and must
-    # keep distinct keys -- preserve the real value in that case.
-    if "SupportUserArgs" not in state["ProblemType"] or not state["ProblemType"]["SupportUserArgs"]:
-      state["ProblemType"]["GroupedGemm"] = False
-    if splitGSU:
-      state["GlobalSplitU"] = "M" if (state["GlobalSplitU"] > 1 or state["GlobalSplitU"] == -1) else state["GlobalSplitU"]
-
-  requiredParametersTemp = set(requiredParameters.union(["GlobalSplitU"]))
-  if state.get("TileProcessingStrategy", "None") != "StreamK":
-    requiredParametersTemp.difference_update({"StreamKAtomic", "StreamKFixupTreeReduction", "DebugStreamK"})
-  if state.get("TileProcessingStrategy", "None") == "None":
-    requiredParametersTemp.difference_update({"WorkAssignment", "PersistentXCCMapping", "WorkQueueStealing"})
-
-
-  if ignoreInternalArgs:
-    if state["GlobalSplitU"] > 0 or state["GlobalSplitU"] == -1:
-      requiredParametersTemp.discard("GlobalSplitU")
-  else:
-    requiredParametersTemp = requiredParametersTemp.union(["WorkGroupMapping",
-                                                          #  "WorkGroupMappingXCC", # WGMXCC affects asm code gen
-                                                           "WorkGroupMappingXCCGroup",
-                                                           "StaggerU",
-                                                           "StaggerUStride",
-                                                           "StaggerUMapping",
-                                                           "GlobalSplitUCoalesced",
-                                                           "GlobalSplitUWorkGroupMappingRoundRobin"])
-  pt = state["ProblemType"]
-  if isinstance(pt, ProblemType):
-    components = [str(pt)]
-  else:
-    components = [str(ProblemType(pt, printIndexAssignmentInfo=False))]
-
-  if "MacroTile0" in state \
-      and "MacroTile1" in state \
-      and "DepthU" in state:
-    components.append(f'{getParameterNameAbbreviation("MacroTile")}{state["MacroTile0"]}x{state["MacroTile1"]}x{state["DepthU"]}')
-
-  if "MatrixInstM" in state:
-    # Use the physical opcode dims (MIBlock) for the name, not the possibly-swapped
-    # effective MatrixInstM/N, so the kernel identity matches the user-specified MI.
-    _miName = state.get("MIBlock", [state["MatrixInstM"], state["MatrixInstN"]])
-    components.append(f'{getParameterNameAbbreviation("MatrixInstruction")}{_miName[0]}x{_miName[1]}x{state["MatrixInstB"]}')
-    requiredParametersTemp.add("MIWaveTile")
-  else:
-    requiredParametersTemp.add("ThreadTile")
-
-  if state["UseCustomMainLoopSchedule"]:
-    components.append('CMS')
-
-  components.append('SN')
-  if state.get("TileProcessingStrategy") == "DataParallel":
-    components.append(f'PLAV{state.get("InternalSupportParams", {}).get("PersistentLoopArgsVersion", 0)}')
-
-  # Skip SFA tag if using default wgm algo
-  if "SpaceFillingAlgo" in requiredParametersTemp and len(state["SpaceFillingAlgo"]) == 0:
-    requiredParametersTemp.discard("SpaceFillingAlgo")
+  try:
+    # Fixed XCC mappings share code except when dynamic/hybrid work queues
+    # need to preserve the raw launch rank for a multi-XCC mapping.
+    if state["WorkGroupMappingXCC"] != -1:
+      state["WorkGroupMappingXCC"] = (
+          2 if state.get("WorkAssignment") in ("DynamicWorkQueue", "Hybrid")
+          and state.get("PersistentXCCMapping", 0) != 0
+          and state["WorkGroupMappingXCC"] > 1 else 1
+      )
+
+    if ignoreInternalArgs:
+      # GroupedGemm is masked so kernels differing only in GroupedGemm dedup to
+      # one key. When SupportUserArgs is set the batch-offset codegen depends on
+      # GroupedGemm, so grouped and non-grouped kernels are not identical and must
+      # keep distinct keys -- preserve the real value in that case.
+      if "SupportUserArgs" not in state["ProblemType"] or not state["ProblemType"]["SupportUserArgs"]:
+        state["ProblemType"]["GroupedGemm"] = False
+      if splitGSU:
+        state["GlobalSplitU"] = "M" if (state["GlobalSplitU"] > 1 or state["GlobalSplitU"] == -1) else state["GlobalSplitU"]
+
+    requiredParametersTemp = set(requiredParameters.union(["GlobalSplitU"]))
+    if state.get("TileProcessingStrategy", "None") != "StreamK":
+      requiredParametersTemp.difference_update({"StreamKAtomic", "StreamKFixupTreeReduction", "DebugStreamK"})
+    if state.get("TileProcessingStrategy", "None") == "None":
+      requiredParametersTemp.difference_update({"WorkAssignment", "PersistentXCCMapping", "WorkQueueStealing"})
+
+
+    if ignoreInternalArgs:
+      if state["GlobalSplitU"] > 0 or state["GlobalSplitU"] == -1:
+        requiredParametersTemp.discard("GlobalSplitU")
+    else:
+      requiredParametersTemp = requiredParametersTemp.union(["WorkGroupMapping",
+                                                            #  "WorkGroupMappingXCC", # WGMXCC affects asm code gen
+                                                             "WorkGroupMappingXCCGroup",
+                                                             "StaggerU",
+                                                             "StaggerUStride",
+                                                             "StaggerUMapping",
+                                                             "GlobalSplitUCoalesced",
+                                                             "GlobalSplitUWorkGroupMappingRoundRobin"])
+    pt = state["ProblemType"]
+    if isinstance(pt, ProblemType):
+      components = [str(pt)]
+    else:
+      components = [str(ProblemType(pt, printIndexAssignmentInfo=False))]
+
+    if "MacroTile0" in state \
+        and "MacroTile1" in state \
+        and "DepthU" in state:
+      components.append(f'{getParameterNameAbbreviation("MacroTile")}{state["MacroTile0"]}x{state["MacroTile1"]}x{state["DepthU"]}')
+
+    if "MatrixInstM" in state:
+      # Use the physical opcode dims (MIBlock) for the name, not the possibly-swapped
+      # effective MatrixInstM/N, so the kernel identity matches the user-specified MI.
+      _miName = state.get("MIBlock", [state["MatrixInstM"], state["MatrixInstN"]])
+      components.append(f'{getParameterNameAbbreviation("MatrixInstruction")}{_miName[0]}x{_miName[1]}x{state["MatrixInstB"]}')
+      requiredParametersTemp.add("MIWaveTile")
+    else:
+      requiredParametersTemp.add("ThreadTile")
 
-  # TDMFuse=0 is the arrangement every shipped kernel already has, so naming it
-  # would rename all of them.
-  if state.get("TDMFuse", 0):
-    requiredParametersTemp.add("TDMFuse")
-  else:
-    requiredParametersTemp.discard("TDMFuse")
+    if state["UseCustomMainLoopSchedule"]:
+      components.append('CMS')
 
-  # LDSSegmentBRotate=0 is the default layout; naming it would rename every kernel.
-  if state.get("LDSSegmentBRotate", 0):
-    requiredParametersTemp.add("LDSSegmentBRotate")
-  else:
-    requiredParametersTemp.discard("LDSSegmentBRotate")
+    components.append('SN')
+    if state.get("TileProcessingStrategy") == "DataParallel":
+      components.append(f'PLAV{state.get("InternalSupportParams", {}).get("PersistentLoopArgsVersion", 0)}')
 
-  for key in sorted(requiredParametersTemp):
-    if key not in state or key == "CustomKernel":
-      continue
-    components.append(f'{getParameterNameAbbreviation(key)}{getParameterValueAbbreviation(key, state[key])}')
+    # Skip SFA tag if using default wgm algo
+    if "SpaceFillingAlgo" in requiredParametersTemp and len(state["SpaceFillingAlgo"]) == 0:
+      requiredParametersTemp.discard("SpaceFillingAlgo")
 
-  # Tuning-only StinkyTofu overrides change the asm, so tag the name; absent when unused.
-  stinkyTune = state.get("_StinkyTofuParameters")
-  if stinkyTune:
-    components.append("STP")
-    components.extend(f"{k}{_nameSafe(v)}" for k, v in sorted(stinkyTune.items()))
+    # Non-default debug modes replace or remove generated control flow.
+    for key in ("DebugStreamK", "StreamKAtomic", "MbskPrefetchMethod", "DebugPersistentKernelLoopForever"):
+      if state.get(key, 0) != 0:
+        requiredParametersTemp.add(key)
 
-  state["GlobalSplitU"] = gsuBackup
-  state["ProblemType"]["GroupedGemm"] = ggBackup
-  state["WorkGroupMappingXCC"] = wgmxccBackup
+    # TDMFuse=0 is the arrangement every shipped kernel already has, so naming it
+    # would rename all of them.
+    if state.get("TDMFuse", 0):
+      requiredParametersTemp.add("TDMFuse")
+    else:
+      requiredParametersTemp.discard("TDMFuse")
 
-  return '_'.join(components)
+    # LDSSegmentBRotate=0 is the default layout; naming it would rename every kernel.
+    if state.get("LDSSegmentBRotate", 0):
+      requiredParametersTemp.add("LDSSegmentBRotate")
+    else:
+      requiredParametersTemp.discard("LDSSegmentBRotate")
+
+    for key in sorted(requiredParametersTemp):
+      if key not in state or key == "CustomKernel":
+        continue
+      components.append(f'{getParameterNameAbbreviation(key)}{getParameterValueAbbreviation(key, state[key])}')
+
+    # Tuning-only StinkyTofu overrides change the asm, so tag the name; absent when unused.
+    stinkyTune = state.get("_StinkyTofuParameters")
+    if stinkyTune:
+      components.append("STP")
+      components.extend(f"{k}{_nameSafe(v)}" for k, v in sorted(stinkyTune.items()))
+
+    return '_'.join(components)
+  finally:
+    state["GlobalSplitU"] = gsuBackup
+    state["ProblemType"]["GroupedGemm"] = ggBackup
+    state["WorkGroupMappingXCC"] = wgmxccBackup
 
 
 def shortenFileBase(splitGSU, kernel):