autotune.py

""" Autotuning for sum factorization kernels """

import os
import re
import subprocess
import filelock
import hashlib
import logging
import json
from operator import mul
from six.moves import reduce
import time

import loopy as lp
from pytools import product
from cgen import ArrayOf, AlignedAttribute, Initializer

from dune.codegen.generation import cache_restoring, delete_cache_items
from dune.codegen.loopy.target import DuneTarget, type_floatingpoint
from dune.codegen.sumfact.realization import realize_sumfact_kernel_function
from dune.codegen.options import get_option, set_option
from dune.codegen.error import CodegenAutotuneError


def get_cmake_cache_entry(entry):
    for line in open(os.path.join(get_option("project_basedir"), "CMakeCache.txt"), "r"):
        match = re.match("{}:[INTERNAL|FILEPATH|BOOL|STRING|PATH|UNINITIALIZED|STATIC]+=(.*)".format(entry), line)
        if match:
            return match.groups()[0]


def get_dune_codegen_dir():
    if get_cmake_cache_entry("CMAKE_PROJECT_NAME") == "dune-codegen":
        return get_option("project_basedir")
    else:
        return get_cmake_cache_entry("dune-codegen_DIR")


def compiler_invocation(name, filename):
    # Determine the CMake Generator in use
    gen = get_cmake_cache_entry("CMAKE_GENERATOR")
    assert(gen == "Unix Makefiles")

    # Find compiler path
    compiler = get_cmake_cache_entry("CMAKE_CXX_COMPILER")
    compile_flags = [compiler]

    # Parse compiler flags
    for line in open(os.path.join(get_dune_codegen_dir(), "python", "CMakeFiles", "_autotune_target.dir", "flags.make"), "r"):
        match = re.match("([^=]*)=(.*)", line)
        if match:
            compile_flags.extend(match.groups()[1].split())

    # Add the source file
    compile_flags.append(filename)

    # Parse linker flags
    for line in open(os.path.join(get_dune_codegen_dir(), "python", "CMakeFiles", "_autotune_target.dir", "link.txt"), "r"):
        match = re.match(".*_autotune_target (.*)", line)
        if match:
            for flag in match.groups()[0].split():
                if flag.startswith("-") or os.path.isabs(flag):
                    compile_flags.append(flag)
                else:
                    compile_flags.append(os.path.join(get_dune_codegen_dir(), "python", flag))

    # Set an output name
    compile_flags.append("-o")
    compile_flags.append(name)

    return compile_flags


def write_global_data(sf, filename):
    opcounting = get_option("opcounter")
    with open(filename, "a") as f:
        # Get kernel
        from dune.codegen.pdelab.localoperator import extract_kernel_from_cache
        knl = realize_sumfact_kernel_function(sf)
        constructor_knl = extract_kernel_from_cache("operator", "constructor_kernel", None, wrap_in_cgen=False, add_timings=False)
        constructor_knl = constructor_knl.copy(target=DuneTarget(declare_temporaries=False))
        constructor_knl = lp.get_one_scheduled_kernel(constructor_knl)

        target = DuneTarget()
        from loopy.codegen import CodeGenerationState
        codegen_state = CodeGenerationState(kernel=constructor_knl,
                                            implemented_data_info=None,
                                            implemented_domain=None,
                                            implemented_predicates=frozenset(),
                                            seen_dtypes=frozenset(),
                                            seen_functions=frozenset(),
                                            seen_atomic_dtypes=frozenset(),
                                            var_subst_map={},
                                            allow_complex=False,
                                            is_generating_device_code=True,
                                            )

        for decl in target.get_device_ast_builder().get_temporary_decls(codegen_state, 0):
            f.write("{}\n".format(next(iter(decl.generate()))))


