perf(tensilelite): overlap helper-kernel compile with asm codegen

projects/hipblaslt/tensilelite/tensilelite/TensileCreateLibrary/Run.py

@@ -37,6 +37,7 @@
 import pickle
 import sys
 import zlib
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
 from pathlib import Path
 from timeit import default_timer as timer
@@ -1089,36 +1090,44 @@ def composed_function(kernel):
         splitGSU,
         compress=memcompress,
     )
-    results = _mapKernelsByBackend(
-        compose(unaryAssemble, unaryWriteAssembly, unaryProcessKernelSource),
-        compose(unaryAssemble, unaryWriteAssembly, unaryProcessNative),
-        uniqueAsmKernels,
-    )
-
-    buildAssemblyCodeObjectFiles(
-        asmToolchain.linker,
-        asmToolchain.bundler,
-        asmGroups,
-        destRoot,
-        assemblyTmpPath,
-        compress,
-        archNames=archNames,
-    )
 
     # The basename has to stay "Kernels": the emitted code object is named from
     # it, and the runtime opens Kernels.so-000-<arch>.hsaco by that name.
     writeHelpers(buildTmpPath, kernelHelperObjs, KERNEL_HELPER_FILENAME_CPP, KERNEL_HELPER_FILENAME_H)
     srcKernelFile = buildTmpPath / "Kernels.cpp"
 
-    buildSourceCodeObjectFiles(
-        srcToolchain.compiler,
-        srcToolchain.bundler,
-        destRoot,
-        objectTmpPath,
-        buildTmpPath,
-        srcKernelFile,
-        cmdlineArchs,
-    )
+    # The helper kernels in Kernels.cpp do not depend on the assembly kernels,
+    # so compile them (one long single-threaded clang job) while the assembly
+    # kernels are being generated instead of afterwards.
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        srcCodeObjects = executor.submit(
+            buildSourceCodeObjectFiles,
+            srcToolchain.compiler,
+            srcToolchain.bundler,
+            destRoot,
+            objectTmpPath,
+            buildTmpPath,
+            srcKernelFile,
+            cmdlineArchs,
+        )
+
+        results = _mapKernelsByBackend(
+            compose(unaryAssemble, unaryWriteAssembly, unaryProcessKernelSource),
+            compose(unaryAssemble, unaryWriteAssembly, unaryProcessNative),
+            uniqueAsmKernels,
+        )
+
+        buildAssemblyCodeObjectFiles(
+            asmToolchain.linker,
+            asmToolchain.bundler,
+            asmGroups,
+            destRoot,
+            assemblyTmpPath,
+            compress,
+            archNames=archNames,
+        )
+
+        srcCodeObjects.result()
 
     return len(uniqueAsmKernels), uniqueAsmKernels, results
 

projects/hipblaslt/tensilelite/tensilelite/Tests/unit/test_tcl_helper_overlap.py

@@ -0,0 +1,65 @@
+# Copyright Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+"""Unit tests for compiling the helper kernels (Kernels.cpp) while
+TensileCreateLibrary generates the assembly kernels."""
+import threading
+from unittest.mock import MagicMock
+
+import pytest
+
+import tensilelite.TensileCreateLibrary.Run as Run
+
+pytestmark = pytest.mark.unit
+
+
+def _write(tmp_path):
+    return Run.writeSolutionsAndKernelsTCL(
+        str(tmp_path),  # outputPath
+        MagicMock(),  # asmToolchain
+        MagicMock(),  # srcToolchain
+        [],  # solutions
+        [],  # kernels
+        [],  # kernelHelperObjs
+        MagicMock(),  # kernelWriterAssembly
+        ["gfx942"],  # cmdlineArchs
+    )
+
+
+@pytest.fixture
+def builders(monkeypatch):
+    monkeypatch.setattr(Run, "rocisa", MagicMock())
+    monkeypatch.setattr(Run, "writeHelpers", lambda *a, **kw: None)
+    monkeypatch.setattr(Run, "buildAssemblyCodeObjectFiles", lambda *a, **kw: [])
+    return monkeypatch
+
+
+def test_helper_compile_runs_while_assembly_kernels_are_generated(tmp_path, builders):
+    compileStarted = threading.Event()
+    codegenDone = threading.Event()
+    seen = {}
+
+    def compileHelpers(*a, **kw):
+        compileStarted.set()
+        codegenDone.wait(timeout=10)
+        return []
+
+    def generate(*a, **kw):
+        seen["compileStartedDuringCodegen"] = compileStarted.wait(timeout=10)
+        codegenDone.set()
+        return []
+
+    builders.setattr(Run, "buildSourceCodeObjectFiles", compileHelpers)
+    builders.setattr(Run, "ParallelMap2", generate)
+    _write(tmp_path)
+
+    assert seen["compileStartedDuringCodegen"]
+
+
+def test_helper_compile_failure_is_raised(tmp_path, builders):
+    def compileHelpers(*a, **kw):
+        raise RuntimeError("Kernels.cpp failed to compile")
+
+    builders.setattr(Run, "buildSourceCodeObjectFiles", compileHelpers)
+    builders.setattr(Run, "ParallelMap2", lambda *a, **kw: [])
+    with pytest.raises(RuntimeError, match="Kernels.cpp failed to compile"):
+        _write(tmp_path)