options.py

""" Manage the command line options to the form compiler executable """

from argparse import ArgumentParser
from os.path import abspath
from pytools import ImmutableRecord, memoize
from contextlib import contextmanager

from dune.testtools.parametertree.parser import parse_ini_file


class CodegenOption(ImmutableRecord):
    """ Data structure representing a single formcompiler option """
    def __init__(self,
                 default=None,
                 helpstr="Undocumented feature!",
                 process=lambda x: x,
                 _type=type(None),
                 ):
        _type = type(default)
        if issubclass(_type, type(None)):
            _type = str
        ImmutableRecord.__init__(self,
                                 helpstr=helpstr,
                                 default=default,
                                 type=_type,
                                 )


class CodegenGlobalOptionsArray(ImmutableRecord):
    """ A collection of form compiler arguments """
    def __init__(self, **kwargs):
        opts = {k: v.default for k, v in CodegenGlobalOptionsArray.__dict__.items() if isinstance(v, CodegenOption)}
        opts.update(**kwargs)
        ImmutableRecord.__init__(self, **opts)

    # Arguments that are to be set from the outside
    uflfile = CodegenOption(helpstr="the UFL file to compile")
    debug_cache_with_stack = CodegenOption(default=False, helpstr="Store stack along with cache objects. Makes debugging caching issues easier.")
    driver_file = CodegenOption(helpstr="The filename for the generated driver header")
    explicit_time_stepping = CodegenOption(default=False, helpstr="use explicit time stepping")
    time_stepping_order = CodegenOption(default=1, helpstr="Order of the time stepping method")
    exact_solution_expression = CodegenOption(helpstr="name of the exact solution expression in the ufl file")
    compare_l2errorsquared = CodegenOption(helpstr="maximal allowed l2 error squared of difference between numerical solution and interpolation of exact solution (NOTE: requires --exact-solution-expression)")
    grid_info = CodegenOption(default=None, helpstr="Path to file with information about facedir and facemod variations. This can be used to limit the generation of skeleton kernels.")
    l2error_tree_path = CodegenOption(default=None, helpstr="Tree pathes that should be considered for l2 error calculation. Default None means we take all of them into account.")
    ini_file = CodegenOption(helpstr="An inifile to use. A generated driver will be hard-coded to it, a [formcompiler] section will be used as default values to form compiler arguments (use snake case)")
    opcounter = CodegenOption(default=False, helpstr="Count operations. Note: In this case only operator applications are generated since solving and operator counting does not work. You probably want to set instrumentation level>0.")
    performance_measuring = CodegenOption(default=False, helpstr="Generate opcounter codepath, but only measure times!")
    instrumentation_level = CodegenOption(default=0, helpstr="Control time/opcounter measurements. 0-do nothing, 1-measure program as a whole, 2-operator applications, 3-measure kernel (eg. alpha-volume, ...), 4-parts of kernel (eg. stage 1-3 of SF)")
    project_basedir = CodegenOption(helpstr="The base (build) directory of the dune-codegen project")
    architecture = CodegenOption(default="haswell", helpstr="The architecture to optimize for. Possible values: haswell|knl|skylake")
    yaspgrid_offset = CodegenOption(default=False, helpstr="Set to true if you want a yasp grid where the lower left corner is not in the origin.")
    grid_unstructured = CodegenOption(default=False, helpstr="Set to true if you want to use an unstructured grid.")
    grid_consistent = CodegenOption(default=False, helpstr="The used grid is already consistent.")
    precision_bits = CodegenOption(default=64, helpstr="The number of bits for the floating point type")
    overlapping = CodegenOption(default=False, helpstr="Use an overlapping solver and constraints. You still need to make sure to construct a grid with overlap! The parallel option will be set automatically.")
    operators = CodegenOption(default="r", helpstr="A comma separated list of operators, each name will be interpreted as a subsection name within the formcompiler section")
    target_name = CodegenOption(default=None, helpstr="The target name from CMake")
    operator_to_build = CodegenOption(default=None, helpstr="The operators from the list that is about to be build now. CMake sets this one!!!")
    debug_interpolate_input = CodegenOption(default=False, helpstr="Should the input for printresidual and printmatix be interpolated (instead of random input).")
    use_likwid = CodegenOption(default=False, helpstr="Use likwid instead of own performance measurements.")
    use_sde = CodegenOption(default=False, helpstr="Use sde instead of own performance measurements.")
    autotune_google_benchmark = CodegenOption(default=False, helpstr="Use google-benchmark library for autotuning (when autotuning is activated).")
    with_mpi = CodegenOption(default=True, helpstr="The module was configured with mpi")
    permuting_horizontal_add = CodegenOption(default=True, helpstr="Whether SIMD horizontal_add should use a permuting implementation.")

    # Arguments that are mainly to be set by logic depending on other options
    max_vector_width = CodegenOption(default=256, helpstr=None)
    parallel = CodegenOption(default=False, helpstr="Mark that this program should be run in parallel. If set to true the c++ code will check that there are more than 1 MPI-ranks involved and the error computation will use communication.")


