gador/nixpkgs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
current
nixpkgs/pkgs/development/rocm-modules/hipblaslt/TensileCreateLibrary-refactor.patch
T
Luna Nova 691e492590 rocmPackages: 7.1.1 -> 7.2.0 
Thanks @GZGavinZhao for providing continued support in the form of
the ISA compat patch set.

Co-Authored-By: Gavin Zhao <git@gzgz.dev>
2026-02-28 21:51:01 -08:00

850 lines
33 KiB
Diff
Raw Permalink Blame History

diff --git a/tensilelite/Tensile/Common/Parallel.py b/tensilelite/Tensile/Common/Parallel.py
index 1a2bf9e119..f46100c7b8 100644
--- a/tensilelite/Tensile/Common/Parallel.py
+++ b/tensilelite/Tensile/Common/Parallel.py
@@ -22,43 +22,58 @@
#
################################################################################
-import concurrent.futures
-import itertools
+import multiprocessing
import os
+import re
import sys
import time
-
-from joblib import Parallel, delayed
+from functools import partial
+from typing import Any, Callable
from .Utilities import tqdm
-def joblibParallelSupportsGenerator():
- import joblib
- from packaging.version import Version
+def get_inherited_job_limit() -> int:
+ # 1. Check CMAKE_BUILD_PARALLEL_LEVEL (CMake 3.12+)
+ if 'CMAKE_BUILD_PARALLEL_LEVEL' in os.environ:
+ try:
+ return int(os.environ['CMAKE_BUILD_PARALLEL_LEVEL'])
+ except ValueError:
+ pass
- joblibVer = joblib.__version__
- return Version(joblibVer) >= Version("1.4.0")
+ # 2. Parse MAKEFLAGS for -jN
+ makeflags = os.environ.get('MAKEFLAGS', '')
+ match = re.search(r'-j\s*(\d+)', makeflags)
+ if match:
+ return int(match.group(1))
+ return -1
-def CPUThreadCount(enable=True):
- from .GlobalParameters import globalParameters
+def CPUThreadCount(enable=True):
if not enable:
return 1
- else:
+ from .GlobalParameters import globalParameters
+
+ # Priority order:
+ # 1. Inherited from build system (CMAKE_BUILD_PARALLEL_LEVEL or MAKEFLAGS)
+ # 2. Explicit --jobs flag
+ # 3. Auto-detect
+ inherited_limit = get_inherited_job_limit()
+ cpuThreads = inherited_limit if inherited_limit > 0 else globalParameters["CpuThreads"]
+
+ if cpuThreads < 1:
if os.name == "nt":
- # Windows supports at most 61 workers because the scheduler uses
- # WaitForMultipleObjects directly, which has the limit (the limit
- # is actually 64, but some handles are needed for accounting).
- cpu_count = min(os.cpu_count(), 61)
+ cpuThreads = os.cpu_count()
else:
- cpu_count = len(os.sched_getaffinity(0))
- cpuThreads = globalParameters["CpuThreads"]
- if cpuThreads == -1:
- return cpu_count
+ cpuThreads = len(os.sched_getaffinity(0))
- return min(cpu_count, cpuThreads)
+ if os.name == "nt":
+ # Windows supports at most 61 workers because the scheduler uses
+ # WaitForMultipleObjects directly, which has the limit (the limit
+ # is actually 64, but some handles are needed for accounting).
+ cpuThreads = min(cpuThreads, 61)
+ return max(1, cpuThreads)
def pcallWithGlobalParamsMultiArg(f, args, newGlobalParameters):
@@ -71,19 +86,22 @@ def pcallWithGlobalParamsSingleArg(f, arg, newGlobalParameters):
return f(arg)
-def apply_print_exception(item, *args):
- # print(item, args)
+def OverwriteGlobalParameters(newGlobalParameters):
+ from . import GlobalParameters
+
+ GlobalParameters.globalParameters.clear()
+ GlobalParameters.globalParameters.update(newGlobalParameters)
+
+
+def worker_function(args, function, multiArg):
+ """Worker function that executes in the pool process."""
try:
- if len(args) > 0:
- func = item
- args = args[0]
- return func(*args)
+ if multiArg:
+ return function(*args)
else:
- func, item = item
- return func(item)
+ return function(args)
except Exception:
import traceback
-
traceback.print_exc()
raise
finally:
@@ -98,154 +116,121 @@ def OverwriteGlobalParameters(newGlobalParameters):
GlobalParameters.globalParameters.update(newGlobalParameters)
-def ProcessingPool(enable=True, maxTasksPerChild=None):
- import multiprocessing
- import multiprocessing.dummy
-
- threadCount = CPUThreadCount()
-
- if (not enable) or threadCount <= 1:
- return multiprocessing.dummy.Pool(1)
-
- if multiprocessing.get_start_method() == "spawn":
- from . import GlobalParameters
-
- return multiprocessing.Pool(
- threadCount,
- initializer=OverwriteGlobalParameters,
- maxtasksperchild=maxTasksPerChild,
- initargs=(GlobalParameters.globalParameters,),
- )
- else:
- return multiprocessing.Pool(threadCount, maxtasksperchild=maxTasksPerChild)
+def progress_logger(iterable, total, message, min_log_interval=5.0):
+ """
+ Generator that wraps an iterable and logs progress with time-based throttling.
+ Only logs progress if at least min_log_interval seconds have passed since last log.
+ Only prints completion message if task took >= min_log_interval seconds.
-def ParallelMap(function, objects, message="", enable=True, method=None, maxTasksPerChild=None):
+ Yields (index, item) tuples.
"""
- Generally equivalent to list(map(function, objects)), possibly executing in parallel.
-
- message: A message describing the operation to be performed.
- enable: May be set to false to disable parallelism.
- method: A function which can fetch the mapping function from a processing pool object.
- Leave blank to use .map(), other possiblities:
- - `lambda x: x.starmap` - useful if `function` takes multiple parameters.
- - `lambda x: x.imap` - lazy evaluation
- - `lambda x: x.imap_unordered` - lazy evaluation, does not preserve order of return value.
- """
- from .GlobalParameters import globalParameters
+ start_time = time.time()
+ last_log_time = start_time
+ log_interval = 1 + (total // 100)
- threadCount = CPUThreadCount(enable)
- pool = ProcessingPool(enable, maxTasksPerChild)
-
- if threadCount <= 1 and globalParameters["ShowProgressBar"]:
- # Provide a progress bar for single-threaded operation.
- # This works for method=None, and for starmap.
- mapFunc = map
- if method is not None:
- # itertools provides starmap which can fill in for pool.starmap. It provides imap on Python 2.7.
- # If this works, we will use it, otherwise we will fallback to the "dummy" pool for single threaded
- # operation.
- try:
- mapFunc = method(itertools)
- except NameError:
- mapFunc = None
-
- if mapFunc is not None:
- return list(mapFunc(function, tqdm(objects, message)))
-
- mapFunc = pool.map
- if method:
- mapFunc = method(pool)
-
- objects = zip(itertools.repeat(function), objects)
- function = apply_print_exception
-
- countMessage = ""
- try:
- countMessage = " for {} tasks".format(len(objects))
- except TypeError:
- pass
+ for idx, item in enumerate(iterable):
+ if idx % log_interval == 0:
+ current_time = time.time()
+ if (current_time - last_log_time) >= min_log_interval:
+ print(f"{message}\t{idx+1: 5d}/{total: 5d}")
+ last_log_time = current_time
+ yield idx, item
- if message != "":
- message += ": "
+ elapsed = time.time() - start_time
+ final_idx = idx + 1 if 'idx' in locals() else 0
- print("{0}Launching {1} threads{2}...".format(message, threadCount, countMessage))
- sys.stdout.flush()
- currentTime = time.time()
- rv = mapFunc(function, objects)
- totalTime = time.time() - currentTime
- print("{0}Done. ({1:.1f} secs elapsed)".format(message, totalTime))
- sys.stdout.flush()
- pool.close()
- return rv
+ if elapsed >= min_log_interval or last_log_time > start_time:
+ print(f"{message} done in {elapsed:.1f}s!\t{final_idx: 5d}/{total: 5d}")
-def ParallelMapReturnAsGenerator(function, objects, message="", enable=True, multiArg=True):
- from .GlobalParameters import globalParameters
+def imap_with_progress(pool, func, iterable, total, message, chunksize):
+ results = []
+ for _, result in progress_logger(pool.imap(func, iterable, chunksize=chunksize), total, message):
+ results.append(result)
+ return results
- threadCount = CPUThreadCount(enable)
- print("{0}Launching {1} threads...".format(message, threadCount))
- if threadCount <= 1 and globalParameters["ShowProgressBar"]:
- # Provide a progress bar for single-threaded operation.
- callFunc = lambda args: function(*args) if multiArg else lambda args: function(args)
- return [callFunc(args) for args in tqdm(objects, message)]
+def _ParallelMap_generator(worker, objects, objLen, message, chunksize, threadCount, globalParameters, maxtasksperchild):
+ # separate fn because yield makes the entire fn a generator even if unreachable
+ ctx = multiprocessing.get_context('forkserver' if os.name != 'nt' else 'spawn')
- with concurrent.futures.ProcessPoolExecutor(max_workers=threadCount) as executor:
- resultFutures = (executor.submit(function, *arg if multiArg else arg) for arg in objects)
- for result in concurrent.futures.as_completed(resultFutures):
- yield result.result()
+ with ctx.Pool(processes=threadCount, maxtasksperchild=maxtasksperchild,
+ initializer=OverwriteGlobalParameters, initargs=(globalParameters,)) as pool:
+ for _, result in progress_logger(pool.imap_unordered(worker, objects, chunksize=chunksize), objLen, message):
+ yield result
def ParallelMap2(
- function, objects, message="", enable=True, multiArg=True, return_as="list", procs=None
+ function: Callable,
+ objects: Any,
+ message: str = "",
+ enable: bool = True,
+ multiArg: bool = True,
+ minChunkSize: int = 1,
+ maxWorkers: int = -1,
+ maxtasksperchild: int = 1024,
+ return_as: str = "list"
):
+ """Executes a function over a list of objects in parallel or sequentially.
+
+ This function is generally equivalent to ``list(map(function, objects))``. However, it provides
+ additional functionality to run in parallel, depending on the 'enable' flag and available CPU
+ threads.
+
+ Args:
+ function: The function to apply to each item in 'objects'. If 'multiArg' is True, 'function'
+ should accept multiple arguments.
+ objects: An iterable of objects to be processed by 'function'. If 'multiArg' is True, each
+ item in 'objects' should be an iterable of arguments for 'function'.
+ message: Optional; a message describing the operation. Default is an empty string.
+ enable: Optional; if False, disables parallel execution and runs sequentially. Default is True.
+ multiArg: Optional; if True, treats each item in 'objects' as multiple arguments for
+ 'function'. Default is True.
+ return_as: Optional; "list" (default) or "generator_unordered" for streaming results
+
+ Returns:
+ A list or generator containing the results of applying **function** to each item in **objects**.
"""
- Generally equivalent to list(map(function, objects)), possibly executing in parallel.
+ from .GlobalParameters import globalParameters
- message: A message describing the operation to be performed.
- enable: May be set to false to disable parallelism.
- multiArg: True if objects represent multiple arguments
- (differentiates multi args vs single collection arg)
- """
- if return_as in ("generator", "generator_unordered") and not joblibParallelSupportsGenerator():
- return ParallelMapReturnAsGenerator(function, objects, message, enable, multiArg)
+ threadCount = CPUThreadCount(enable)
- from .GlobalParameters import globalParameters
+ if not hasattr(objects, "__len__"):
+ objects = list(objects)
- threadCount = procs if procs else CPUThreadCount(enable)
+ objLen = len(objects)
+ if objLen == 0:
+ return [] if return_as == "list" else iter([])
- threadCount = CPUThreadCount(enable)
+ f = (lambda x: function(*x)) if multiArg else function
+ if objLen == 1:
+ print(f"{message}: (1 task)")
+ result = [f(x) for x in objects]
+ return result if return_as == "list" else iter(result)
- if threadCount <= 1 and globalParameters["ShowProgressBar"]:
- # Provide a progress bar for single-threaded operation.
- return [function(*args) if multiArg else function(args) for args in tqdm(objects, message)]
+ extra_message = (
+ f": {threadCount} thread(s)" + f", {objLen} tasks"
+ if objLen
+ else ""
+ )
- countMessage = ""
- try:
- countMessage = " for {} tasks".format(len(objects))
- except TypeError:
- pass
-
- if message != "":
- message += ": "
- print("{0}Launching {1} threads{2}...".format(message, threadCount, countMessage))
- sys.stdout.flush()
- currentTime = time.time()
-
- pcall = pcallWithGlobalParamsMultiArg if multiArg else pcallWithGlobalParamsSingleArg
- pargs = zip(objects, itertools.repeat(globalParameters))
-
- if joblibParallelSupportsGenerator():
- rv = Parallel(n_jobs=threadCount, timeout=99999, return_as=return_as)(
- delayed(pcall)(function, a, params) for a, params in pargs
- )
+ print(f"ParallelMap {message}{extra_message}")
+
+ if threadCount <= 1:
+ result = [f(x) for x in objects]
+ return result if return_as == "list" else iter(result)
+
+ if maxWorkers > 0:
+ threadCount = min(maxWorkers, threadCount)
+
+ chunksize = max(minChunkSize, objLen // 2000)
+ worker = partial(worker_function, function=function, multiArg=multiArg)
+ if return_as == "generator_unordered":
+ # yield results as they complete without buffering
+ return _ParallelMap_generator(worker, objects, objLen, message, chunksize, threadCount, globalParameters, maxtasksperchild)
else:
- rv = Parallel(n_jobs=threadCount, timeout=99999)(
- delayed(pcall)(function, a, params) for a, params in pargs
- )
-
- totalTime = time.time() - currentTime
- print("{0}Done. ({1:.1f} secs elapsed)".format(message, totalTime))
- sys.stdout.flush()
- return rv
+ ctx = multiprocessing.get_context('forkserver' if os.name != 'nt' else 'spawn')
+ with ctx.Pool(processes=threadCount, maxtasksperchild=maxtasksperchild,
+ initializer=OverwriteGlobalParameters, initargs=(globalParameters,)) as pool:
+ return list(imap_with_progress(pool, worker, objects, objLen, message, chunksize))
diff --git a/tensilelite/Tensile/CustomKernels.py b/tensilelite/Tensile/CustomKernels.py
index ffceb636f5..127b3386a1 100644
--- a/tensilelite/Tensile/CustomKernels.py
+++ b/tensilelite/Tensile/CustomKernels.py
@@ -24,7 +24,9 @@
from . import CUSTOM_KERNEL_PATH
from Tensile.Common.ValidParameters import checkParametersAreValid, validParameters, newMIValidParameters
+from Tensile.CustomYamlLoader import DEFAULT_YAML_LOADER
+from functools import lru_cache
import yaml
import os
@@ -58,10 +60,13 @@ def getCustomKernelConfigAndAssembly(name, directory=CUSTOM_KERNEL_PATH):
return (config, assembly)
+# getCustomKernelConfig will get called repeatedly on the same file
+# 20x logic loading speedup for aquavanjaram_Cijk_Ailk_Bljk_F8NH_HHS_BH_Bias_HAS_SAB_SAV_freesize_custom_GSUs
+@lru_cache
def readCustomKernelConfig(name, directory=CUSTOM_KERNEL_PATH):
rawConfig, _ = getCustomKernelConfigAndAssembly(name, directory)
try:
- return yaml.safe_load(rawConfig)["custom.config"]
+ return yaml.load(rawConfig, Loader=DEFAULT_YAML_LOADER)["custom.config"]
except yaml.scanner.ScannerError as e:
raise RuntimeError("Failed to read configuration for custom kernel: {0}\nDetails:\n{1}".format(name, e))
diff --git a/tensilelite/Tensile/CustomYamlLoader.py b/tensilelite/Tensile/CustomYamlLoader.py
index e03f456fbe..ed7510f2ce 100644
--- a/tensilelite/Tensile/CustomYamlLoader.py
+++ b/tensilelite/Tensile/CustomYamlLoader.py
@@ -1,6 +1,7 @@
# Copyright © Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
+import sys
import yaml
from pathlib import Path
@@ -70,7 +71,7 @@ def parse_scalar(loader: yaml.Loader):
elif is_float(value_lower):
return float(value_lower)
- return value
+ return sys.intern(value)
def load_yaml_stream(yaml_path: Path, loader_type: yaml.Loader):
with open(yaml_path, 'r') as f:
diff --git a/tensilelite/Tensile/TensileCreateLibrary/Run.py b/tensilelite/Tensile/TensileCreateLibrary/Run.py
index 22d19851a3..348068b3bf 100644
--- a/tensilelite/Tensile/TensileCreateLibrary/Run.py
+++ b/tensilelite/Tensile/TensileCreateLibrary/Run.py
@@ -26,8 +26,10 @@ import rocisa
import functools
import glob
+import gc
import itertools
import os
+import resource
import shutil
from pathlib import Path
from timeit import default_timer as timer
@@ -78,6 +80,25 @@ from Tensile.Utilities.Decorators.Timing import timing
from .ParseArguments import parseArguments
+def getMemoryUsage():
+ """Get peak and current memory usage in MB."""
+ rusage = resource.getrusage(resource.RUSAGE_SELF)
+ peak_memory_mb = rusage.ru_maxrss / 1024 # KB to MB on Linux
+
+ # Get current memory from /proc/self/status
+ current_memory_mb = 0
+ try:
+ with open('/proc/self/status') as f:
+ for line in f:
+ if line.startswith('VmRSS:'):
+ current_memory_mb = int(line.split()[1]) / 1024 # KB to MB
+ break
+ except:
+ current_memory_mb = peak_memory_mb # Fallback
+
+ return (peak_memory_mb, current_memory_mb)
+
+
class KernelCodeGenResult(NamedTuple):
err: int
src: str
@@ -115,6 +136,29 @@ def processKernelSource(kernelWriterAssembly, data, outOptions, splitGSU, kernel
)
+def processAndAssembleKernelTCL(kernelWriterAssembly, rocisa_data, outOptions, splitGSU, kernel, assemblyTmpPath, assembler):
+ """
+ Pipeline function for TCL mode that:
+ 1. Generates kernel source
+ 2. Writes .s file to disk
+ 3. Assembles to .o file
+ 4. Deletes .s file
+ """
+ result = processKernelSource(kernelWriterAssembly, rocisa_data, outOptions, splitGSU, kernel)
+ return writeAndAssembleKernel(result, assemblyTmpPath, assembler)
+
+
+def writeMasterSolutionLibrary(name_lib_tuple, newLibraryDir, splitGSU, libraryFormat):
+ """
+ Write a master solution library to disk.
+ Module-level function to support multiprocessing.
+ """
+ name, lib = name_lib_tuple
+ filename = os.path.join(newLibraryDir, name)
+ lib.applyNaming(splitGSU)
+ LibraryIO.write(filename, state(lib), libraryFormat)
+
+
def removeInvalidSolutionsAndKernels(results, kernels, solutions, errorTolerant, printLevel: bool, splitGSU: bool):
removeKernels = []
removeKernelNames = []
@@ -189,6 +233,24 @@ def writeAssembly(asmPath: Union[Path, str], result: KernelCodeGenResult):
return path, isa, wfsize, minResult
+def writeAndAssembleKernel(result: KernelCodeGenResult, asmPath: Union[Path, str], assembler):
+ """Write assembly file and immediately assemble it to .o file"""
+ if result.err:
+ printExit(f"Failed to build kernel {result.name} because it has error code {result.err}")
+
+ path = Path(asmPath) / f"{result.name}.s"
+ with open(path, "w", encoding="utf-8") as f:
+ f.write(result.src)
+
+ # Assemble .s -> .o
+ assembler(isaToGfx(result.isa), result.wavefrontSize, str(path), str(path.with_suffix(".o")))
+
+ # Delete assembly file immediately to save disk space
+ path.unlink()
+
+ return KernelMinResult(result.err, result.cuoccupancy, result.pgr, result.mathclk)
+
+
def writeHelpers(
outputPath, kernelHelperObjs, KERNEL_HELPER_FILENAME_CPP, KERNEL_HELPER_FILENAME_H
):
@@ -272,14 +334,15 @@ def writeSolutionsAndKernels(
numAsmKernels = len(asmKernels)
numKernels = len(asmKernels)
assert numKernels == numAsmKernels, "Only assembly kernels are supported in TensileLite"
- asmIter = zip(
- itertools.repeat(kernelWriterAssembly),
- itertools.repeat(rocisa.rocIsa.getInstance().getData()),
- itertools.repeat(outOptions),
- itertools.repeat(splitGSU),
- asmKernels
+
+ processKernelFn = functools.partial(
+ processKernelSource,
+ kernelWriterAssembly=kernelWriterAssembly,
+ data=rocisa.rocIsa.getInstance().getData(),
+ outOptions=outOptions,
+ splitGSU=splitGSU
)
- asmResults = ParallelMap2(processKernelSource, asmIter, "Generating assembly kernels", return_as="list")
+ asmResults = ParallelMap2(processKernelFn, asmKernels, "Generating assembly kernels", return_as="list", multiArg=False)
removeInvalidSolutionsAndKernels(
asmResults, asmKernels, solutions, errorTolerant, getVerbosity(), splitGSU
)
@@ -287,19 +350,21 @@ def writeSolutionsAndKernels(
asmResults, asmKernels, solutions, splitGSU
)
- def assemble(ret):
- p, isa, wavefrontsize, _ = ret
- asmToolchain.assembler(isaToGfx(isa), wavefrontsize, str(p), str(p.with_suffix(".o")))
-
- unaryWriteAssembly = functools.partial(writeAssembly, assemblyTmpPath)
- compose = lambda *F: functools.reduce(lambda f, g: lambda x: f(g(x)), F)
+ # Use functools.partial to bind assemblyTmpPath and assembler
+ writeAndAssembleFn = functools.partial(
+ writeAndAssembleKernel,
+ asmPath=assemblyTmpPath,
+ assembler=asmToolchain.assembler
+ )
ret = ParallelMap2(
- compose(assemble, unaryWriteAssembly),
+ writeAndAssembleFn,
asmResults,
"Writing assembly kernels",
return_as="list",
multiArg=False,
)
+ del asmResults
+ gc.collect()
writeHelpers(outputPath, kernelHelperObjs, KERNEL_HELPER_FILENAME_CPP, KERNEL_HELPER_FILENAME_H)
srcKernelFile = Path(outputPath) / "Kernels.cpp"
@@ -376,40 +441,35 @@ def writeSolutionsAndKernelsTCL(
uniqueAsmKernels = [k for k in asmKernels if not k.duplicate]
- def assemble(ret, removeTemporaries: bool):
- asmPath, isa, wavefrontsize, result = ret
- asmToolchain.assembler(isaToGfx(isa), wavefrontsize, str(asmPath), str(asmPath.with_suffix(".o")))
- if removeTemporaries:
- asmPath.unlink()
- return result
-
- unaryAssemble = functools.partial(assemble, removeTemporaries=removeTemporaries)
-
outOptions = rocisa.rocIsa.getInstance().getOutputOptions()
outOptions.outputNoComment = not disableAsmComments
- unaryProcessKernelSource = functools.partial(
- processKernelSource,
+ processKernelFn = functools.partial(
+ processAndAssembleKernelTCL,
kernelWriterAssembly,
rocisa.rocIsa.getInstance().getData(),
outOptions,
splitGSU,
+ assemblyTmpPath=assemblyTmpPath,
+ assembler=asmToolchain.assembler
)
- unaryWriteAssembly = functools.partial(writeAssembly, assemblyTmpPath)
- compose = lambda *F: functools.reduce(lambda f, g: lambda x: f(g(x)), F)
- ret = ParallelMap2(
- compose(unaryAssemble, unaryWriteAssembly, unaryProcessKernelSource),
+ results = ParallelMap2(
+ processKernelFn,
uniqueAsmKernels,
"Generating assembly kernels",
multiArg=False,
return_as="list"
)
+ del processKernelFn
+ gc.collect()
+
passPostKernelInfoToSolution(
- ret, uniqueAsmKernels, solutions, splitGSU
+ results, uniqueAsmKernels, solutions, splitGSU
)
- # result.src is very large so let garbage collector know to clean up
- del ret
+ del results
+ gc.collect()
+
buildAssemblyCodeObjectFiles(
asmToolchain.linker,
asmToolchain.bundler,
@@ -508,6 +568,15 @@ def generateKernelHelperObjects(solutions: List[Solution], cxxCompiler: str, isa
return sorted(khos, key=sortByEnum, reverse=True) # Ensure that we write Enum kernel helpers are first in list
+def libraryIter(lib: MasterSolutionLibrary):
+ if len(lib.solutions):
+ for i, s in enumerate(lib.solutions.items()):
+ yield (i, *s)
+ else:
+ for _, lazyLib in lib.lazyLibraries.items():
+ yield from libraryIter(lazyLib)
+
+
@timing
def generateLogicDataAndSolutions(logicFiles, args, assembler: Assembler, isaInfoMap):
@@ -523,26 +592,23 @@ def generateLogicDataAndSolutions(logicFiles, args, assembler: Assembler, isaInf
printSolutionRejectionReason = True
printIndexAssignmentInfo = False
- fIter = zip(
- logicFiles,
- itertools.repeat(assembler),
- itertools.repeat(splitGSU),
- itertools.repeat(printSolutionRejectionReason),
- itertools.repeat(printIndexAssignmentInfo),
- itertools.repeat(isaInfoMap),
- itertools.repeat(args["LazyLibraryLoading"]),
+ parseLogicFn = functools.partial(
+ LibraryIO.parseLibraryLogicFile,
+ assembler=assembler,
+ splitGSU=splitGSU,
+ printSolutionRejectionReason=printSolutionRejectionReason,
+ printIndexAssignmentInfo=printIndexAssignmentInfo,
+ isaInfoMap=isaInfoMap,
+ lazyLibraryLoading=args["LazyLibraryLoading"]
)
- def libraryIter(lib: MasterSolutionLibrary):
- if len(lib.solutions):
- for i, s in enumerate(lib.solutions.items()):
- yield (i, *s)
- else:
- for _, lazyLib in lib.lazyLibraries.items():
- yield from libraryIter(lazyLib)
-
for library in ParallelMap2(
- LibraryIO.parseLibraryLogicFile, fIter, "Loading Logics...", return_as="generator_unordered"
+ parseLogicFn, logicFiles, "Loading Logics...",
+ return_as="generator_unordered",
+ minChunkSize=24,
+ maxWorkers=32,
+ maxtasksperchild=1,
+ multiArg=False,
):
_, architectureName, _, _, _, newLibrary = library
@@ -554,6 +620,9 @@ def generateLogicDataAndSolutions(logicFiles, args, assembler: Assembler, isaInf
else:
masterLibraries[architectureName] = newLibrary
masterLibraries[architectureName].version = args["CodeObjectVersion"]
+ del library, newLibrary
+
+ gc.collect()
# Sort masterLibraries to make global soln index values deterministic
solnReIndex = 0
@@ -751,6 +820,9 @@ def run():
)
stop_wsk = timer()
print(f"Time to generate kernels (s): {(stop_wsk-start_wsk):3.2f}")
+ numKernelHelperObjs = len(kernelHelperObjs)
+ del kernelWriterAssembly, kernelHelperObjs
+ gc.collect()
archs = [ # is this really different than the other archs above?
isaToGfx(arch)
@@ -768,13 +840,10 @@ def run():
if kName not in solDict:
solDict["%s"%kName] = kernel
- def writeMsl(name, lib):
- filename = os.path.join(newLibraryDir, name)
- lib.applyNaming(splitGSU)
- LibraryIO.write(filename, state(lib), arguments["LibraryFormat"])
-
filename = os.path.join(newLibraryDir, "TensileLiteLibrary_lazy_Mapping")
LibraryIO.write(filename, libraryMapping, "msgpack")
+ del libraryMapping
+ gc.collect()
start_msl = timer()
for archName, newMasterLibrary in masterLibraries.items():
@@ -791,12 +860,22 @@ def run():
kName = getKeyNoInternalArgs(s.originalSolution, splitGSU)
s.sizeMapping.CUOccupancy = solDict["%s"%kName]["CUOccupancy"]
- ParallelMap2(writeMsl,
+ writeFn = functools.partial(
+ writeMasterSolutionLibrary,
+ newLibraryDir=newLibraryDir,
+ splitGSU=splitGSU,
+ libraryFormat=arguments["LibraryFormat"]
+ )
+
+ ParallelMap2(writeFn,
newMasterLibrary.lazyLibraries.items(),
"Writing master solution libraries",
+ multiArg=False,
return_as="list")
stop_msl = timer()
print(f"Time to write master solution libraries (s): {(stop_msl-start_msl):3.2f}")
+ del masterLibraries, solutions, kernels, solDict
+ gc.collect()
if not arguments["KeepBuildTmp"]:
buildTmp = Path(arguments["OutputPath"]).parent / "library" / "build_tmp"
@@ -813,8 +892,11 @@ def run():
print("")
stop = timer()
+ peak_memory_mb, current_memory_mb = getMemoryUsage()
print(f"Total time (s): {(stop-start):3.2f}")
print(f"Total kernels processed: {numKernels}")
print(f"Kernels processed per second: {(numKernels/(stop-start)):3.2f}")
- print(f"KernelHelperObjs: {len(kernelHelperObjs)}")
+ print(f"KernelHelperObjs: {numKernelHelperObjs}")
+ print(f"Peak memory usage (MB): {peak_memory_mb:,.1f}")
+ print(f"Current memory usage (MB): {current_memory_mb:,.1f}")
diff --git a/tensilelite/Tensile/TensileMergeLibrary.py b/tensilelite/Tensile/TensileMergeLibrary.py
index e33c617b6f..ba163e9918 100644
--- a/tensilelite/Tensile/TensileMergeLibrary.py
+++ b/tensilelite/Tensile/TensileMergeLibrary.py
@@ -303,8 +303,7 @@ def avoidRegressions(originalDir, incrementalDir, outputPath, forceMerge, noEff=
logicsFiles[origFile] = origFile
logicsFiles[incFile] = incFile
- iters = zip(logicsFiles.keys())
- logicsList = ParallelMap2(loadData, iters, "Loading Logics...", return_as="list")
+ logicsList = ParallelMap2(loadData, logicsFiles.keys(), "Loading Logics...", return_as="list", multiArg=False)
logicsDict = {}
for i, _ in enumerate(logicsList):
logicsDict[logicsList[i][0]] = logicsList[i][1]
diff --git a/tensilelite/Tensile/TensileUpdateLibrary.py b/tensilelite/Tensile/TensileUpdateLibrary.py
index 5ff265d0ed..c1803a6349 100644
--- a/tensilelite/Tensile/TensileUpdateLibrary.py
+++ b/tensilelite/Tensile/TensileUpdateLibrary.py
@@ -26,7 +26,7 @@ from . import LibraryIO
from .Tensile import addCommonArguments, argUpdatedGlobalParameters
from .Common import assignGlobalParameters, print1, restoreDefaultGlobalParameters, HR, \
- globalParameters, architectureMap, ensurePath, ParallelMap, __version__
+ globalParameters, architectureMap, ensurePath, ParallelMap2, __version__
import argparse
import copy
@@ -149,7 +149,7 @@ def TensileUpdateLibrary(userArgs):
for logicFile in logicFiles:
print("# %s" % logicFile)
fIter = zip(logicFiles, itertools.repeat(args.logic_path), itertools.repeat(outputPath))
- libraries = ParallelMap(UpdateLogic, fIter, "Updating logic files", method=lambda x: x.starmap)
+ libraries = ParallelMap2(UpdateLogic, fIter, "Updating logic files", multiArg=True, return_as="list")
def main():
diff --git a/tensilelite/Tensile/Toolchain/Assembly.py b/tensilelite/Tensile/Toolchain/Assembly.py
index a8b91e8d62..265e1d532c 100644
--- a/tensilelite/Tensile/Toolchain/Assembly.py
+++ b/tensilelite/Tensile/Toolchain/Assembly.py
@@ -30,7 +30,7 @@ import subprocess
from pathlib import Path
from typing import List, Union, NamedTuple
-from Tensile.Common import print2
+from Tensile.Common import print1, print2
from Tensile.Common.Architectures import isaToGfx
from ..SolutionStructs import Solution
@@ -92,8 +92,26 @@ def buildAssemblyCodeObjectFiles(
if coName:
coFileMap[asmDir / (coName + extCoRaw)].add(str(asmDir / (kernel["BaseName"] + extObj)))
+ # Build reference count map for .o files to handle shared object files
+ # (.o files from kernels marked .duplicate in TensileCreateLibrary)
+ objFileRefCount = collections.Counter()
+ for coFileRaw, objFiles in coFileMap.items():
+ for objFile in objFiles:
+ objFileRefCount[objFile] += 1
+
+ sharedObjFiles = {objFile: count for objFile, count in objFileRefCount.items() if count > 1}
+ if sharedObjFiles:
+ print1(f"Found {len(sharedObjFiles)} .o files shared across multiple code objects:")
+
for coFileRaw, objFiles in coFileMap.items():
linker(objFiles, str(coFileRaw))
+
+ # Delete .o files after linking once usage count reaches 0
+ for objFile in objFiles:
+ objFileRefCount[objFile] -= 1
+ if objFileRefCount[objFile] == 0:
+ Path(objFile).unlink()
+
coFile = destDir / coFileRaw.name.replace(extCoRaw, extCo)
if compress:
bundler.compress(str(coFileRaw), str(coFile), gfx)
diff --git a/tensilelite/Tensile/Toolchain/Component.py b/tensilelite/Tensile/Toolchain/Component.py
index 67fa35e2d8..dde83af4c3 100644
--- a/tensilelite/Tensile/Toolchain/Component.py
+++ b/tensilelite/Tensile/Toolchain/Component.py
@@ -355,6 +355,7 @@ class Linker(Component):
when invoking the linker, LLVM allows the provision of arguments via a "response file"
Reference: https://llvm.org/docs/CommandLine.html#response-files
"""
+ # FIXME: this prevents threading as clang_args.txt is overwritten
with open(Path.cwd() / "clang_args.txt", "wt") as file:
file.write(" ".join(srcPaths).replace('\\', '\\\\') if os_name == "nt" else " ".join(srcPaths))
return [*(self.default_args), "-o", destPath, "@clang_args.txt"]
diff --git a/tensilelite/requirements.txt b/tensilelite/requirements.txt
index 60c4c11445..5c8fd66a88 100644
--- a/tensilelite/requirements.txt
+++ b/tensilelite/requirements.txt
@@ -2,8 +2,6 @@ dataclasses; python_version == '3.6'
packaging
pyyaml
msgpack
-joblib>=1.4.0; python_version >= '3.8'
-joblib>=1.1.1; python_version < '3.8'
simplejson
ujson
orjson
Reference in New Issue View Git Blame Copy Permalink