register_optimization_preambles for all kernels

Author: isuruf

sumpy/e2e.py

@@ -29,6 +29,7 @@
 
 from loopy.version import MOST_RECENT_LANGUAGE_VERSION
 from sumpy.tools import KernelCacheMixin, to_complex_dtype
+from sumpy.codegen import register_optimization_preambles
 from pytools import memoize_method
 
 import logging
@@ -145,6 +146,7 @@ def get_optimized_kernel(self):
         # FIXME
         knl = self.get_kernel()
         knl = lp.split_iname(knl, "itgt_box", 64, outer_tag="g.0", inner_tag="l.0")
+        knl = register_optimization_preambles(knl, self.device)
 
         return knl
 
@@ -279,6 +281,7 @@ def get_optimized_kernel(self):
         # FIXME
         knl = self.get_kernel()
         knl = lp.split_iname(knl, "itgt_box", 64, outer_tag="g.0", inner_tag="l.0")
+        knl = register_optimization_preambles(knl, self.device)
 
         return knl
 
@@ -518,6 +521,7 @@ def get_optimized_kernel(self, result_dtype):
         knl = self.get_kernel(result_dtype)
         knl = self.tgt_expansion.m2l_translation.optimize_loopy_kernel(
                 knl, self.tgt_expansion, self.src_expansion)
+        knl = register_optimization_preambles(knl, self.device)
 
         return knl
 
@@ -627,6 +631,7 @@ def get_optimized_kernel(self, result_dtype):
         knl = self.get_kernel(result_dtype)
         knl = lp.tag_inames(knl, "idim*:unr")
         knl = lp.tag_inames(knl, {"itr_class": "g.0"})
+        knl = register_optimization_preambles(knl, self.device)
 
         return knl
 
@@ -732,6 +737,7 @@ def get_optimized_kernel(self, result_dtype):
         _, optimizations = self.get_inner_knl_and_optimizations(result_dtype)
         for optimization in optimizations:
             knl = optimization(knl)
+        knl = register_optimization_preambles(knl, self.device)
         return knl
 
     def __call__(self, queue, **kwargs):
@@ -831,6 +837,7 @@ def get_optimized_kernel(self, result_dtype):
         for optimization in optimizations:
             knl = optimization(knl)
         knl = lp.add_inames_for_unused_hw_axes(knl)
+        knl = register_optimization_preambles(knl, self.device)
         return knl
 
     def __call__(self, queue, **kwargs):

sumpy/e2p.py

@@ -26,6 +26,7 @@
 import loopy as lp
 
 from sumpy.tools import KernelCacheMixin, gather_loopy_arguments
+from sumpy.codegen import register_optimization_preambles
 from loopy.version import MOST_RECENT_LANGUAGE_VERSION
 
 
@@ -198,6 +199,7 @@ def get_optimized_kernel(self):
         knl = lp.add_inames_to_insn(knl, "itgt_box", "id:kernel_scaling")
         knl = lp.set_options(knl,
                 enforce_variable_access_ordered="no_check")
+        knl = register_optimization_preambles(knl, self.device)
 
         return knl
 
@@ -324,6 +326,8 @@ def get_optimized_kernel(self):
         knl = lp.add_inames_to_insn(knl, "itgt_box", "id:kernel_scaling")
         knl = lp.set_options(knl,
                 enforce_variable_access_ordered="no_check")
+        knl = register_optimization_preambles(knl, self.device)
+
         return knl
 
     def __call__(self, queue, **kwargs):

sumpy/p2e.py

@@ -25,6 +25,7 @@
 from loopy.version import MOST_RECENT_LANGUAGE_VERSION
 
 from sumpy.tools import KernelCacheMixin, KernelComputation
+from sumpy.codegen import register_optimization_preambles
 
 import logging
 logger = logging.getLogger(__name__)
@@ -118,6 +119,7 @@ def get_optimized_kernel(self, sources_is_obj_array, centers_is_obj_array):
         knl = self._allow_redundant_execution_of_knl_scaling(knl)
         knl = lp.set_options(knl,
                 enforce_variable_access_ordered="no_check")
+        knl = register_optimization_preambles(knl, self.device)
         return knl
 
     def __call__(self, queue, **kwargs):

sumpy/p2p.py

@@ -29,6 +29,7 @@
 
 from sumpy.tools import (
         KernelComputation, KernelCacheMixin, is_obj_array_like)
+from sumpy.codegen import register_optimization_preambles
 
 
 __doc__ = """
@@ -190,7 +191,6 @@ def get_optimized_kernel(self, targets_is_obj_array, sources_is_obj_array):
         knl = lp.set_options(knl,
                 enforce_variable_access_ordered="no_check")
 
-        from sumpy.codegen import register_optimization_preambles
         knl = register_optimization_preambles(knl, self.device)
 
         return knl
@@ -411,6 +411,8 @@ def get_optimized_kernel(self, targets_is_obj_array, sources_is_obj_array):
         knl = self._allow_redundant_execution_of_knl_scaling(knl)
         knl = lp.set_options(knl,
                 enforce_variable_access_ordered="no_check")
+        knl = register_optimization_preambles(knl, self.device)
+
         return knl
 
     def __call__(self, queue, targets, sources, tgtindices, srcindices, **kwargs):
@@ -717,7 +719,6 @@ def get_optimized_kernel(self, max_nsources_in_one_box,
         knl = lp.set_options(knl,
                 enforce_variable_access_ordered="no_check")
 
-        from sumpy.codegen import register_optimization_preambles
         knl = register_optimization_preambles(knl, self.device)
 
         return knl