fix(tensilelite): find msgpack master libraries in the LibraryClient step

projects/hipblaslt/tensilelite/Tensile/ClientWriter.py

@@ -45,7 +45,7 @@
 from Tensile.Common.Architectures import ARCH_COMPILER_TARGET, baseArchName, gfxToIsa, isaToGfx
 from Tensile.Common.GlobalParameters import globalParameters
 from Tensile.Common.TimingInstrumentation import timing_context
-from .TensileCreateLibrary import copyStaticFiles, libraryDir
+from .TensileCreateLibrary import computeOutputArchNames, copyStaticFiles, libraryDir
 from .ParallelExecution import detectAvailableGpus, runClientParallel
 from .Contractions import FreeIndex, BatchIndex
 from .Contractions import ProblemType as ContractionsProblemType
@@ -124,6 +124,32 @@ def buildTargetGfx(isaInfoMap, archNames=None) -> str:
   return requested.get(isa, isaToGfx(isa))
 
 
+def clientLibraryFiles(clientLibraryPath, archs, outArchNames=None):
+  """Return the code objects and master library files built for `archs`.
+
+  Kernels fan out into one per-base subdir per arch, so the lists are unioned
+  across them. `outArchNames` maps an ISA-derived arch to the subdir it was
+  written to (see computeOutputArchNames): a stepping such as gfx1250v0 lands
+  in `library/gfx1250v0/` while its master keeps the `gfx1250` suffix. The
+  master is matched by its exact name because lazy-loading shards share the
+  `TensileLibrary_` prefix. Msgpack is written to disk as `<name>.dat.zlib`,
+  but the client must be given the logical `.dat` name; it probes for the
+  `.zlib` variant itself.
+  """
+  libraryExt = ".yaml" if globalParameters["LibraryFormat"] == "yaml" else ".dat"
+  masterPrefix = "TensileLibrary_lazy_" if globalParameters["LazyLibraryLoading"] else "TensileLibrary_"
+  outArchNames = outArchNames or {}
+  coList = []
+  libraryList = []
+  for arch in archs:
+    archDir = libraryDir(clientLibraryPath, outArchNames.get(arch, arch))
+    coList.extend(glob(os.path.join(archDir, "*.co")))
+    master = os.path.join(archDir, masterPrefix + arch + libraryExt)
+    if os.path.exists(master) or os.path.exists(master + ".zlib"):
+      libraryList.append(master)
+  return coList, libraryList
+
+
 def main(config, assembler: Assembler, cCompiler: str, isaInfoMap, outputPath: Path, deviceId: int, gfxName: str, archNames=None):
 
   libraryLogicPath = ensurePath(outputPath / LIBRARY_LOGIC_DIR)
@@ -155,13 +181,7 @@ def main(config, assembler: Assembler, cCompiler: str, isaInfoMap, outputPath: P
   createLibraryScript = getBuildClientLibraryScript(clientLibraryPath, libraryLogicPath, str(assembler.path), targetGfx)
   subprocess.run(shlex.split(createLibraryScript), env=env, cwd=clientLibraryPath)
   archs = [isaToGfx(isa) for isa in isaInfoMap.keys()]
-  # Kernels fan out into one per-base subdir per arch; union the globs across them.
-  coList = []
-  yamlList = []
-  for arch in archs:
-    archDir = libraryDir(clientLibraryPath, arch)
-    coList.extend(glob(os.path.join(archDir, "*.co")))
-    yamlList.extend(glob(os.path.join(archDir, "*.yaml")))
+  coList, libraryList = clientLibraryFiles(clientLibraryPath, archs, computeOutputArchNames([targetGfx]))
 
   clientParametersPaths = []
   splitGSU = False
@@ -227,7 +247,7 @@ def main(config, assembler: Assembler, cCompiler: str, isaInfoMap, outputPath: P
                                   deviceId=deviceId,
                                   gfxName=gfxName,
                                   tileAwareSelection=False,
-                                  libraryFile=yamlList[0]))
+                                  libraryFile=libraryList[0]))
 
   forBenchmark = False
   problemSizes = None

projects/hipblaslt/tensilelite/Tensile/TensileCreateLibrary/__init__.py

@@ -1 +1 @@
-from .Run import copyStaticFiles, libraryDir, libraryRoot, run, tensileLibraryFile, writeSolutionsAndKernels
+from .Run import computeOutputArchNames, copyStaticFiles, libraryDir, libraryRoot, run, tensileLibraryFile, writeSolutionsAndKernels

projects/hipblaslt/tensilelite/Tensile/Tests/unit/test_ClientWriter_library_files.py