def write_setup_code(sf, filename, define_thetas=True):
    with open(filename, "a") as f:
        # Setup a polynomial object (normally done in the LocalOperator members)
        from dune.codegen.loopy.target import type_floatingpoint
        real = type_floatingpoint()
        f.write("  using RF = {};\n".format(real))
        f.write("  using DF = {};\n".format(real))

        from dune.codegen.sumfact.tabulation import name_polynomials
        degs = tuple(m.basis_size - 1 for m in sf.matrix_sequence_quadrature_permuted)
        for deg in set(degs):
            f.write("  Dune::QkStuff::EquidistantLagrangePolynomials<DF, RF, {}> {};\n".format(deg, name_polynomials(deg)))

        # Get kernel
        from dune.codegen.pdelab.localoperator import extract_kernel_from_cache
        knl = realize_sumfact_kernel_function(sf)
        constructor_knl = extract_kernel_from_cache("operator", "constructor_kernel", None, wrap_in_cgen=False, add_timings=False)
        constructor_knl = constructor_knl.copy(target=DuneTarget(declare_temporaries=False))
        constructor_knl = lp.get_one_scheduled_kernel(constructor_knl)

        # Allocate buffers
        alignment = get_option("max_vector_width") // 8
        size = max(product(m.quadrature_size for m in sf.matrix_sequence_quadrature_permuted) * sf.vector_width,
                   product(m.basis_size for m in sf.matrix_sequence_quadrature_permuted) * sf.vector_width)
        size = int(size * (get_option("precision_bits") / 8))
        f.writelines(["  char buffer0[{}] __attribute__ ((aligned ({})));\n".format(size, alignment),
                      "  char buffer1[{}] __attribute__ ((aligned ({})));\n".format(size, alignment),
                      ])

        # Setup fastdg inputs
        for arg in sf.interface.signature_args:
            if "jacobian" in arg:
                f.write("{} = 0;\n".format(arg))
            else:
                size = sf.interface.fastdg_interface_object_size
                f.write("  RF {}[{}] __attribute__ ((aligned ({})));\n".format(arg.split()[-1], size, alignment))

        # Write stuff into the input buffer
        f.writelines(["  {0} *input = ({0} *)buffer0;\n".format(real),
                      "  {0} *output = ({0} *)buffer{1};\n".format(real, sf.length % 2),
                      "  for(int i=0; i<{}; ++i)\n".format(size / (get_option("precision_bits") / 8)),
                      "    input[i] = ({})(i+1);\n".format(real),
                      ])

        target = DuneTarget()
        from loopy.codegen import CodeGenerationState
        codegen_state = CodeGenerationState(kernel=constructor_knl,
                                            implemented_data_info=None,
                                            implemented_domain=None,
                                            implemented_predicates=frozenset(),
                                            seen_dtypes=frozenset(),
                                            seen_functions=frozenset(),
                                            seen_atomic_dtypes=frozenset(),
                                            var_subst_map={},
                                            allow_complex=False,
                                            is_generating_device_code=True,
                                            )

        if define_thetas:
            for decl in target.get_device_ast_builder().get_temporary_decls(codegen_state, 0):
                f.write("  {}\n".format(next(iter(decl.generate()))))

        for _, line in constructor_knl.preambles:
            if "gfsu" not in line:
                f.write("  {}\n".format(line))

        # Add setup code for theta matrices. We add some lines not necessary,
        # but it would be more work to remove them than keeping them.
        for line in lp.generate_body(constructor_knl).split("\n")[1:-1]:
            if "gfsu" not in line and "meshwidth" not in line and "geometry" not in line:
                f.write("  {}\n".format(line))

        # INtroduces a variable that makes sure that the kernel cannot be optimized away
        f.writelines(["  {} accum;\n".format(real),
                      "  std::mt19937 rng;\n",
                      "  rng.seed(42);\n",
                      "  std::uniform_int_distribution<> dis(0, {});\n".format(size / (get_option("precision_bits") / 8)),
                      ])


