Subject: Re: [FFC-dev] [DOLFIN-dev] Compile time of forms and linear algebra operators

From: Johan Hake <[email-masked]>
Date: Thu, 3 Jul 2008 21:34:55 +0200

On Monday 30 June 2008 22:28:42 Anders Logg wrote:
> On Mon, Jun 30, 2008 at 10:26:19PM +0200, Martin Sandve Alnæs wrote:
> > 2008/6/30 Anders Logg <[email-masked]>:
> > > On Mon, Jun 30, 2008 at 09:47:31PM +0200, Martin Sandve Alnæs wrote:
> > >> The best way to achieve good UFC compatibility is to limit
> > >> assumptions about the form compilers to a minimum, and
> > >> (a) ufl form
> > >> (b) which compiler?
> > >> (c) a generic set of compiler options (any kind of object)
> > >> sounds to me like a good generalization.
> > >
> > > Sounds good. We can just have a parameter
> > >
> > >  form_compiler_options
> > >
> > > which is passed on to the form compiler in assemble().
> > >
> > > Should the choice of form compiler be an argument to assemble() or
> > > should it be a global option set by dolfin_set()?
> >
> > I think that can safely stay a global option, at least for now.
> >
> > In particular, mixing form compilers in one application will
> > be messy because of dofmaps, and I don't think fixing that
> > is where we want to spend our time.
>
> ok.
>
> Hake, do you want to supply a patch for this (a slight modification of
> your previous patch). I'm on vacation.

Ok, a patch is provided. 

As the suggested interface to jit was jit(form,options), I also needed to 
update ffc.jit. I hope this was what you have intetended. 

In addition to the update in jit functionality I also cleaned up the extracton 
of coefficients.

All python demos are running. I do not know if there are any demos on the 
latets assemble_system that Kent added. If you have a python demo on that 
could you please verify that it runs Kent?

A copy of the dolfin commit message which explain what I did.
********************************************************************
A clean up in assemble.py
  - Putted code for extracting coeffisients in own function
  - Added the parameter "form compiler", which can be set by, e.g.,
    dolfin_set("form compiler","ffc")
  - Dolfin assume that the form compiler implements the jit function, 
    which takes a form as first argument and an optional option parameter
    as second argument.
  - The options parameter is optional. If provided it must be a dict that is
    handed directly to the registered form compilers jit function.
  - Added a private function _jit(form, options) which uses the jit compiler
    that is provided by the module, registered in dolfin_get("form compiler").
  - This commit requires that also ffc is updated to newest version, as 
    the jit interface needed to be updated.
********************************************************************

The change in FFC was done so it was as noninvasive for the rest of the code 
as possible. This made it a bit hackish. I suggest that the parameter system 
for ffc should be looked upon.

It would be nice if ffc kept its options in one dict that could be updated
by users either from jit() or from the command line. As it is now, there are 
several parameters that is provided from function to function.

I also suggest that the names of these options should be intuitive, so e.g., 
changing "precision=" to "precision" and "quadrature_points=" to "quadrature 
points".

Johan
# HG changeset patch
# User "Johan Hake <[email-masked]>"
# Date 1215111425 -7200
# Node ID c60b15e025c0a10ede4e457514b2e6602f69ad65
# Parent  d65f637ae991396daf448dc455cd24402d68ffdc
A clean up in assemble.py
  - Putted code for extracting coeffisients in own function
  - Added the parameter "form compiler", which can be set by, e.g.,
    dolfin_set("form compiler","ffc")
  - Dolfin assume that the form compiler implements the jit function, which takes
    a form as first argument and an optional option parameter as second argument.
  - The options parameter is optional. If provided it must be a dict that is
    handed directly to the registered form compilers jit function.
  - Added a private function _jit(form, options) which uses the jit compiler
    that is provided by the module, registered in dolfin_get("form compiler").
  - This commit requires that also ffc is updated to newest version, as the jit
    interface needed to be updated.

