You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
339 lines
11 KiB
339 lines
11 KiB
from IPython.core.magic import register_cell_magic, needs_local_scope
|
|
|
|
import importlib
|
|
import glob
|
|
import sys, os
|
|
import random
|
|
import string
|
|
import hashlib
|
|
import shutil
|
|
import subprocess
|
|
import shlex
|
|
|
|
from lief import parse, ELF
|
|
|
|
compiler_cmd = "c++"
|
|
|
|
optimize = "-O3"
|
|
|
|
cmd_flags = ["-std=c++11", "$(pkg-config --cflags eigen3)", "-I."] #, "-Wall"
|
|
link_flags = [""]
|
|
|
|
template = """
|
|
#include <pybind11/pybind11.h>
|
|
namespace py = pybind11;
|
|
|
|
{cell_code}
|
|
|
|
PYBIND11_MODULE({name}, m)
|
|
{defs}
|
|
"""
|
|
|
|
def_template = '\nm.def("{function}", &{function});\n'
|
|
|
|
TMP_PATH = "pbm_modules"
|
|
|
|
if os.path.exists(TMP_PATH):
|
|
shutil.rmtree(TMP_PATH)
|
|
|
|
BUILD_DIR = os.path.join(TMP_PATH, "build")
|
|
MODULES_DIR = os.path.join(TMP_PATH, "modules")
|
|
LIB_DIR = os.path.join(TMP_PATH, "lib")
|
|
# Make sure the correct python executeable is used using cmake
|
|
os.environ["PATH"] = os.path.dirname(sys.executable) + os.pathsep + os.environ["PATH"]
|
|
|
|
# def run_command(command):
|
|
# process = subprocess.Popen(shlex.split(command), stdout=subprocess.PIPE)
|
|
# while True:
|
|
# output = process.stdout.readline()
|
|
# if output == '' and process.poll() is not None:
|
|
# break
|
|
# if output:
|
|
# print output.strip()
|
|
# rc = process.poll()
|
|
# return rc
|
|
|
|
def exec_command(cmd, cwd=None,**kwargs):
|
|
if isinstance(cmd, str):
|
|
cmd = shlex.split(cmd)
|
|
try:
|
|
output = subprocess.check_output(cmd,cwd=cwd,stderr=subprocess.STDOUT,**kwargs)
|
|
if isinstance(output,bytes):
|
|
return False, output.decode("utf-8")
|
|
return False, output
|
|
except subprocess.CalledProcessError as e:
|
|
if isinstance(e.output,bytes):
|
|
return True, e.output.decode("utf-8")
|
|
return True, e.output
|
|
|
|
|
|
def ensure_tmp_folder(path=TMP_PATH):
|
|
path = os.path.join(os.getcwd(), path)
|
|
if not os.path.exists(path):
|
|
os.makedirs(path)
|
|
|
|
|
|
def ensure_pybind(path, force=False):
|
|
path = os.path.join(path,"pybind11")
|
|
exists = os.path.exists(path)
|
|
if exists and force:
|
|
shutil.rmtree(os.path.join(path,"pybind11"))
|
|
if not exists:
|
|
err, output = exec_command(["git", "clone", "https://github.com/pybind/pybind11.git", path])
|
|
|
|
|
|
def ensure_build_dir(path, clear=False):
|
|
path = os.path.join(path, "build")
|
|
exists = os.path.exists(path)
|
|
if exists and clear:
|
|
shutil.rmtree(path)
|
|
if not exists or clear:
|
|
os.makedirs(path)
|
|
return exists and not clear
|
|
|
|
|
|
def parse_args(line):
|
|
args = line.split(" ")
|
|
ret = {}
|
|
ret["cmake_path"] = None
|
|
ret["functions"] = []
|
|
ret["rebuild"] = False
|
|
# TODO: Add full argument parser
|
|
for i,arg in enumerate(args):
|
|
if arg == "-c":
|
|
ret["cmake_path"] = args[i+1].strip()
|
|
if arg == "-f":
|
|
ret["functions"].append(args[i+1].strip())
|
|
if arg == "-rebuild":
|
|
ret["rebuild"] = True
|
|
return ret
|
|
|
|
|
|
def get_CMakeLists(args):
|
|
CMakeLists = ""
|
|
if args["cmake_path"] is not None:
|
|
with open(os.path.join(args["cmake_path"],"CMakeLists.txt"), 'r') as f:
|
|
CMakeLists = f.read()
|
|
return CMakeLists
|
|
|
|
|
|
def cleanup(args, name, CMakeLists):
|
|
if args["cmake_path"] is not None:
|
|
with open(os.path.join(args["cmake_path"],"CMakeLists.txt"), 'w') as f:
|
|
f.write(CMakeLists)
|
|
if os.path.exists(os.path.join(args["path"],f"{name}.cpp")):
|
|
os.remove(os.path.join(args["path"],f"{name}.cpp"))
|
|
|
|
|
|
def get_gcc_cmd(name):
|
|
pyexe = sys.executable
|
|
return f"{compiler_cmd} {' '.join(cmd_flags)} {optimize} -fPIC $({pyexe} -m pybind11 --includes) -c {BUILD_DIR}/{name}.cpp -o {BUILD_DIR}/{name}.o"
|
|
|
|
|
|
def get_linker_cmd(name, others=None):
|
|
if others is None:
|
|
others = " ".join([f"-l{os.path.basename(n)[3:-3]}" for n in sorted(glob.glob(f"{LIB_DIR}/*.so"), key=os.path.getmtime)[::-1]])
|
|
return f"{compiler_cmd} -Wl,-rpath {LIB_DIR} -shared {BUILD_DIR}/{name}.o -o {MODULES_DIR}/{name}.so -L{LIB_DIR} {others} {' '.join(link_flags)}"
|
|
|
|
|
|
def import_to_ip(name):
|
|
# get a handle on the module
|
|
mdl = importlib.import_module(name)
|
|
# is there an __all__? if so respect it
|
|
if "__all__" in mdl.__dict__:
|
|
names = mdl.__dict__["__all__"]
|
|
else:
|
|
# otherwise we import all names that don't begin with _
|
|
names = [x for x in mdl.__dict__ if not x.startswith("_")]
|
|
get_ipython().push({k: getattr(mdl, k) for k in names})
|
|
|
|
|
|
def get_importexport(file):
|
|
binary = parse(file)
|
|
symbols = binary.symbols
|
|
imported = set()
|
|
exported = set()
|
|
for i,symbol in enumerate(symbols):
|
|
if not symbol.imported and not symbol.exported and symbol.binding == ELF.SYMBOL_BINDINGS.GLOBAL:
|
|
imported.add(symbol.name)
|
|
if symbol.exported and symbol.binding == ELF.SYMBOL_BINDINGS.GLOBAL:
|
|
exported.add(symbol.name)
|
|
return imported, exported
|
|
|
|
|
|
export_graph = {}
|
|
import_graph = {}
|
|
symbols = {}
|
|
|
|
def update_dependencies(fresh_compiled):
|
|
global export_graph, import_graph
|
|
imports, exports = get_importexport(fresh_compiled)
|
|
symbols[fresh_compiled] = (imports, exports)
|
|
for inp in imports:
|
|
if inp.startswith("Py") or inp.startswith("_Py"):
|
|
continue
|
|
if inp in import_graph:
|
|
import_graph[inp].append(fresh_compiled)
|
|
else:
|
|
import_graph[inp] = [fresh_compiled]
|
|
|
|
for ex in exports:
|
|
if ex.startswith("PyInit"):
|
|
continue
|
|
export_graph[ex] = fresh_compiled
|
|
|
|
|
|
def get_list_of_rebuilds(fresh_compiled, need_recompilation=None):
|
|
global export_graph, import_graph
|
|
if need_recompilation is None:
|
|
need_recompilation = set()
|
|
need_recompilation.add(fresh_compiled)
|
|
exports = symbols[fresh_compiled][1]
|
|
for ex in exports:
|
|
if ex in import_graph:
|
|
for lib in import_graph[ex]:
|
|
if lib not in need_recompilation:
|
|
need_recompilation.add(lib)
|
|
# recursive collect all libs, that depend on this
|
|
get_list_of_rebuilds(lib, need_recompilation)
|
|
return need_recompilation
|
|
|
|
def get_list_of_links(fresh_compiled):
|
|
global export_graph, import_graph
|
|
imports = symbols[fresh_compiled][0]
|
|
res = []
|
|
for imp in imports:
|
|
if imp in export_graph:
|
|
lib = export_graph[imp]
|
|
if lib not in res:
|
|
res.append(lib)
|
|
return res
|
|
|
|
|
|
def bulk_import(need_reimport=None):
|
|
if need_reimport is None:
|
|
need_reimport = sorted(glob.glob(f"{MODULES_DIR}/*.so"), key=os.path.getmtime)
|
|
|
|
# create unique folder
|
|
rnd = ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(5))
|
|
rnd_path = os.path.join(TMP_PATH, rnd)
|
|
if not os.path.exists(rnd_path):
|
|
os.makedirs(rnd_path)
|
|
with open(os.path.join(rnd_path,"__init__.py"), 'w+') as f:
|
|
pass
|
|
|
|
for lib_path in need_reimport:
|
|
lib = os.path.basename(lib_path)[:-3]
|
|
shutil.copyfile(lib_path, os.path.join(rnd_path,f"{lib}.so"))
|
|
|
|
for lib_path in need_reimport:
|
|
lib = os.path.basename(lib_path)[:-3]
|
|
import_to_ip(f"{rnd}.{lib}")
|
|
shutil.rmtree(rnd_path)
|
|
|
|
|
|
def compile_module(cell, args):
|
|
name = args["name"]
|
|
ensure_tmp_folder(BUILD_DIR)
|
|
|
|
create_cpp_file(BUILD_DIR, cell, args)
|
|
|
|
ip = get_ipython()
|
|
# We compile the C++ code into an object file
|
|
command = get_gcc_cmd(name)
|
|
compile = ip.getoutput(command)
|
|
# currently warnings will lead to an abort
|
|
if len(compile) != 0:
|
|
raise Exception("\n".join(compile))
|
|
|
|
# link into a shared object, without links
|
|
linker = get_linker_cmd(name, "")
|
|
link = ip.getoutput(linker)
|
|
if len(link) != 0:
|
|
raise Exception("\n".join(link))
|
|
|
|
|
|
def link_module(name, others=None):
|
|
other_name = os.path.basename(name)[:-3]
|
|
linker = get_linker_cmd(other_name, others)
|
|
link = get_ipython().getoutput(linker)
|
|
if len(link) != 0:
|
|
raise Exception("\n".join(link))
|
|
|
|
|
|
version = 0
|
|
|
|
def create_cpp_file(path, cell, args):
|
|
name = args["name"]
|
|
# add pybind11 function defs
|
|
split = cell.split("defs")
|
|
if len(split) == 1 and len(args["functions"]) == 0:
|
|
raise Exception("You have to name a function as argument ('%%cpp -f <functionname>') or manual set defs")
|
|
# make sure that split is at least length 2
|
|
split.append("")
|
|
curr_defs = split[1]
|
|
for function in args["functions"]:
|
|
curr_defs += def_template.format(function=function)
|
|
with open(os.path.join(path,f"{name}.cpp"), 'w') as f:
|
|
f.write(template.format(cell_code=split[0],name=name,defs="{"+curr_defs+"}"))
|
|
|
|
|
|
def get_newest(name):
|
|
return sorted(glob.glob(f"{LIB_DIR}/lib{name}*.so"), key=os.path.getmtime)[-1]
|
|
|
|
history = []
|
|
|
|
|
|
# TODO: TEST
|
|
@register_cell_magic
|
|
def cpp(line, cell, *what):
|
|
global version
|
|
args = parse_args(line)
|
|
ip = get_ipython()
|
|
# get name of module
|
|
if args["cmake_path"] is None:
|
|
compile_id = get_gcc_cmd("dummy")
|
|
else:
|
|
# TODO!
|
|
compile_id = get_CMakeLists(args)
|
|
|
|
hash_object = hashlib.sha256((compile_id+"|"+line+"|"+cell).encode())
|
|
name = "cpp_magic_"+hash_object.hexdigest()[:10]
|
|
args["name"] = name
|
|
|
|
version += 1
|
|
history.append(os.path.join(MODULES_DIR, f"{name}.so"))
|
|
|
|
ensure_tmp_folder(TMP_PATH)
|
|
# Check if neccessary here
|
|
if not TMP_PATH in sys.path:
|
|
sys.path.append(TMP_PATH)
|
|
ensure_tmp_folder(MODULES_DIR)
|
|
ensure_tmp_folder(LIB_DIR)
|
|
|
|
# if not yet compiled, do so
|
|
if not os.path.exists(os.path.join(MODULES_DIR,f"{name}.so")):
|
|
# TODO: extend to cmake building
|
|
compile_module(cell, args)
|
|
|
|
update_dependencies(os.path.join(MODULES_DIR,f"{name}.so"))
|
|
relinks = list(get_list_of_rebuilds(os.path.join(MODULES_DIR,f"{name}.so")))
|
|
relinks = sorted(relinks, key=os.path.getmtime)[::-1]
|
|
for re in relinks:
|
|
other_name = os.path.basename(re)[:-3]
|
|
shutil.copyfile(os.path.join(MODULES_DIR,f"{other_name}.so"), os.path.join(LIB_DIR,f"lib{other_name}{version}.so"))
|
|
# touch
|
|
open(os.path.join(LIB_DIR,f"lib{other_name}{version}.so"), 'a').close()
|
|
# TODO: delete old
|
|
|
|
|
|
# modules = sorted(glob.glob(f"{MODULES_DIR}/*.so"), key=os.path.getmtime)[::-1]
|
|
# others = " ".join([f"-l{os.path.basename(get_newest(os.path.basename(n)[:-3]))[3:-3]}" for n in modules])
|
|
# print(others)
|
|
for re in relinks:
|
|
others = " ".join([f"-l{os.path.basename(get_newest(os.path.basename(n)[:-3]))[3:-3]}" for n in sorted(get_list_of_links(re), key=os.path.getmtime)[::-1]])
|
|
link_module(re, others)
|
|
other_name = os.path.basename(re)[:-3]
|
|
shutil.copyfile(os.path.join(MODULES_DIR,f"{other_name}.so"), os.path.join(LIB_DIR,f"lib{other_name}{version}.so"))
|
|
|
|
bulk_import([h for h in history if h in relinks])
|
|
|