def load(jitmodule, extension, fn_name, cppargs=[], ldargs=[], argtypes=None, restype=None, compiler=None, comm=None): """Build a shared library and return a function pointer from it. :arg jitmodule: The JIT Module which can generate the code to compile, or the string representing the source code. :arg extension: extension of the source file (c, cpp) :arg fn_name: The name of the function to return from the resulting library :arg cppargs: A list of arguments to the C compiler (optional) :arg ldargs: A list of arguments to the linker (optional) :arg argtypes: A list of ctypes argument types matching the arguments of the returned function (optional, pass ``None`` for ``void``). This is only used when string is passed in instead of JITModule. :arg restype: The return type of the function (optional, pass ``None`` for ``void``). :arg compiler: The name of the C compiler (intel, ``None`` for default). :kwarg comm: Optional communicator to compile the code on (only rank 0 compiles code) (defaults to COMM_WORLD). """ if isinstance(jitmodule, str): class StrCode(object): def __init__(self, code, argtypes): self.code_to_compile = code self.cache_key = (None, code) # We peel off the first # entry, since for a jitmodule, it's a process-local # cache key self.argtypes = argtypes code = StrCode(jitmodule, argtypes) elif isinstance(jitmodule, JITModule): code = jitmodule else: raise ValueError("Don't know how to compile code of type %r" % type(jitmodule)) platform = sys.platform cpp = extension == "cpp" if not compiler: compiler = configuration["compiler"] if platform.find('linux') == 0: if compiler == 'icc': compiler = LinuxIntelCompiler(cppargs, ldargs, cpp=cpp, comm=comm) elif compiler == 'gcc': compiler = LinuxCompiler(cppargs, ldargs, cpp=cpp, comm=comm) else: raise CompilationError("Unrecognized compiler name '%s'" % compiler) elif platform.find('darwin') == 0: compiler = MacCompiler(cppargs, ldargs, cpp=cpp, comm=comm) else: raise CompilationError("Don't know what compiler to use for platform '%s'" % platform) dll = compiler.get_so(code, extension) fn = getattr(dll, fn_name) fn.argtypes = code.argtypes fn.restype = restype return fn
def compilation_comm(comm): """Get a communicator for compilation. :arg comm: The input communicator. :returns: A communicator used for compilation (may be smaller) """ # Should we try and do node-local compilation? if not configuration["node_local_compilation"]: return comm retcomm = get_compilation_comm(comm) if retcomm is not None: debug("Found existing compilation communicator") return retcomm if MPI.VERSION >= 3: debug("Creating compilation communicator using MPI_Split_type") retcomm = comm.Split_type(MPI.COMM_TYPE_SHARED) debug( "Finished creating compilation communicator using MPI_Split_type") set_compilation_comm(comm, retcomm) return retcomm debug("Creating compilation communicator using MPI_Split + filesystem") import tempfile if comm.rank == 0: if not os.path.exists(configuration["cache_dir"]): os.makedirs(configuration["cache_dir"], exist_ok=True) tmpname = tempfile.mkdtemp(prefix="rank-determination-", dir=configuration["cache_dir"]) else: tmpname = None tmpname = comm.bcast(tmpname, root=0) if tmpname is None: raise CompilationError("Cannot determine sharedness of filesystem") # Touch file debug("Made tmpdir %s" % tmpname) with open(os.path.join(tmpname, str(comm.rank)), "wb"): pass comm.barrier() import glob ranks = sorted( int(os.path.basename(name)) for name in glob.glob("%s/[0-9]*" % tmpname)) debug("Creating compilation communicator using filesystem colors") retcomm = comm.Split(color=min(ranks), key=comm.rank) debug("Finished creating compilation communicator using filesystem colors") set_compilation_comm(comm, retcomm) return retcomm
def load(src, extension, fn_name, cppargs=[], ldargs=[], argtypes=None, restype=None, compiler=None, comm=None): """Build a shared library and return a function pointer from it. :arg src: A string containing the source to build :arg extension: extension of the source file (c, cpp) :arg fn_name: The name of the function to return from the resulting library :arg cppargs: A list of arguments to the C compiler (optional) :arg ldargs: A list of arguments to the linker (optional) :arg argtypes: A list of ctypes argument types matching the arguments of the returned function (optional, pass ``None`` for ``void``). :arg restype: The return type of the function (optional, pass ``None`` for ``void``). :arg compiler: The name of the C compiler (intel, ``None`` for default). :kwarg comm: Optional communicator to compile the code on (only rank 0 compiles code) (defaults to COMM_WORLD). """ platform = sys.platform cpp = extension == "cpp" if platform.find('linux') == 0: if compiler == 'intel': compiler = LinuxIntelCompiler(cppargs, ldargs, cpp=cpp, comm=comm) else: compiler = LinuxCompiler(cppargs, ldargs, cpp=cpp, comm=comm) elif platform.find('darwin') == 0: compiler = MacCompiler(cppargs, ldargs, cpp=cpp, comm=comm) else: raise CompilationError( "Don't know what compiler to use for platform '%s'" % platform) dll = compiler.get_so(src, extension) fn = getattr(dll, fn_name) fn.argtypes = argtypes fn.restype = restype return fn
def get_so(self, src, extension): """Build a shared library and load it :arg src: The source string to compile. :arg extension: extension of the source file (c, cpp). Returns a :class:`ctypes.CDLL` object of the resulting shared library.""" # Determine cache key hsh = md5(src.encode()) hsh.update(self._cc.encode()) if self._ld: hsh.update(self._ld.encode()) hsh.update("".join(self._cppargs).encode()) hsh.update("".join(self._ldargs).encode()) basename = hsh.hexdigest() cachedir = configuration['cache_dir'] pid = os.getpid() cname = os.path.join(cachedir, "%s_p%d.%s" % (basename, pid, extension)) oname = os.path.join(cachedir, "%s_p%d.o" % (basename, pid)) soname = os.path.join(cachedir, "%s.so" % basename) # Link into temporary file, then rename to shared library # atomically (avoiding races). tmpname = os.path.join(cachedir, "%s_p%d.so.tmp" % (basename, pid)) if configuration['check_src_hashes'] or configuration['debug']: matching = self.comm.allreduce(basename, op=_check_op) if matching != basename: # Dump all src code to disk for debugging output = os.path.join(cachedir, "mismatching-kernels") srcfile = os.path.join(output, "src-rank%d.c" % self.comm.rank) if self.comm.rank == 0: if not os.path.exists(output): os.makedirs(output, exist_ok=True) self.comm.barrier() with open(srcfile, "w") as f: f.write(src) self.comm.barrier() raise CompilationError("Generated code differs across ranks (see output in %s)" % output) try: # Are we in the cache? return ctypes.CDLL(soname) except OSError: # No, let's go ahead and build if self.comm.rank == 0: # No need to do this on all ranks if not os.path.exists(cachedir): os.makedirs(cachedir, exist_ok=True) logfile = os.path.join(cachedir, "%s_p%d.log" % (basename, pid)) errfile = os.path.join(cachedir, "%s_p%d.err" % (basename, pid)) with progress(INFO, 'Compiling wrapper'): with open(cname, "w") as f: f.write(src) # Compiler also links if self._ld is None: cc = [self._cc] + self._cppargs + \ ['-o', tmpname, cname] + self._ldargs debug('Compilation command: %s', ' '.join(cc)) with open(logfile, "w") as log: with open(errfile, "w") as err: log.write("Compilation command:\n") log.write(" ".join(cc)) log.write("\n\n") try: if configuration['no_fork_available']: cc += ["2>", errfile, ">", logfile] cmd = " ".join(cc) status = os.system(cmd) if status != 0: raise subprocess.CalledProcessError(status, cmd) else: subprocess.check_call(cc, stderr=err, stdout=log) except subprocess.CalledProcessError as e: raise CompilationError( """Command "%s" return error status %d. Unable to compile code Compile log in %s Compile errors in %s""" % (e.cmd, e.returncode, logfile, errfile)) else: cc = [self._cc] + self._cppargs + \ ['-c', '-o', oname, cname] ld = self._ld.split() + ['-o', tmpname, oname] + self._ldargs debug('Compilation command: %s', ' '.join(cc)) debug('Link command: %s', ' '.join(ld)) with open(logfile, "w") as log: with open(errfile, "w") as err: log.write("Compilation command:\n") log.write(" ".join(cc)) log.write("\n\n") log.write("Link command:\n") log.write(" ".join(ld)) log.write("\n\n") try: if configuration['no_fork_available']: cc += ["2>", errfile, ">", logfile] ld += ["2>", errfile, ">", logfile] cccmd = " ".join(cc) ldcmd = " ".join(ld) status = os.system(cccmd) if status != 0: raise subprocess.CalledProcessError(status, cccmd) status = os.system(ldcmd) if status != 0: raise subprocess.CalledProcessError(status, ldcmd) else: subprocess.check_call(cc, stderr=err, stdout=log) subprocess.check_call(ld, stderr=err, stdout=log) except subprocess.CalledProcessError as e: raise CompilationError( """Command "%s" return error status %d. Unable to compile code Compile log in %s Compile errors in %s""" % (e.cmd, e.returncode, logfile, errfile)) # Atomically ensure soname exists os.rename(tmpname, soname) # Wait for compilation to complete self.comm.barrier() # Load resulting library return ctypes.CDLL(soname)