diff -r d65f637ae991 -r c60b15e025c0 site-packages/dolfin/assemble.py
--- a/site-packages/dolfin/assemble.py	Thu Jul 03 20:45:38 2008 +0200
+++ b/site-packages/dolfin/assemble.py	Thu Jul 03 20:57:05 2008 +0200
@@ -36,56 +36,30 @@ _dof_map_cache = {}
 # Cache for tensors
 _tensor_cache = {}
 
+# Temporary storage of compiled coeffisient functions avoiding
+# Python deleting our function objects before assembly!
+_compiled_coefficients = []
+
+# Add a parameter for form compiler
+dolfin_add("form compiler","ffc")
+
 # JIT assembler
 def assemble(form, mesh, coefficients=None, dof_maps=None,
     cell_domains=None, exterior_facet_domains=None, interior_facet_domains=None, reset_tensor=None,
-    tensor=None, backend=None, return_dofmaps=False):
+    tensor=None, backend=None, return_dofmaps=False, form_compiler_option = None):
     "Assemble form over mesh and return tensor"
-
-    # Create empty list of coefficients, filled below
-    _coefficients = ArrayFunctionPtr()
-        
-    # Extract coefficients
-    if not coefficients is None: 
-        # Compile all strings as dolfin::Function
-        string_expressions = []
-        for c in coefficients:
-            # Note: To allow tuples of floats or ints below, this logic becomes more involved...
-            if isinstance(c, (tuple, str)):
-                string_expressions.append(c)
-        if string_expressions:
-            compiled_functions = compile_functions(string_expressions, mesh)
-            compiled_functions.reverse()
-        
-        # Avoid Python deleting our function objects before assembly!
-        remember_coefficients = []
-
-        # Build list of coefficients
-        for c in coefficients:
-            # Note: We could generalize this to support more objects 
-            # like sympy expressions, tuples for constant vectors, etc...
-            if isinstance(c, (float, int)):
-                c = cpp_Function(mesh, float(c))
-            elif isinstance(c, (tuple, str)):
-                c = compiled_functions.pop()
-            _coefficients.push_back(c)
-            remember_coefficients.append(c)
 
     # Check if we need to compile the form (JIT)
     if not hasattr(form, "create_cell_integral"):
         # FFC form, call JIT compile
-        optimize = dolfin_get("optimize form") or dolfin_get("optimize")
-        (compiled_form, module, form_data) = jit(form, optimize=optimize)
-        
-        # Extract coefficients from form if no coefficients are provided
-        if coefficients is None:
-            for c in form_data.coefficients:
-                _coefficients.push_back(c.f)
-        
+        (compiled_form, module, form_data) = _jit(form, form_compiler_option)
     else:    
         # UFC form, no need to compile
         compiled_form = form
-        
+        form_data     = None
+
+    # Extract coefficients
+    _coefficients = _extract_coefficients(coefficients, form_data)
 
     # FIXME: do we need these lines now? None works fine with Assembler
     # Create dummy arguments for domains (not yet supported in Python)
