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