def generate_standalone_code_google_benchmark(sf, filename):
    delete_cache_items("kernel_default")

    # Turn off opcounting
    opcounting = get_option("opcounter")
    set_option("opcounter", False)

    # Extract sum factorization kernel
    from dune.codegen.pdelab.localoperator import extract_kernel_from_cache
    knl = realize_sumfact_kernel_function(sf)

    # Add the implementation of the kernel.
    # TODO: This can probably done in a safer way?
    first_line = knl.member.lines[0]
    arguments = first_line[first_line.find("(") + 1:first_line.find(")")]

    with open(filename, "w") as f:
        f.writelines(["// {}".format(first_line),
                      "\n",
                      "#include \"config.h\"\n",
                      "#include \"benchmark/benchmark.h\"\n",
                      "#include<dune/pdelab/finiteelementmap/qkdg.hh>\n",
                      "#include<dune/codegen/common/vectorclass.hh>\n",
                      "#include<dune/codegen/sumfact/onedquadrature.hh>\n",
                      "#include<dune/codegen/sumfact/horizontaladd.hh>\n",
                      "#include<random>\n",
                      "#include<fstream>\n",
                      "#include<iostream>\n",
                      "\n"
                      ])

    write_global_data(sf, filename)

    with open(filename, "a") as f:
        arguments = ', '.join(sf.interface.signature_args)
        if len(arguments) > 0:
            arguments = ', ' + arguments
        arguments = 'const char* buffer0, const char* buffer1' + arguments
        f.write("void sumfact_kernel({})\n".format(arguments))
        for line in knl.member.lines[1:]:
            f.write("{}\n".format(line))

        f.write("\n\n")
        f.write("static void BM_sumfact_kernel(benchmark::State& state){\n")

    write_setup_code(sf, filename, define_thetas=False)

    additional_arguments = [i.split()[-1] for i in sf.interface.signature_args]
    additional_arguments = ', '.join(additional_arguments)
    if len(additional_arguments) > 0:
        additional_arguments = ', ' + additional_arguments
    with open(filename, "a") as f:
        f.writelines(["  for (auto _ : state){\n",
                      "    sumfact_kernel(buffer0, buffer1{});\n".format(additional_arguments),
                      "  }\n",
                      "}\n",
                      "BENCHMARK(BM_sumfact_kernel);\n",
                      "\n",
                      "BENCHMARK_MAIN();"
                      ])

    # Maybe turn opcounting on again
    set_option("opcounter", opcounting)


def generate_standalone_code(sf, filename):
    delete_cache_items("kernel_default")

    # Turn off opcounting
    opcounting = get_option("opcounter")
    set_option("opcounter", False)

    # Extract sum factorization kernel
    from dune.codegen.pdelab.localoperator import extract_kernel_from_cache
    knl = realize_sumfact_kernel_function(sf)
    first_line = knl.member.lines[0]

    with open(filename, "w") as f:
        f.writelines(["// {}".format(first_line),
                      "\n",
                      "#include \"config.h\"\n",
                      "#include<dune/pdelab/finiteelementmap/qkdg.hh>\n",
                      "#include<dune/codegen/common/tsc.hh>\n",
                      "#include<dune/codegen/common/vectorclass.hh>\n",
                      "#include<dune/codegen/sumfact/onedquadrature.hh>\n",
                      "#include<dune/codegen/sumfact/horizontaladd.hh>\n",
                      "#include<random>\n",
                      "#include<fstream>\n",
                      "#include<iostream>\n",
                      "\n"
                      ])

        f.writelines(["int main(int argc, char** argv)\n",
                      "{\n",
                      ])

    write_setup_code(sf, filename)

    # Write measurement
    with open(filename, "a") as f:
        # Start a TSC timer
        f.writelines(["  auto start = Dune::PDELab::TSC::start();\n",
                      ])

        # Add the implementation of the kernel.
        repeats = int(1e9 / sf.operations)
        f.write("  for(int i=0; i<{}; ++i)\n".format(repeats))
        f.write("  {\n")
        for line in knl.member.lines[1:]:
            f.write("    {}\n".format(line))
        f.write("  }\n")

        # Stop the TSC timer and write the result to a file
        f.writelines(["  auto stop = Dune::PDELab::TSC::stop();\n",
                      "  std::ofstream file;\n",
                      "  file.open(argv[1]);\n",
                      "  file << Dune::PDELab::TSC::elapsed(start, stop) / {} << std::endl;\n".format(str(float(repeats))),
                      "  file.close();\n",
                      "  accum += output[dis(rng)];\n",
                      "  std::cout << accum;\n",
                      "}\n",
                      ])

    # Maybe turn opcounting on again
    set_option("opcounter", opcounting)


