pystencils: enhanced windows support

cpu/cpujit.py

@@ -124,8 +124,9 @@ def readConfig():
     elif platform.system().lower() == 'windows':
         defaultCompilerConfig = OrderedDict([
             ('os', 'windows'),
+            ('msvcVersion', 'latest'),
             ('arch', 'x64'),
-            ('flags', '/Ox /fp:fast /openmp'),
+            ('flags', '/Ox /fp:fast /openmp /arch:avx'),
             ('restrictQualifier', '__restrict')
         ])
         defaultCacheConfig = OrderedDict([
@@ -143,28 +144,26 @@ def readConfig():
     if configExists:
         loadedConfig = json.load(open(configPath, 'r'))
         config = _recursiveDictUpdate(config, loadedConfig)
+    else:
+        createFolder(configPath, True)
+        json.dump(config, open(configPath, 'w'), indent=4)
 
     config['cache']['sharedLibrary'] = os.path.expanduser(config['cache']['sharedLibrary'])
     config['cache']['objectCache'] = os.path.expanduser(config['cache']['objectCache'])
 
-    # create folders if they don't exist yet
-    createFolder(configPath, True)
-
     if config['cache']['clearCacheOnStart']:
         shutil.rmtree(config['cache']['objectCache'], ignore_errors=True)
 
     createFolder(config['cache']['objectCache'], False)
     createFolder(config['cache']['sharedLibrary'], True)
 
-    json.dump(config, open(configPath, 'w'), indent=4)
-
     if 'env' not in config['compiler']:
         config['compiler']['env'] = {}
 
     if config['compiler']['os'] == 'windows':
-        from setuptools.msvc import msvc14_get_vc_env
-        msvcEnv = msvc14_get_vc_env(config['compiler']['arch'])
-        config['compiler']['env'].update({k.upper(): v for k, v in msvcEnv.items()})
+        from pystencils.cpu.msvc_detection import getEnvironment
+        msvcEnv = getEnvironment(config['compiler']['msvcVersion'], config['compiler']['arch'])
+        config['compiler']['env'].update(msvcEnv)
 
     return config
 
@@ -325,10 +324,12 @@ def compileAndLoad(ast):
     else:
         if getCompilerConfig()['os'].lower() == 'windows':
             libFile = os.path.join(cacheConfig['objectCache'], codeHashStr + ".dll")
-            compileWindows(ast, codeHashStr, srcFile, libFile)
+            if not os.path.exists(libFile):
+                compileWindows(ast, codeHashStr, srcFile, libFile)
         else:
             libFile = os.path.join(cacheConfig['objectCache'], codeHashStr + ".so")
-            compileLinux(ast, codeHashStr, srcFile, libFile)
+            if not os.path.exists(libFile):
+                compileLinux(ast, codeHashStr, srcFile, libFile)
         return cdll.LoadLibrary(libFile)[ast.functionName]
 
 

cpu/msvc_detection.py

+import subprocess
+import os
+
+
+def getEnvironment(versionSpecifier, arch='x64'):
+    """
+    Returns an environment dictionary, for activating the Visual Studio compiler
+    :param versionSpecifier: either a version number, year number, 'auto' or 'latest' for automatic detection of latest
+                             installed version or 'setuptools' for setuptools-based detection
+    :param arch: x86 or x64
+    """
+    if versionSpecifier == 'setuptools':
+        return getEnvironmentFromSetupTools(arch)
+    else:
+        if versionSpecifier in ('auto', 'latest'):
+            versionNr = findLatestMsvcVersionUsingEnvironmentVariables()
+        else:
+            versionNr = normalizeMsvcVersion(versionSpecifier)
+        vcVarsPath = getVcVarsPath(versionNr)
+        return getEnvironmentFromVcVarsFile(vcVarsPath, arch)
+
+
+def findLatestMsvcVersionUsingEnvironmentVariables():
+    import re
+    regex = re.compile('VS(\d\d)\dCOMNTOOLS')
+    versions = []
+    for key, value in os.environ.items():
+        match = regex.match(key)
+        if match:
+            versions.append(int(match.group(1)))
+    if len(versions) == 0:
+        raise ValueError("Visual Studio not found.")
+    versions.sort()
+    return versions[-1]
+
+
+def normalizeMsvcVersion(version):
+    """
+    Takes version specifiers in the following form:
+        - year: 2012, 2013, 2015, either as int or string
+        - version numbers with or without dot i.e. 11.0 or 11
+    :return: integer version number
+    """
+    if isinstance(version, str) and '.' in version:
+        version = version.split('.')[0]
+
+    version = int(version)
+    mapping = {
+        2015: 14,
+        2013: 12,
+        2012: 11
+    }
+    if version in mapping:
+        return mapping[version]
+    else:
+        return version
+
+
+def getEnvironmentFromVcVarsFile(vcVarsFile, arch):
+    out = subprocess.check_output(
+        'cmd /u /c "{}" {} && set'.format(vcVarsFile, arch),
+        stderr=subprocess.STDOUT,
+    ).decode('utf-16le', errors='replace')
+
+    env = {key.upper(): value for key, _, value in (line.partition('=') for line in out.splitlines()) if key and value}
+    return env
+
+
+def getVcVarsPath(versionNr):
+    environmentVarName = 'VS%d0COMNTOOLS' % (versionNr,)
+    vcPath = os.environ[environmentVarName]
+    path = os.path.join(vcPath, '..', '..', 'VC', 'vcvarsall.bat')
+    return os.path.abspath(path)
+
+
+def getEnvironmentFromSetupTools(arch):
+    from setuptools.msvc import msvc14_get_vc_env
+    msvcEnv = msvc14_get_vc_env(arch)
+    return {k.upper(): v for k, v in msvcEnv.items()}
+
+