class CodegenFormOptionsArray(ImmutableRecord):
    """ A collection of form-specific form compiler arguments """
    def __init__(self, **kwargs):
        opts = {k: v.default for k, v in CodegenFormOptionsArray.__dict__.items() if isinstance(v, CodegenOption)}
        opts.update(**kwargs)
        ImmutableRecord.__init__(self, **opts)

    # Form specific options
    form = CodegenOption(default=None, helpstr="The name of the UFL object representing the form in the UFL file")
    filename = CodegenOption(default=None, helpstr="The filename to use for this LocalOperator")
    classname = CodegenOption(default=None, helpstr="The name of the C++ class to generate")
    numerical_jacobian = CodegenOption(default=False, helpstr="use numerical jacobians (only makes sense, if uflpdelab for some reason fails to generate analytic jacobians)")
    matrix_free = CodegenOption(default=False, helpstr="Generate jacobian_apply_* methods for matrix free solvers")
    print_transformations = CodegenOption(default=False, helpstr="print out dot files after ufl tree transformations")
    print_transformations_dir = CodegenOption(default=".", helpstr="place where to put dot files (can be omitted)")
    quadrature_order = CodegenOption(_type=int, helpstr="Quadrature order used for all integrals.")
    diagonal_transformation_matrix = CodegenOption(default=False, helpstr="set option if the jacobian of the transformation is diagonal (axiparallel grids)")
    constant_transformation_matrix = CodegenOption(default=False, helpstr="set option if the jacobian of the transformation is constant on a cell")
    fastdg = CodegenOption(default=False, helpstr="Use FastDGGridOperator from PDELab.")
    sumfact = CodegenOption(default=False, helpstr="Use sumfactorization")
    sumfact_regular_jacobians = CodegenOption(default=False, helpstr="Generate non sum-factorized jacobians (only useful if sumfact is set)")
    sumfact_on_boundary = CodegenOption(default=True, helpstr="Whether boundary integrals should be vectorized. It might not be worth the hassle...")
    sumfact_optimize_loop_order = CodegenOption(default=False, helpstr="Optimize order of loops in sumf factorization function using autotuning.")
    sumfact_performance_transformations = CodegenOption(default=False, helpstr="Apply sum factorization specific performance transformations.")
    sumfact_performance_transformations_testrun = CodegenOption(default=0, helpstr="If larger than zero determines test case to run.")
    vectorization_quadloop = CodegenOption(default=False, helpstr="whether to generate code with explicit vectorization")
    vectorization_strategy = CodegenOption(default="none", helpstr="The identifier of the vectorization cost model. Possible values: none|explicit|model|target|autotune")
    vectorization_not_fully_vectorized_error = CodegenOption(default=False, helpstr="throw an error if nonquadloop vectorization did not fully vectorize")
    vectorization_horizontal = CodegenOption(default=None, helpstr="an explicit value for horizontal vectorization read by the 'explicit' strategy")
    vectorization_vertical = CodegenOption(default=None, helpstr="an explicit value for vertical vectorization read by the 'explicit' strategy")
    vectorization_padding = CodegenOption(default=None, helpstr="an explicit value for the allowed padding in vectorization")
    vectorization_allow_quadrature_changes = CodegenOption(default=False, helpstr="whether the vectorization strategy is allowed to alter quadrature point numbers")
    vectorization_list_index = CodegenOption(default=None, helpstr="Which vectorization to pick from a list (only valid with vectorization_strategy=fromlist).")
    vectorization_jacobians = CodegenOption(default=True, helpstr="Whether to attempt to vectorize jacobians (takes time, often not needed)")
    vectorization_target = CodegenOption(_type=float, helpstr="The cost function target for the 'target' cost model. Only needed to verify the cost model itself, do not use light-heartedly!!!")
    simplify = CodegenOption(default=False, helpstr="Whether to simplify expressions using sympy")
    generate_jacobians = CodegenOption(default=True, helpstr="Whether jacobian_* methods should be generated. This is set to false automatically, when numerical_jacobian is set to true.")
    generate_jacobian_apply = CodegenOption(default=False, helpstr="Wether jacobian_allpy_* methods should be generated.")
    generate_residuals = CodegenOption(default=True, helpstr="Whether alpha_* methods should be generated.")
    unroll_dimension_loops = CodegenOption(default=False, helpstr="whether loops over the geometric dimension should be unrolled")
    blockstructured = CodegenOption(default=False, helpstr="Use block structure")
    number_of_blocks = CodegenOption(default=1, helpstr="Number of sub blocks in one direction")
    vectorization_blockstructured = CodegenOption(default=False, helpstr="Vectorize block structuring")
    vectorization_blockstructured_tail = CodegenOption(default=True, helpstr="Try to fully vectorize block structuring even when 'nunmber_of_blocks' is not divisible by vector length")
    vectorization_blockstructured_tail_ordering = CodegenOption(default='consecutive', helpstr="Ordering of the tail w.r.t the vectorized loop. Possible values: consecutive|blocked")
    adjoint = CodegenOption(default=False, helpstr="Generate adjoint operator")
    control = CodegenOption(default=False, helpstr="Generate operator of derivative w.r.t. the control variable")
    objective_function = CodegenOption(default=None, helpstr="Name of form representing the objective function in UFL file")
    control_variable = CodegenOption(default=None, helpstr="Name of control variable in UFL file")
    block_preconditioner_diagonal = CodegenOption(default=False, helpstr="Whether this operator should implement the diagonal part of a block preconditioner")
    block_preconditioner_offdiagonal = CodegenOption(default=False, helpstr="Whether this operator should implement the off-diagonal part of a block preconditioner")
    block_preconditioner_pointdiagonal = CodegenOption(default=False, helpstr="Whether this operator should implement the point diagonal part of a block preconditioner")
    geometry_mixins = CodegenOption(default="generic", helpstr="A comma separated list of mixin identifiers to use for geometries. Currently implemented mixins: generic, axiparallel, equidistant, sumfact_multilinear, sumfact_axiparallel, sumfact_equidistant")
    quadrature_mixins = CodegenOption(default="generic", helpstr="A comma separated list of mixin identifiers to use for quadrature. Currently implemented: generic, sumfact")
    basis_mixins = CodegenOption(default="generic", helpstr="A comma separated list of mixin identifiers to use for basis function evaluation. Currently implemented: generic, sumfact")
    accumulation_mixins = CodegenOption(default="generic", helpstr="A comma separated list of mixin identifiers to use for accumulation. Currently implemented: generic, sumfact, control, blockstructured")
    enable_volume = CodegenOption(default=True, helpstr="Whether to assemble volume integrals")
    enable_skeleton = CodegenOption(default=True, helpstr="Whether to assemble skeleton integrals")
    enable_boundary = CodegenOption(default=True, helpstr="Whether to assemble boundary integrals")