def generate_standalone_kernel_code(kernel, signature, filename, transformations=None):
    # Turn off opcounting
    opcounting = get_option("opcounter")
    set_option("opcounter", False)

    # Remove opcounter from signature
    p = re.compile('OpCounter::OpCounter<([^>]*)>')
    assert len(signature) == 1
    sig = signature[0]
    sig = p.sub(r'\1', sig)
    assert 'OpCounter' not in signature

    with open(filename, 'w') as f:
        if transformations:
            f.write('// Transformations:\n')
            for trafo in transformations:
                f.write('// {}\n'.format(trafo))
            f.write('\n')

        # Write headers
        headers = ['#include "config.h"',
                   '#include <iostream>',
                   '#include <fstream>',
                   '#include <random>',
                   '#include "benchmark/benchmark.h"',
                   '#include <dune/codegen/common/vectorclass.hh>',
                   '#include <dune/codegen/sumfact/horizontaladd.hh>',
                   ]
        f.write("\n".join(headers))

        # Get a list of the function argument names
        arguments = sig[sig.find('(') + 1:sig.find(')')].split(',')
        arguments = [a.split(' ')[-1] for a in arguments]

        global_args = [a for a in kernel.args if a.name not in arguments]

        # Declare global arguments
        f.write('\n\n')
        target = DuneTarget()
        for g in global_args:
            decl_info = g.decl_info(target, True, g.dtype)
            for idi in decl_info:
                ast_builder = target.get_device_ast_builder()
                arg_decl = lp.target.c.POD(ast_builder, idi.dtype, idi.name)
                arg_decl = ArrayOf(arg_decl, reduce(mul, g.shape))
                arg_decl = AlignedAttribute(g.dtype.itemsize * g.vector_size(target), arg_decl)
                f.write('{}\n'.format(arg_decl))

        # Generate function we want to benchmark
        f.write('\n')
        f.write(sig[0:sig.find(')') + 1])
        f.writelines(lp.generate_body(kernel))
        f.write('\n\n')

        # Generate function that will do the benchmarking
        f.write('static void BM_sumfact_kernel(benchmark::State& state){\n')

        # Declare random generators
        real = type_floatingpoint()
        lines = ['  std::uniform_real_distribution<{}> unif(0,1);'.format(real),
                 '  std::uniform_int_distribution<int> unif_int(0,1);',
                 '  std::default_random_engine re;']
        f.write('\n'.join(lines) + '\n')

        # Declare function arguments
        function_arguments = [a for a in kernel.args if a.name in arguments]
        for arg in function_arguments:
            if 'buffer' in arg.name:
                byte_size = reduce(mul, arg.shape) * 8
                f.write('  char {}[{}] __attribute__ ((aligned ({})));\n'.format(arg.name,
                                                                                 byte_size,
                                                                                 arg.alignment),)
            elif isinstance(arg, lp.ValueArg):
                assert 'jacobian_offset' in arg.name
                decl = arg.get_arg_decl(ast_builder)
                decl = Initializer(decl, 'unif_int(re)')
                f.write('  {}\n'.format(decl))
            else:
                assert 'fastdg' in arg.name
                size = reduce(mul, arg.shape)
                min_stride = min([tag.stride for tag in arg.dim_tags])
                size *= min_stride
                alignment = arg.dtype.itemsize
                f.write('  {} {}[{}] __attribute__ ((aligned ({})));\n'.format(real,
                                                                               arg.name,
                                                                               size,
                                                                               alignment))

        # Initialize arguments
        def _initialize_arg(arg):
            if isinstance(arg, lp.ValueArg):
                return []
            real = type_floatingpoint()
            size = reduce(mul, arg.shape)
            fill_name = arg.name + '_fill'
            lines = ['  {}* {} = (double *) {};'.format(real, fill_name, arg.name),
                     '  for (std::size_t i=0; i<{}; ++i){{'.format(size),
                     '    {}[i] = unif(re);'.format(fill_name),
                     '  }']
            return lines

        for arg in kernel.args:
            lines = _initialize_arg(arg)
            f.write('\n'.join(lines) + '\n')

        # Benchmark loop
        function_call = kernel.name + '({})'.format(','.join(arguments))
        f.writelines(['  for (auto _ : state){\n',
                      '    {};\n'.format(function_call),
                      '  }\n',
                      ])
        f.write('}\n')

        # Benchmark main
        main = ['',
                'BENCHMARK(BM_sumfact_kernel);',
                '',
                'BENCHMARK_MAIN();']
        f.write('\n'.join(main))

    # Maybe turn opcounting on again
    set_option("opcounter", opcounting)


