Skip to content

Commit b178e00

Browse files
committed
api: Start enabling specialization at operator apply
1 parent e70b6d5 commit b178e00

1 file changed

Lines changed: 22 additions & 3 deletions

File tree

devito/operator/operator.py

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from devito.ir.clusters import ClusterGroup, clusterize
2222
from devito.ir.iet import (Callable, CInterface, EntryFunction, DeviceFunction,
2323
FindSymbols, MetaCall, derive_parameters, iet_build)
24+
from devito.ir.iet.visitors import Specializer
2425
from devito.ir.support import AccessMode, SymbolRegistry
2526
from devito.ir.stree import stree_build
2627
from devito.operator.profiling import create_profile
@@ -985,16 +986,34 @@ def apply(self, **kwargs):
985986
>>> op = Operator(Eq(u3.forward, u3 + 1))
986987
>>> summary = op.apply(time_M=10)
987988
"""
988-
# Compile the operator before building the arguments list
989-
# to avoid out of memory with greedy compilers
990-
cfunction = self.cfunction
989+
# Get items expected to be specialized
990+
specialize = as_tuple(kwargs.pop('specialize', []))
991+
992+
if not specialize:
993+
# Compile the operator before building the arguments list
994+
# to avoid out of memory with greedy compilers
995+
cfunction = self.cfunction
991996

992997
# Build the arguments list to invoke the kernel function
993998
with self._profiler.timer_on('arguments-preprocess'):
994999
args = self.arguments(**kwargs)
9951000
with switch_log_level(comm=args.comm):
9961001
self._emit_args_profiling('arguments-preprocess')
9971002

1003+
# In the case of specialization, arguments must be processed before
1004+
# the operator can be compiled
1005+
if specialize:
1006+
specialized_args = {p: sympify(args.pop(p.name))
1007+
for p in self.parameters if p.name in specialize}
1008+
1009+
op = Specializer(specialized_args).visit(self)
1010+
else:
1011+
op = self
1012+
1013+
from IPython import embed; embed()
1014+
1015+
# TODO: Whose profiler should get used here?
1016+
9981017
# Invoke kernel function with args
9991018
arg_values = [args[p.name] for p in self.parameters]
10001019
try:

0 commit comments

Comments
 (0)