# Until more sophisticated logic is needed, we keep the actual option data in this module
_global_options = CodegenGlobalOptionsArray()
_form_options = {}


def show_options():
    def subopt(arr):
        for k, v in arr.__dict__.items():
            if isinstance(v, CodegenOption) and v.helpstr is not None:
                print("{}\n    {}".format(k, v.helpstr))

    print("This is a summary of options available for the code generation process:\n")
    print("The following options can be given in the [formcompiler] section:")
    subopt(CodegenGlobalOptionsArray)

    print("\nThefollowing options can be given in a form-specific subsection of [formcompiler]:")
    subopt(CodegenFormOptionsArray)


def initialize_options():
    """ Initialize the options from the command line """
    global _global_options
    _global_options = update_options_from_commandline(_global_options)
    _global_options = update_options_from_inifile(_global_options)


def update_options_from_commandline(opt):
    """ Return an options array object with updated values from the commandline """
    assert isinstance(opt, CodegenGlobalOptionsArray)
    parser = ArgumentParser(description="Compile UFL files to PDELab C++ code",
                            epilog="Please report bugs to dominic.kempf@iwr.uni-heidelberg.de",
                            )
    parser.add_argument('--version', action='version', version='%(prog)s 0.1')
    for k, v in type(opt).__dict__.items():
        if isinstance(v, CodegenOption) and v.helpstr is not None:
            cmdopt = "--{}".format(k.replace('_', '-'))
            parser.add_argument(cmdopt, help=v.helpstr, type=v.type)
    parsedargs = {k: v for k, v in vars(parser.parse_args()).items() if v is not None}
    return opt.copy(**parsedargs)


