openmp test

test_openmp.py

"""Check  whether the default compiler supports OpenMP.

    This routine is adapted from yt, thanks to Nathan
    Goldbaum. See https://github.com/pynbody/pynbody/issues/124"""

import tempfile 
import os
import sysconfig
import subprocess
import shutil 

# Create a temporary directory
tmpdir = tempfile.mkdtemp()
curdir = os.getcwd()
os.chdir(tmpdir)

# Get compiler invocation
compiler = os.environ.get('CC',sysconfig.get_config_var('CC'))
print 'Using compiler: ', compiler
# Attempt to compile a test script.
# See http://openmp.org/wp/openmp-compilers/
filename = r'test.c'
file = open(filename,'w', 0)
file.write(
    "#include <omp.h>\n"
    "#include <stdio.h>\n"
    "int main() {\n"
    "#pragma omp parallel\n"
    "printf(\"Hello from thread %d, nthreads %d\\n\", omp_get_thread_num(), omp_get_num_threads());\n"
    "}"
    )
try:
    with open(os.devnull, 'w') as fnull:
        exit_code = subprocess.call([compiler, '-fopenmp', filename],
                                    stdout=fnull, stderr=fnull)
except OSError :
    exit_code = 1
    
# Clean up
file.close()
os.chdir(curdir)
shutil.rmtree(tmpdir)

if exit_code == 0:
    print 'OpenMP works!'
else:
    print 'OpenMP error'