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
import pkg_resources
from six.moves import reduce
import textwrap
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, option_context
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
    with option_context(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();"
                          ])


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

    # Turn off opcounting
    with option_context(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",
                          ])


def generate_standalone_kernel_code(kernel, signature, filename, transformations=None):
    # Turn off opcounting
    with option_context(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

        # Which transformations were applied
        codegen_transformations = ''
        if transformations:
            codegen_transformations = ''
            for trafo in transformations:
                codegen_transformations += '// {}\n'.format(trafo)

        template = 'kernel_benchmark_template1.cc.in'
        use_datasets = True

        # Old benchmark template
        # template = 'kernel_benchmark_template0.cc.in'
        # use_datasets = False

        template_filename = pkg_resources.resource_filename(__name__, template)
        with open(template_filename, 'r') as f:
            benchmark = f.read()

        # Find function arguments and global arguments
        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]
        buffer_arguments = [a for a in arguments if a.startswith('buff')]
        input_arguments = [a for a in arguments if a not in buffer_arguments]

        # Declare global arguments
        codegen_declare_global_arguments = ''
        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)
                codegen_declare_global_arguments += '{}\n'.format(arg_decl)
        codegen_declare_global_arguments = textwrap.indent(codegen_declare_global_arguments, '  ')

        # Helper function for argument initialization
        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

        # Initialize global arguments
        codegen_initialize_global_arguments = ''
        for arg in global_args:
            lines = _initialize_arg(arg)
            codegen_initialize_global_arguments += '\n'.join(lines) + '\n'
        codegen_initialize_global_arguments = textwrap.indent(codegen_initialize_global_arguments, '  ')

        codegen_initialize_input = ''

        # Function we want to benchmark
        codegen_benchmark_function = ''
        codegen_benchmark_function += sig[0:sig.find(')') + 1]
        codegen_benchmark_function += lp.generate_body(kernel)
        codegen_benchmark_function = textwrap.indent(codegen_benchmark_function, '  ')

        # Declare function arguments
        codegen_declare_arguments = []
        codegen_declare_input = []
        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
                codegen_declare_arguments.append('  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)')
                codegen_declare_arguments.append(('  {}\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
                real = type_floatingpoint()
                if use_datasets:
                    codegen_declare_input.append(('{} {}[datasets][{}] __attribute__ ((aligned ({})));\n'.format(real,
                                                                                                                 arg.name,
                                                                                                                 size,
                                                                                                                 alignment)))
                else:
                    codegen_declare_input.append(('{} {}[{}] __attribute__ ((aligned ({})));\n'.format(real,
                                                                                                       arg.name,
                                                                                                       size,
                                                                                                       alignment)))

        codegen_declare_arguments = ''.join(codegen_declare_arguments)
        codegen_declare_arguments = textwrap.indent(codegen_declare_arguments, '  ')
        codegen_declare_input = ''.join(codegen_declare_input)
        codegen_declare_input = textwrap.indent(codegen_declare_input, '  ')

        # Initialize function arguments
        codegen_initialize_arguments = ''
        codegen_initialize_input = ''
        for arg in function_arguments:
            if 'fastdg' in arg.name:
                if use_datasets:
                    lines = _initialize_arg(arg)
                    lines = ['  ' + a for a in lines]
                    lines = [a.replace(arg.name + ';', arg.name + '[i];') for a in lines]
                    lines.insert(0, 'for(std::size_t i=0; i<datasets; ++i){')
                    lines.append('}')
                    codegen_initialize_input += '\n'.join(lines) + '\n'
                else:
                    lines = _initialize_arg(arg)
                    codegen_initialize_arguments += '\n'.join(lines) + '\n'
            else:
                lines = _initialize_arg(arg)
                codegen_initialize_arguments += '\n'.join(lines) + '\n'
        codegen_initialize_arguments = textwrap.indent(codegen_initialize_arguments, '  ')
        codegen_initialize_input = textwrap.indent(codegen_initialize_input, '  ')

        # Call the benchmark function
        if use_datasets:
            arguments_with_datasets = arguments.copy()
            arguments_with_datasets = [a if 'fastdg' not in a else a + '[i]' for a in arguments]
            codegen_call_benchmark_function = 'for (std::size_t i=0; i<datasets; ++i){\n'
            codegen_call_benchmark_function += '  ' + kernel.name + '({})'.format(','.join(arguments_with_datasets)) + ';\n'
            for arg in input_arguments:
                codegen_call_benchmark_function += 'benchmark::DoNotOptimize({}[i][0]);\n'.format(arg)
            codegen_call_benchmark_function += '}'
        else:
            codegen_call_benchmark_function = kernel.name + '({})'.format(','.join(arguments)) + ';\n'
        codegen_call_benchmark_function = textwrap.indent(codegen_call_benchmark_function, '    ')

        # Replace placeholders in benchmark template
        benchmark = benchmark.replace('${CODEGEN_TRANSFORMATIONS}', codegen_transformations)
        benchmark = benchmark.replace('${CODEGEN_DECLARE_GLOBAL_ARGUMENTS}', codegen_declare_global_arguments)
        benchmark = benchmark.replace('${CODEGEN_DECLARE_INPUT}', codegen_declare_input)
        benchmark = benchmark.replace('${CODEGEN_INITIALIZE_GLOBAL_ARGUMENTS}', codegen_initialize_global_arguments)
        benchmark = benchmark.replace('${CODEGEN_INITIALIZE_INPUT}', codegen_initialize_input)
        benchmark = benchmark.replace('${CODEGEN_BENCHMARK_FUNCTION}', codegen_benchmark_function)
        benchmark = benchmark.replace('${CODEGEN_DECLARE_ARGUMENTS}', codegen_declare_arguments)
        benchmark = benchmark.replace('${CODEGEN_INITIALIZE_ARGUMENTS}', codegen_initialize_arguments)
        benchmark = benchmark.replace('${CODEGEN_CALL_BENCHMARK_FUNCTION}', codegen_call_benchmark_function)

        # Write benchmark source file
        with open(filename, 'w') as f:
            f.writelines(benchmark)


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')
                os.environ['DUNE_CODEGEN_THREADS'] = '1'
                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=5")
                    # 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'].endswith('_mean') or b['name'].endswith('_median') or b['name'].endswith('_stddev'):
                                pass
                            else:
                                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