@@ -110,7 +84,10 @@ def assemble(form, mesh, coefficients=No
     # Assemble tensor from compiled form
     cpp_assemble(tensor, compiled_form, mesh, _coefficients, dof_maps,
                  cell_domains, exterior_facet_domains, interior_facet_domains, reset_tensor)
-    
+
+    # Clear any temporary stored coeffisient functions
+    _clear_compiled_coefficients()
+
     # Convert to float for scalars
     if rank == 0:
         tensor = tensor.getval()
@@ -121,91 +98,34 @@ def assemble(form, mesh, coefficients=No
     else:
         return tensor
 
+
 # JIT system assembler
 def assemble_system(A_form, b_form, bc, mesh, A_coefficients=None, b_coefficients=None, A_dof_maps=None, b_dof_maps=None,
     cell_domains=None, exterior_facet_domains=None, interior_facet_domains=None, reset_tensor=None,
-    A_tensor=None, b_tensor=None, backend=None, return_dofmaps=False):
+    A_tensor=None, b_tensor=None, backend=None, return_dofmaps=False, form_compiler_option = None):
     "Assemble form over mesh and return tensor"
 
-    # Create empty list of coefficients, filled below
-    _A_coefficients = ArrayFunctionPtr()
-    _b_coefficients = ArrayFunctionPtr()
+    # Check if we need to compile the A_form (JIT)
+    if not hasattr(A_form, "create_cell_integral"):
+        # FFC form, call JIT compile
+        (A_compiled_form, module, A_form_data) = _jit(A_form, form_compiler_option)
+    else:    
+        # UFC form, no need to compile
+        compiled_form = form
+        A_form_data     = None
+
+    # Check if we need to compile the b_form (JIT)
+    if not hasattr(b_form, "create_cell_integral"):
+        # FFC form, call JIT compile
+        (b_compiled_form, module, b_form_data) = _jit(b_form, form_compiler_option)
+    else:    
+        # UFC form, no need to compile
+        compiled_form = form
+        b_form_data     = None
 
     # Extract coefficients
-    if not A_coefficients is None: 
-        # Compile all strings as dolfin::Function
-        string_expressions = []
-        for c in A_coefficients:
-            # Note: To allow tuples of floats or ints below, this logic becomes more involved...
-            if isinstance(c, (tuple, str)):
-                string_expressions.append(c)
-        if string_expressions:
-            compiled_functions = compile_functions(string_expressions, mesh)
-            compiled_functions.reverse()
-        
-        # Build list of coefficients
-        for c in A_coefficients:
-            # Note: We could generalize this to support more objects 
-            # like sympy expressions, tuples for constant vectors, etc...
-            if isinstance(c, (float, int)):
-                c = cpp_Function(mesh, float(c))
-            elif isinstance(c, (tuple, str)):
-                c = compiled_functions.pop()
-            _A_coefficients.push_back(c)
-
-    # Extract coefficients
-    if not b_coefficients is None: 
-        # Compile all strings as dolfin::Function
-        string_expressions = []
-        for c in b_coefficients:
-            # Note: To allow tuples of floats or ints below, this logic becomes more involved...
-            if isinstance(c, (tuple, str)):
-                string_expressions.append(c)
-        if string_expressions:
-            compiled_functions = compile_functions(string_expressions, mesh)
-            compiled_functions.reverse()
-        
-        # Build list of coefficients
-        for c in b_coefficients:
-            # Note: We could generalize this to support more objects 
-            # like sympy expressions, tuples for constant vectors, etc...
-            if isinstance(c, (float, int)):
-                c = cpp_Function(mesh, float(c))
-            elif isinstance(c, (tuple, str)):
-                c = compiled_functions.pop()
-            _b_coefficients.push_back(c)
-
-
-    # Check if we need to compile the form (JIT)
-    if not hasattr(A_form, "create_cell_integral"):
-        # FFC form, call JIT compile
-        optimize = dolfin_get("optimize form") or dolfin_get("optimize")
-        (A_compiled_form, module, form_data) = jit(A_form, optimize=optimize)
-        
-        # Extract coefficients from form if no coefficients are provided
-        if A_coefficients is None:
-            for c in form_data.coefficients:
-                _A_coefficients.push_back(c.f)
-        
-    else:    
-        # UFC form, no need to compile
-        A_compiled_form = form
-
-    # Check if we need to compile the form (JIT)
-    if not hasattr(b_form, "create_cell_integral"):
-        # FFC form, call JIT compile
-        optimize = dolfin_get("optimize form") or dolfin_get("optimize")
-        (b_compiled_form, module, form_data) = jit(b_form, optimize=optimize)
-        
-        # Extract coefficients from form if no coefficients are provided
-        if b_coefficients is None:
-            for c in form_data.coefficients:
-                _b_coefficients.push_back(c.f)
-        
-    else:    
-        # UFC form, no need to compile
-        b_compiled_form = form
-        
+    _A_coefficients = _extract_coefficients(A_coefficients, A_form_data)
+    _b_coefficients = _extract_coefficients(b_coefficients, b_form_data)
 
     # FIXME: do we need these lines now? None works fine with Assembler
     # Create dummy arguments for domains (not yet supported in Python)
@@ -235,6 +155,8 @@ def assemble_system(A_form, b_form, bc, 
                         b_tensor, b_compiled_form, _b_coefficients, b_dof_maps,  
                         mesh, bc, cell_domains, exterior_facet_domains, interior_facet_domains, reset_tensor)
     
+    # Clear any temporary stored coeffisient functions
+    _clear_compiled_coefficients()
     
     # Return value
     if return_dofmaps:
@@ -263,6 +185,68 @@ def _create_dof_map_set(form, compiled_f
         _dof_map_cache[form] = dof_maps
 
     return dof_maps
+
+def _jit(form, options = None):
+    """ Just in time compile any provided form
+
+    It uses the jit function from the form compiler registered by
+    dolfin_set("form compiler")"""
+
+    # Extract the form compiler
+    compiler_str = dolfin_get("form compiler")
+    try:
+        exec("import %s"%compiler_str)
+        jit = eval("%s.jit"%compiler_str)
+    except ImportError:
+        raise RuntimeError, "Could not import %s form compiler"%compiler_str
+    except AttributeError:
+        raise RuntimeError, "Form compiler must implement the jit function"
+    
+    if options is None:
+        options = {}
+    options["cpp optimize"] = dolfin_get("optimize form") or dolfin_get("optimize")
+    return jit(form, options)
+
+def _extract_coefficients(coefficients, form_data):
+    "Extract provided coefficients"
+    
+    # Create empty list of coefficients, filled below
+    _coefficients = ArrayFunctionPtr()
+        
+    # Extract coefficients
+    if coefficients is None:
+        # If no coefficients are porvided try to extract them from form_data
+        if hasattr(form_data,"coefficients"):
+            for c in form_data.coefficients:
+                _coefficients.push_back(c.f)
+    else:
+        # Compile all strings as dolfin::Function
+        string_expressions = []
+        for c in coefficients:
+            # Note: To allow tuples of floats or ints below, this logic becomes more involved...
+            if isinstance(c, (tuple, str)):
+                string_expressions.append(c)
+        if string_expressions:
+            compiled_functions = compile_functions(string_expressions, mesh)
+            compiled_functions.reverse()
+        
+        # Build list of coefficients
+        for c in coefficients:
+            # Note: We could generalize this to support more objects 
+            # like sympy expressions, tuples for constant vectors, etc...
+            if isinstance(c, (float, int)):
+                c = cpp_Function(mesh, float(c))
+            elif isinstance(c, (tuple, str)):
+                c = compiled_functions.pop()
+            _coefficients.push_back(c)
+            _compiled_coefficients.append(c)
+
+    return _coefficients
+
+def _clear_compiled_coefficients():
+    "Clear stored compiled coeffisient functions" 
+    while _compiled_coefficients:
+        _compiled_coefficients.pop()
 
 def _create_tensor(form, rank, backend, tensor, reset_tensor):
     "Create tensor for form"
@@ -318,7 +302,7 @@ class Function(ffc_Function, cpp_Functio
             else:
                 form = TestFunction(element)*dx
             # Compile form and create dof map
-            (compiled_form, module, form_data) = jit(form)
+            (compiled_form, module, form_data) = _jit(form)
             self.dof_maps = DofMapSet(compiled_form, mesh)
             # Initialize FFC and DOLFIN Function
             ffc_Function.__init__(self, element)
@@ -423,7 +407,7 @@ class LinearPDE:
 
         # FIXME: Maybe there is a better solution?
         # Compile form, needed to create discrete function
-        (compiled_form, module, form_data) = jit(self.a)
+        (compiled_form, module, form_data) = _jit(self.a)
 
         # Apply boundary conditions
         for bc in self.bcs:
@@ -478,7 +462,7 @@ class DirichletBC(cpp_DirichletBC):
         "Apply boundary condition to linear system"
         
         # Compile form
-        (compiled_form, module, form_data) = jit(form)        
+        (compiled_form, module, form_data) = _jit(form)        
 
         # Create dof maps
         dof_maps = DofMapSet(compiled_form, self.mesh())
@@ -490,7 +474,7 @@ class DirichletBC(cpp_DirichletBC):
         "Apply boundary condition to linear system"
         
         # Compile form
-        (compiled_form, module, form_data) = jit(form)        
+        (compiled_form, module, form_data) = _jit(form)        
 
         # Create dof maps
         dof_maps = DofMapSet(compiled_form, self.mesh())
@@ -509,7 +493,7 @@ class PeriodicBC(cpp_PeriodicBC):
         "Apply boundary condition to linear system"
         
         # Compile form
-        (compiled_form, module, form_data) = jit(form)        
+        (compiled_form, module, form_data) = _jit(form)        
 
         # Create dof maps
         dof_maps = DofMapSet(compiled_form, self.mesh())
# HG changeset patch
# User "Johan Hake <[email-masked]>"
# Date 1215111845 -7200
# Node ID 77e33e2dbc3ccd7abfe644d4e7a00a8bf7fd2155
# Parent  ca574155bd2d52db55beea7d04b0ca19894d3264
Change the interface to jit. It now takes a form as first argument and
and optional options argument at second. It now complies to the interface
used in pyDOLFIN's assemble function.

diff -r ca574155bd2d -r 77e33e2dbc3c src/ffc/jit/jit.py
--- a/src/ffc/jit/jit.py	Tue Jun 24 07:31:48 2008 +0200
+++ b/src/ffc/jit/jit.py	Thu Jul 03 21:04:05 2008 +0200
@@ -33,11 +33,36 @@ FFC_OPTIONS_JIT = FFC_OPTIONS.copy()
 #FFC_OPTIONS_JIT["no-evaluate_basis"] = True
 FFC_OPTIONS_JIT["no-evaluate_basis_derivatives"] = True
 
-def jit(input_form, representation=FFC_REPRESENTATION, language=FFC_LANGUAGE, options=FFC_OPTIONS_JIT, optimize=False):
-    "Just-in-time compile the given form or element"
+def jit(input_form, options = None):
+    """ Just-in-time compile the given form or element
+    
+    Parameters:
+    input_form : The form
+    options    : An option dict. 
+    """
 
+    # Collect options
+    _options = FFC_OPTIONS_JIT.copy()
+    if options is None:
+        # Default options
+        cpp_optimize   = False
+        representation = FFC_REPRESENTATION
+        language       = FFC_LANGUAGE
+    elif isinstance(options,dict):
+        cpp_optimize   = options.pop("cpp optimize",False)
+        representation = options.pop("representation",FFC_REPRESENTATION)
+        language       = options.pop("language",FFC_LANGUAGE)
+        for key, value in options.iteritems():
+            if _options.has_key(key):
+                _options[key] = value
+            else:
+                # FIXME: Warn that options is not set?
+                pass
+    else:
+        raise RuntimeError, "options must be a dict"
+        
     # Set C++ compiler options
-    if optimize:
+    if cpp_optimize: 
         cpp_args = "-O2"
     else:
         cpp_args = "-O0"
@@ -57,7 +82,7 @@ def jit(input_form, representation=FFC_R
     # Compute md5 checksum of form signature
     signature = " ".join([str(form),
                           ", ".join([element.signature() for element in form_data.elements]),
-                          representation, language, str(options), cpp_args])
+                          representation, language, str(_options), cpp_args])
     md5sum = "form_" + md5.new(signature).hexdigest()
 
     # Get name of form
@@ -102,7 +127,7 @@ def jit(input_form, representation=FFC_R
     if compiled_form is None:
         
         # Build form module
-        build_module(form, representation, language, options, md5sum, form_dir, module_dir, prefix, cpp_args)
+        build_module(form, representation, language, _options, md5sum, form_dir, module_dir, prefix, cpp_args)
 
         # Import form module
         sys.path.append(form_dir)

---

From: Anders Logg <[email-masked]>
Date: Mon, 7 Jul 2008 20:06:14 +0200

Very nice!

I've applied the patches and they seem to work fine.

I agree with your suggestions below, but will fix later.

-- 
Anders