def update_options_from_inifile(opt):
    """ Return an options array object with updated values from an inifile """
    if opt.ini_file:
        def parse_ini(section, opttype):
            def _fix_types(k, v):
                if hasattr(opttype, k) and getattr(opttype, k).type is bool:
                    return bool(eval(v))
                if hasattr(opttype, k):
                    return getattr(opttype, k).type(v)
                return v
            ini = parse_ini_file(opt.ini_file).get(section, {})
            return {k: _fix_types(k, v) for k, v in ini.items()}

        opt = opt.copy(**parse_ini("formcompiler", CodegenGlobalOptionsArray))
        # Also parse form-specific options
        for form in [i.strip() for i in opt.operators.split(",")]:
            _form_options[form] = CodegenFormOptionsArray(**parse_ini("formcompiler.{}".format(form), CodegenFormOptionsArray))

    return opt


@memoize
def process_global_options(opt):
    """ Make sure that the options have been fully processed """
    opt = expand_architecture_options(opt)

    if opt.overlapping:
        opt = opt.copy(parallel=True)

    return opt


@memoize
def process_form_options(opt, form):
    if opt.sumfact:
        opt = opt.copy(unroll_dimension_loops=True,
                       quadrature_mixins="sumfact",
                       basis_mixins="sumfact",
                       accumulation_mixins="sumfact",
                       )

    if opt.blockstructured:
        opt = opt.copy(accumulation_mixins="blockstructured",
                       quadrature_mixins="blockstructured",
                       basis_mixins="blockstructured"
                       )

    if opt.control:
        opt = opt.copy(accumulation_mixins="control")

    if opt.numerical_jacobian:
        opt = opt.copy(generate_jacobians=False, generate_jacobian_apply=False)

    if opt.form is None:
        opt = opt.copy(form=form)

    if opt.classname is None:
        opt = opt.copy(classname="{}Operator".format(form))

    if opt.filename is None:
        opt = opt.copy(filename="{}_{}_file.hh".format(get_option("target_name"), opt.classname))

    if opt.block_preconditioner_pointdiagonal:
        opt = opt.copy(generate_jacobians=False,
                       basis_mixins="sumfact_pointdiagonal",
                       accumulation_mixins="sumfact_pointdiagonal",
                       )

    if opt.block_preconditioner_diagonal or opt.block_preconditioner_offdiagonal:
        assert opt.numerical_jacobian is False
        opt = opt.copy(generate_residuals=False,
                       generate_jacobians=True,
                       matrix_free=True,
                       )

    if opt.matrix_free:
        opt = opt.copy(generate_jacobian_apply=True)

    return opt


def expand_architecture_options(opt):
    if opt.architecture == "haswell":
        return opt.copy(max_vector_width=256)
    elif opt.architecture == "knl":
        return opt.copy(max_vector_width=512)
    elif opt.architecture == "skylake":
        return opt.copy(max_vector_width=512)
    else:
        raise NotImplementedError("Architecture {} not known!".format(opt.architecture))


def set_option(key, value):
    """Add the key value pair to the options.

    If the key is already in the options dictionary its value will be
    overwritten.  Form compiler arguments will always be set before
    any other options.
    """
    global _global_options
    _global_options = process_global_options(_global_options).copy(**{key: value})


def set_form_option(key, value, form=None):
    if form is None:
        from dune.codegen.generation import get_global_context_value
        form = get_global_context_value("form_identifier", 0)
    if isinstance(form, int):
        form = get_option("operators").split(",")[form].strip()
    _form_options[form] = _form_options[form].copy(**{key: value})


def get_option(key):
    processed_global_opts = process_global_options(_global_options)
    return getattr(processed_global_opts, key)


def get_form_option(key, form=None):
    if form is None:
        from dune.codegen.generation import get_global_context_value
        form = get_global_context_value("form_identifier", 0)
    if isinstance(form, int):
        form = get_option("operators").split(",")[form].strip()
    processed_form_opts = process_form_options(_form_options[form], form)
    return getattr(processed_form_opts, key)


@contextmanager
def option_context(conditional=True, **opts):
    """ A context manager that sets a given option and restores it on exit. """
    # Backup old values and set to new ones
    if conditional:
        backup = {}
        for k, v in opts.items():
            backup[k] = get_option(k)
            set_option(k, v)

    yield

    if conditional:
        # Restore old values
        for k in opts.keys():
            set_option(k, backup[k])


@contextmanager
def form_option_context(conditional=True, **opts):
    """ A context manager that sets a given form option and restores it on exit """
    if conditional:
        form = opts.pop("form", None)

        # Backup old values and set to new ones
        backup = {}
        for k, v in opts.items():
            backup[k] = get_form_option(k, form=form)
            set_form_option(k, v, form=form)

    yield

    # Restore old values
    if conditional:
        for k in opts.keys():
            set_form_option(k, backup[k], form=form)