Merge pull request #1017 from vmickus/fix/using-func_table-ops

Replace _ops_kernels by _func_table in OPS

Author: FabioLuporini

devito/ops/operator.py

@@ -1,6 +1,7 @@
 from devito import Eq
 from devito.ir.equations import ClusterizedEq
-from devito.ir.iet import Call, Expression, List, derive_parameters, find_affine_trees
+from devito.ir.iet import (Call, Expression, List, MetaCall, derive_parameters,
+                           find_affine_trees)
 from devito.ir.iet.visitors import FindSymbols, Transformer
 from devito.logger import warning
 from devito.operator import Operator
@@ -26,7 +27,6 @@ class OperatorOPS(Operator):
     _default_headers = Operator._default_headers + ['#define restrict __restrict']
 
     def __init__(self, *args, **kwargs):
-        self._ops_kernels = []
         super(OperatorOPS, self).__init__(*args, **kwargs)
         self._compiler = ops_configuration['compiler'].copy()
 
@@ -88,7 +88,8 @@ def _specialize_iet(self, iet, **kwargs):
             )
 
             pre_time_loop.extend(pre_loop)
-            self._ops_kernels.append(ops_kernel)
+            self._func_table[namespace['ops_kernel_file'](ops_kernel.name)] = \
+                MetaCall(ops_kernel, False)
             mapper[tree[0].root] = ops_par_loop_call
             mapper.update({i.root: mapper.get(i.root) for i in tree})  # Drop trees
 
@@ -108,13 +109,13 @@ def _specialize_iet(self, iet, **kwargs):
 
     @property
     def hcode(self):
-        return ''.join(str(kernel) for kernel in self._ops_kernels)
+        return ''.join(str(kernel.root) for kernel in self._func_table.values())
 
     def _finalize(self, iet, parameters):
-        # Applies _finalize for each generated OPS kernel, which will generate the .h file
-        self._ops_kernels = [super(OperatorOPS, self)._finalize(
-            kernel, derive_parameters(kernel, True))
-            for kernel in self._ops_kernels]
+        for k, v in self._func_table.items():
+            parameters = derive_parameters(v.root, True)
+            root = super(OperatorOPS, self)._finalize(v.root, parameters)
+            self._func_table[k] = MetaCall(root, v.local)
 
         return super()._finalize(iet, parameters)
 

devito/ops/utils.py

@@ -50,3 +50,5 @@
 namespace['ops_dat_d_p'] = lambda i: '%s_d_p' % i
 namespace['ops_dat_d_m'] = lambda i: '%s_d_m' % i
 namespace['ops_dat_name'] = lambda i: '%s_dat' % i
+
+namespace['ops_kernel_file'] = lambda i: 'ops_kernel_%s' % i

tests/test_ops.py

@@ -86,7 +86,8 @@ def test_kernel_generation(self, equation, expected):
 
         operator = Operator(eval(equation))
 
-        assert str(operator._ops_kernels[0]) == expected
+        for func in operator._func_table.values():
+            assert str(func.root) == expected
 
     @pytest.mark.parametrize('equation, expected', [
         ('Eq(u,3*a - 4**a)', '{ "ut0": [[0]] }'),