Something went wrong. Try again.
lol
Something went wrong. Try again.
34 kB · 849 lines
at master
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850diff --git a/tensilelite/Tensile/Common/Parallel.py b/tensilelite/Tensile/Common/Parallel.pyindex 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.pyindex 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.pyindex 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.pyindex 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.pyindex 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.pyindex 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.pyindex 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.pyindex 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.txtindex 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