def autotune_realization(sf=None, kernel=None, signature=None, transformations=None):
    """Generate an microbenchmark, compile run and return time

    Parameters
    ----------
    sf: SumfactKernel or VectorizedSumfactKernel
    kernel: loopy.kernel.LoopKernel
    signature: str
    transformation: list of str
        Will be used to distinguish between autotune targets
    """
    if sf is None:
        assert kernel is not None
        assert signature is not None
    else:
        assert kernel is None
        assert signature is None

    logger = logging.getLogger(__name__)

    # Make sure that the benchmark directory exists
    dir = os.path.join(get_option("project_basedir"), "autotune-benchmarks")
    if not os.path.exists(dir):
        os.mkdir(dir)

    if sf is None:
        basename = "autotune_sumfact_{}".format(kernel.name)
    else:
        basename = "autotune_sumfact_{}".format(sf.function_name)
    if transformations:
        for trafo in transformations:
            basename = '{}_{}'.format(basename, trafo)
    basename = hashlib.sha256(basename.encode()).hexdigest()

    filename = os.path.join(dir, "{}.cc".format(basename))
    logname = os.path.join(dir, "{}.log".format(basename))
    lock = os.path.join(dir, "{}.lock".format(basename))
    executable = os.path.join(dir, basename)

    # Generate and compile a benchmark program
    #
    # Note: cache restoring is only necessary when generating from SumfactKernel
    with cache_restoring():
        with filelock.FileLock(lock):
            if not os.path.isfile(logname):
                logger.debug('Generate autotune target in file {}'.format(filename))

                if sf is None:
                    generate_standalone_kernel_code(kernel, signature, filename, transformations)
                elif get_option("autotune_google_benchmark"):
                    generate_standalone_code_google_benchmark(sf, filename)
                else:
                    generate_standalone_code(sf, filename)

                call = []
                wrapper = get_cmake_cache_entry("DUNE_CODEGEN_BENCHMARK_COMPILATION_WRAPPER")
                if wrapper:
                    call.append(wrapper)

                call.extend(compiler_invocation(executable, filename))
                devnull = open(os.devnull, 'w')
                ret = subprocess.call(call, stdout=devnull, stderr=subprocess.STDOUT)
                if ret != 0:
                    raise CodegenAutotuneError("Compilation of autotune executable failed. Invocation: {}".format(" ".join(call)))

                # File system synchronization!
                while not os.path.exists(executable):
                    time.sleep(0.01)

                # Check whether the user specified an execution wrapper
                call = []
                wrapper = get_cmake_cache_entry("DUNE_CODEGEN_BENCHMARK_EXECUTION_WRAPPER")
                if wrapper:
                    call.append(wrapper)

                # Run the benchmark program
                call.append(executable)
                if get_option("autotune_google_benchmark"):
                    call.append("--benchmark_out={}".format(logname))
                    call.append("--benchmark_repetitions=10")
                    # call.append("--benchmark_out_format=csv")
                else:
                    call.append(logname)
                ret = subprocess.call(call, stdout=devnull, stderr=subprocess.STDOUT)
                if ret != 0:
                    raise CodegenAutotuneError("Execution of autotune benchmark failed. Invocation: {}".format(" ".join(call)))

                # File system synchronization!
                while not os.path.exists(logname):
                    time.sleep(0.01)

            # Extract the result form the log file
            if get_option("autotune_google_benchmark"):
                import json
                with open(logname) as json_file:
                    try:
                        data = json.load(json_file)
                        minimal_time = 1e80
                        for b in data['benchmarks']:
                            if b['name'] == 'BM_sumfact_kernel':
                                if b['cpu_time'] < minimal_time:
                                    minimal_time = b['cpu_time']
                        assert minimal_time < 1e80
                        return minimal_time
                    except Exception as e:
                        print("Error while loading file {}".format(logname))
                        raise e
            else:
                return float(next(iter(open(logname, "r")))) / 1000000