@@ -0,0 +1,107 @@
+# Copyright Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+
+"""
+Unit tests for Tensile.ClientWriter.clientLibraryFiles.
+
+The LibraryClient step used to collect only `*.yaml` libraries, so with
+LibraryFormat=msgpack it found nothing and ClientWriter.main crashed on
+`libraryList[0]` with IndexError.
+
+Library files are written with LibraryIO.write so the on-disk names (e.g.
+msgpack's `.dat.zlib`) match what TensileCreateLibrary produces.
+"""
+
+import os
+
+import pytest
+
+import Tensile.ClientWriter as CW
+from Tensile import LibraryIO
+from Tensile.Common.GlobalParameters import globalParameters
+
+pytestmark = pytest.mark.unit
+
+# A lazy-loading shard name; shards share the master's `TensileLibrary_` prefix.
+SHARD = "TensileLibrary_Type_SS_Contraction_l_Ailk_Bljk_Cijk_Dijk_{arch}"
+
+
+def _write_library(directory, stem, library_format):
+    directory.mkdir(parents=True, exist_ok=True)
+    LibraryIO.write(str(directory / stem), {}, library_format)
+
+
+def _build_arch(root, arch, library_format, lazy, out_dir=None):
+    arch_dir = CW.libraryDir(root, out_dir or arch)
+    master = f"TensileLibrary_lazy_{arch}" if lazy else f"TensileLibrary_{arch}"
+    _write_library(arch_dir, master, library_format)
+    if lazy:
+        _write_library(arch_dir, SHARD.format(arch=arch), library_format)
+        # Always written as msgpack, whatever LibraryFormat is; never the master.
+        _write_library(arch_dir, f"TensileLiteLibrary_lazy_{arch}_Mapping", "msgpack")
+    code_object = arch_dir / f"Kernels.so-000-{arch}.co"
+    code_object.write_bytes(b"")
+    ext = ".yaml" if library_format == "yaml" else ".dat"
+    return str(arch_dir / (master + ext)), str(code_object)
+
+
+@pytest.mark.parametrize("lazy", [True, False], ids=["lazy", "not-lazy"])
+@pytest.mark.parametrize("library_format", ["msgpack", "yaml"])
+def test_selects_master_library(tmp_path, monkeypatch, library_format, lazy):
+    monkeypatch.setitem(globalParameters, "LibraryFormat", library_format)
+    monkeypatch.setitem(globalParameters, "LazyLibraryLoading", lazy)
+    master, code_object = _build_arch(tmp_path, "gfx942", library_format, lazy)
+
+    co_list, library_list = CW.clientLibraryFiles(tmp_path, ["gfx942"])
+
+    assert library_list == [master]
+    assert co_list == [code_object]
+
+
+def test_msgpack_master_is_logical_dat_name(tmp_path, monkeypatch):
+    """The client takes the `.dat` name and probes for `.dat.zlib` itself."""
+    monkeypatch.setitem(globalParameters, "LibraryFormat", "msgpack")
+    monkeypatch.setitem(globalParameters, "LazyLibraryLoading", True)
+    master, _ = _build_arch(tmp_path, "gfx942", "msgpack", lazy=True)
+
+    assert not os.path.exists(master)
+    assert os.path.exists(master + ".zlib")
+    assert CW.clientLibraryFiles(tmp_path, ["gfx942"])[1] == [master]
+
+
+def test_unions_files_across_archs(tmp_path, monkeypatch):
+    monkeypatch.setitem(globalParameters, "LibraryFormat", "msgpack")
+    monkeypatch.setitem(globalParameters, "LazyLibraryLoading", True)
+    expected = [_build_arch(tmp_path, arch, "msgpack", lazy=True)[0] for arch in ("gfx942", "gfx950")]
+
+    _, library_list = CW.clientLibraryFiles(tmp_path, ["gfx942", "gfx950"])
+
+    assert library_list == expected
+
+
+def test_skips_arch_without_master(tmp_path, monkeypatch):
+    """Only the first ISA's client library is rebuilt, so other archs can have
+    shards but no master; they must not contribute one."""
+    monkeypatch.setitem(globalParameters, "LibraryFormat", "msgpack")
+    monkeypatch.setitem(globalParameters, "LazyLibraryLoading", True)
+    master, _ = _build_arch(tmp_path, "gfx942", "msgpack", lazy=True)
+    _write_library(CW.libraryDir(tmp_path, "gfx950"), SHARD.format(arch="gfx950"), "msgpack")
+
+    _, library_list = CW.clientLibraryFiles(tmp_path, ["gfx942", "gfx950"])
+
+    assert library_list == [master]
+
+
+def test_finds_stepping_output_subdir(tmp_path, monkeypatch):
+    """A gfx1250v0 client build is written to library/gfx1250v0/, but its master
+    keeps the ISA-derived gfx1250 suffix."""
+    monkeypatch.setitem(globalParameters, "LibraryFormat", "msgpack")
+    monkeypatch.setitem(globalParameters, "LazyLibraryLoading", True)
+    master, code_object = _build_arch(tmp_path, "gfx1250", "msgpack", lazy=True, out_dir="gfx1250v0")
+
+    co_list, library_list = CW.clientLibraryFiles(
+        tmp_path, ["gfx1250"], CW.computeOutputArchNames(["gfx1250v0"])
+    )
+
+    assert library_list == [master]
+    assert co_list == [code_object]