refactoring timings.py

python/dune/codegen/pdelab/driver/timings.py

 """ Timing related generator functions """
 
-from dune.codegen.options import get_option
 from dune.codegen.generation import (cached,
                                      include_file,
                                      pre_include,
-                                     post_include,
                                      preamble,
                                      )
+from dune.codegen.options import get_option
 from dune.codegen.pdelab.driver import (get_form_ident,
                                         is_linear,
                                         name_initree,
@@ -21,7 +20,7 @@ from dune.codegen.pdelab.driver.gridoperator import (name_gridoperator,
                                                      type_gridoperator,
                                                      )
 from dune.codegen.pdelab.driver.solve import (name_vector,
-                                              type_vector,
+                                              define_vector,
                                               )
 
 
@@ -90,6 +89,26 @@ def name_timing_stream():
     return name
 
 
+def name_temporary_vector(name, form):
+    name = "{}_{}".format(name, form)
+    define_vector(name, form)
+    return name
+
+
+@preamble(section="timings")
+def define_jacobian(name, form_ident):
+    t_go = type_gridoperator(form_ident)
+    n_go = name_gridoperator(form_ident)
+    return ["using M_{} = typename {}::Traits::Jacobian;".format(form_ident, t_go),
+            "M_{} {}({});".format(form_ident, name, n_go)]
+
+
+def name_jacobian(form_ident):
+    name = "J_{}".format(form_ident)
+    define_jacobian(name, form_ident)
+    return name
+
+
 @cached
 def setup_timer():
     # TODO check that we are using YASP?
@@ -99,100 +118,88 @@ def setup_timer():
     include_file("dune/codegen/common/timer.hh", filetag="driver")
 
 
-@preamble(section="timings")
-def evaluate_residual_timer():
-    n_go = name_gridoperator(get_form_ident())
-    v = name_vector(get_form_ident())
-    t_v = type_vector(get_form_ident())
+def init_region_timer(region):
     setup_timer()
+    from dune.codegen.generation import post_include
+    post_include("HP_DECLARE_TIMER({});".format(region), filetag="driver")
 
-    if get_option('instrumentation_level') >= 2:
-        # Write back times
-        from dune.codegen.generation import post_include
-        post_include("HP_DECLARE_TIMER(residual_evaluation);", filetag="driver")
-        timestream = name_timing_stream()
-        print_times = []
 
-    lop_name = name_localoperator(get_form_ident())
+def start_region_timer(region):
+    return ["HP_TIMER_START({});".format(region)]
+
+
+def stop_region_timer(region):
+    timestream = name_timing_stream()
+    return ["HP_TIMER_STOP({});".format(region),
+            "DUMP_TIMER({}, {}, {}, true);".format(get_option("instrumentation_level"), region, timestream)]
+
+
+def timed_region(region, actions):
+    if isinstance(actions, str):
+        actions = [actions]
+
+    assert(isinstance(actions, list))
+
+    assembly = []
+    print_times = []
+
+    init_region_timer(region)
+
     if get_option('instrumentation_level') >= 3:
+        timestream = name_timing_stream()
+        lop_name = name_localoperator(get_form_ident())
         print_times.append("{}.dump_timers({}, {}, true);".format(lop_name, timestream, name_timing_identifier()))
 
-    if get_option('instrumentation_level') >= 2:
-        evaluation = ["HP_TIMER_START(residual_evaluation);",
-                      "{}.residual({}, r);".format(n_go, v),
-                      "HP_TIMER_STOP(residual_evaluation);",
-                      "DUMP_TIMER({}, residual_evaluation, {}, true);".format(get_option("instrumentation_level"), timestream)]
-        evaluation.extend(print_times)
-    else:
-        evaluation = ["{}.residual({}, r);".format(n_go, v)]
-
-    evaluation = ["{} r({});".format(t_v, v), "r=0.0;"] + evaluation
+    assembly += start_region_timer(region)
+    assembly += actions
+    assembly += stop_region_timer(region)
 
-    return evaluation
+    return assembly + print_times
 
 
 @preamble(section="timings")
-def apply_jacobian_timer():
+def evaluate_residual_timer():
     n_go = name_gridoperator(get_form_ident())
     v = name_vector(get_form_ident())
-    t_v = type_vector(get_form_ident())
-    setup_timer()
+    r = name_temporary_vector("r", get_form_ident())
+
+    action = "{}.residual({}, {});".format(n_go, v, r)
 
     if get_option('instrumentation_level') >= 2:
-        # Write back times
-        from dune.codegen.generation import post_include
-        post_include("HP_DECLARE_TIMER(apply_jacobian);", filetag="driver")
-        timestream = name_timing_stream()
-        print_times = []
+        return timed_region("residual_evaluation", action)
+    else:
+        return action
 
-    lop_name = name_localoperator(get_form_ident())
-    if get_option('instrumentation_level') >= 3:
-        print_times.append("{}.dump_timers({}, {}, true);".format(lop_name, timestream, name_timing_identifier()))
+
+@preamble(section="timings")
+def apply_jacobian_timer():
+    form = get_form_ident()
+    n_go = name_gridoperator(form)
+    v = name_vector(form)
 
     if is_linear():
-        declaration = ["{} j({});".format(t_v, v), "j=0.0;"]
-        evaluation = ["{}.jacobian_apply({}, j);".format(n_go, v)]
+        j = name_temporary_vector("j", form)
+        action = "{}.jacobian_apply({}, {});".format(n_go, v, j)
     else:
-        declaration = ["{} j0({});".format(t_v, v), "j0=0.0;",
-                       "{} j1({});".format(t_v, v), "j1=0.0;"]
-        evaluation = ["{}.nonlinear_jacobian_apply({}, j0, j1);".format(n_go, v)]
+        j0 = name_temporary_vector("j0", form)
+        j1 = name_temporary_vector("j1", form)
+        action = "{}.nonlinear_jacobian_apply({}, {}, {});".format(n_go, v, j0, j1)
 
     if get_option('instrumentation_level') >= 2:
-        evaluation = ["HP_TIMER_START(apply_jacobian);"] + evaluation + ["HP_TIMER_STOP(apply_jacobian);", "DUMP_TIMER({}, apply_jacobian, {}, true);".format(get_option("instrumentation_level"), timestream)]
-        evaluation.extend(print_times)
-
-    return declaration + evaluation
+        return timed_region("apply_jacobian", action)
+    else:
+        return action
 
 
 @preamble(section="timings")
 def assemble_matrix_timer():
-    t_go = type_gridoperator(get_form_ident())
     n_go = name_gridoperator(get_form_ident())
     v = name_vector(get_form_ident())
-    t_v = type_vector(get_form_ident())
-    setup_timer()
-
-    if get_option('instrumentation_level') >= 2:
-        # Write back times
-        from dune.codegen.generation import post_include
-        post_include("HP_DECLARE_TIMER(matrix_assembly);", filetag="driver")
-        timestream = name_timing_stream()
-        print_times = []
+    m = name_jacobian(get_form_ident())
 
-    lop_name = name_localoperator(get_form_ident())
-    if get_option('instrumentation_level') >= 3:
-        print_times.append("{}.dump_timers({}, {}, true);".format(lop_name, timestream, name_timing_identifier()))
+    action = "{}.jacobian({},{});".format(n_go, v, m)
 
     if get_option('instrumentation_level') >= 2:
-        assembly = ["HP_TIMER_START(matrix_assembly);",
-                    "{}.jacobian({},m);".format(n_go, v),
-                    "HP_TIMER_STOP(matrix_assembly);",
-                    "DUMP_TIMER({}, matrix_assembly, {}, true);".format(get_option("instrumentation_level"), timestream)]
-        assembly.extend(print_times)
+        return timed_region("matrix_assembly", action)
     else:
-        assembly = ["{}.jacobian({},m);".format(n_go, v)]
-
-    assembly = ["using M = typename {}::Traits::Jacobian;".format(t_go),
-                "M m({});".format(n_go)] + assembly
-
-    return assembly
+        return [action]