Skip to content
Snippets Groups Projects
Commit a85e0b4d authored by Markus Holzer's avatar Markus Holzer
Browse files

Clean cpujit a bit

parent 27cbf05c
Branches
Tags
No related merge requests found
Pipeline #53893 passed
...@@ -46,13 +46,14 @@ Then 'cl.exe' is used to compile. ...@@ -46,13 +46,14 @@ Then 'cl.exe' is used to compile.
from appdirs import user_cache_dir, user_config_dir from appdirs import user_cache_dir, user_config_dir
from collections import OrderedDict from collections import OrderedDict
import hashlib import hashlib
import importlib.util
import json import json
import os import os
import platform import platform
import shutil import shutil
import subprocess import subprocess
from sysconfig import get_paths import sysconfig
from tempfile import TemporaryDirectory, NamedTemporaryFile import tempfile
import textwrap import textwrap
import time import time
import warnings import warnings
...@@ -62,9 +63,10 @@ import numpy as np ...@@ -62,9 +63,10 @@ import numpy as np
from pystencils import FieldType from pystencils import FieldType
from pystencils.astnodes import LoopOverCoordinate from pystencils.astnodes import LoopOverCoordinate
from pystencils.backends.cbackend import generate_c, get_headers, CFunction from pystencils.backends.cbackend import generate_c, get_headers, CFunction
from pystencils.typing import BasicType, CastFunc, VectorType, VectorMemoryAccess from pystencils.cpu.msvc_detection import get_environment
from pystencils.include import get_pystencils_include_path from pystencils.include import get_pystencils_include_path
from pystencils.kernel_wrapper import KernelWrapper from pystencils.kernel_wrapper import KernelWrapper
from pystencils.typing import BasicType, CastFunc, VectorType, VectorMemoryAccess
from pystencils.utils import atomic_file_write, recursive_dict_update from pystencils.utils import atomic_file_write, recursive_dict_update
...@@ -218,12 +220,11 @@ def read_config(): ...@@ -218,12 +220,11 @@ def read_config():
shutil.rmtree(config['cache']['object_cache'], ignore_errors=True) shutil.rmtree(config['cache']['object_cache'], ignore_errors=True)
create_folder(config['cache']['object_cache'], False) create_folder(config['cache']['object_cache'], False)
with NamedTemporaryFile('w', dir=os.path.dirname(cache_status_file), delete=False) as f: with tempfile.NamedTemporaryFile('w', dir=os.path.dirname(cache_status_file), delete=False) as f:
json.dump(config['compiler'], f, indent=4) json.dump(config['compiler'], f, indent=4)
os.replace(f.name, cache_status_file) os.replace(f.name, cache_status_file)
if config['compiler']['os'] == 'windows': if config['compiler']['os'] == 'windows':
from pystencils.cpu.msvc_detection import get_environment
msvc_env = get_environment(config['compiler']['msvc_version'], config['compiler']['arch']) msvc_env = get_environment(config['compiler']['msvc_version'], config['compiler']['arch'])
if 'env' not in config['compiler']: if 'env' not in config['compiler']:
config['compiler']['env'] = {} config['compiler']['env'] = {}
...@@ -472,16 +473,15 @@ def create_module_boilerplate_code(module_name, names): ...@@ -472,16 +473,15 @@ def create_module_boilerplate_code(module_name, names):
def load_kernel_from_file(module_name, function_name, path): def load_kernel_from_file(module_name, function_name, path):
from importlib.util import spec_from_file_location, module_from_spec
try: try:
spec = spec_from_file_location(name=module_name, location=path) spec = importlib.util.spec_from_file_location(name=module_name, location=path)
mod = module_from_spec(spec) mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod) spec.loader.exec_module(mod)
except ImportError: except ImportError:
warnings.warn(f"Could not load {path}, trying on more time in 5 seconds ...") warnings.warn(f"Could not load {path}, trying on more time in 5 seconds ...")
time.sleep(5) time.sleep(5)
spec = spec_from_file_location(name=module_name, location=path) spec = importlib.util.spec_from_file_location(name=module_name, location=path)
mod = module_from_spec(spec) mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod) spec.loader.exec_module(mod)
return getattr(mod, function_name) return getattr(mod, function_name)
...@@ -563,7 +563,7 @@ def compile_module(code, code_hash, base_dir, compile_flags=None): ...@@ -563,7 +563,7 @@ def compile_module(code, code_hash, base_dir, compile_flags=None):
compile_flags = [] compile_flags = []
compiler_config = get_compiler_config() compiler_config = get_compiler_config()
extra_flags = ['-I' + get_paths()['include'], '-I' + get_pystencils_include_path()] + compile_flags extra_flags = ['-I' + sysconfig.get_paths()['include'], '-I' + get_pystencils_include_path()] + compile_flags
if compiler_config['os'].lower() == 'windows': if compiler_config['os'].lower() == 'windows':
lib_suffix = '.pyd' lib_suffix = '.pyd'
...@@ -597,7 +597,6 @@ def compile_module(code, code_hash, base_dir, compile_flags=None): ...@@ -597,7 +597,6 @@ def compile_module(code, code_hash, base_dir, compile_flags=None):
# Linking # Linking
if windows: if windows:
import sysconfig
config_vars = sysconfig.get_config_vars() config_vars = sysconfig.get_config_vars()
py_lib = os.path.join(config_vars["installed_base"], "libs", py_lib = os.path.join(config_vars["installed_base"], "libs",
f"python{config_vars['py_version_nodot']}.lib") f"python{config_vars['py_version_nodot']}.lib")
...@@ -631,7 +630,7 @@ def compile_and_load(ast, custom_backend=None): ...@@ -631,7 +630,7 @@ def compile_and_load(ast, custom_backend=None):
compile_flags = ast.instruction_set['compile_flags'] compile_flags = ast.instruction_set['compile_flags']
if cache_config['object_cache'] is False: if cache_config['object_cache'] is False:
with TemporaryDirectory() as base_dir: with tempfile.TemporaryDirectory() as base_dir:
lib_file = compile_module(code, code_hash_str, base_dir, compile_flags=compile_flags) lib_file = compile_module(code, code_hash_str, base_dir, compile_flags=compile_flags)
result = load_kernel_from_file(code_hash_str, ast.function_name, lib_file) result = load_kernel_from_file(code_hash_str, ast.function_name, lib_file)
else: else:
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment