Optimize M2P

Author: isuruf

sumpy/e2p.py

@@ -123,10 +123,9 @@ class E2PFromSingleBox(E2PBase):
     def default_name(self):
         return "e2p_from_single_box"
 
-    def get_kernel(self, max_ntargets_in_one_box):
+    def get_kernel(self, max_ntargets_in_one_box, max_work_items):
         ncoeffs = len(self.expansion)
         loopy_args = self.get_loopy_args()
-        max_work_items = min(256, max(ncoeffs, max_ntargets_in_one_box))
 
         loopy_knl = lp.make_kernel(
                 [
@@ -208,11 +207,16 @@ def get_kernel(self, max_ntargets_in_one_box):
         return loopy_knl
 
     def get_optimized_kernel(self, max_ntargets_in_one_box):
-        inner_knl, optimizations = self.get_loopy_evaluator_and_optimizations()
-        knl = self.get_kernel(max_ntargets_in_one_box=max_ntargets_in_one_box)
+        _, optimizations = self.get_loopy_evaluator_and_optimizations()
+
+        ncoeffs = len(self.expansion)
+        max_work_items = min(256, max(ncoeffs, max_ntargets_in_one_box))
+        knl = self.get_kernel(max_ntargets_in_one_box=max_ntargets_in_one_box,
+                              max_work_items=max_work_items)
+
         knl = lp.tag_inames(knl, {"itgt_box": "g.0"})
-        knl = lp.split_iname(knl, "itgt_offset", 256, inner_tag="l.0")
-        knl = lp.split_iname(knl, "icoeff", 256, inner_tag="l.0")
+        knl = lp.split_iname(knl, "itgt_offset", max_work_items, inner_tag="l.0")
+        knl = lp.split_iname(knl, "icoeff", max_work_items, inner_tag="l.0")
         knl = lp.add_inames_to_insn(knl, "dummy",
             "id:fetch_init* or id:fetch_center or id:kernel_scaling")
         knl = lp.add_inames_to_insn(knl, "itgt_box", "id:kernel_scaling")
@@ -273,10 +277,9 @@ class E2PFromCSR(E2PBase):
     def default_name(self):
         return "e2p_from_csr"
 
-    def get_kernel(self, max_ntargets_in_one_box):
+    def get_kernel(self, max_ntargets_in_one_box, max_work_items):
         ncoeffs = len(self.expansion)
         loopy_args = self.get_loopy_args()
-        max_work_items = min(256, max(ncoeffs, max_ntargets_in_one_box))
 
         loopy_knl = lp.make_kernel(
                 [
@@ -379,12 +382,16 @@ def get_kernel(self, max_ntargets_in_one_box):
 
     def get_optimized_kernel(self, max_ntargets_in_one_box):
         _, optimizations = self.get_loopy_evaluator_and_optimizations()
-        knl = self.get_kernel(max_ntargets_in_one_box=max_ntargets_in_one_box)
+        ncoeffs = len(self.expansion)
+        max_work_items = min(256, max(ncoeffs, max_ntargets_in_one_box))
+
+        knl = self.get_kernel(max_ntargets_in_one_box=max_ntargets_in_one_box,
+                              max_work_items=max_work_items)
         knl = lp.tag_inames(knl, {"itgt_box": "g.0", "dummy": "l.0"})
         knl = lp.unprivatize_temporaries_with_inames(knl,
             "itgt_offset", "result_temp")
-        knl = lp.split_iname(knl, "itgt_offset", 256, inner_tag="l.0")
-        knl = lp.split_iname(knl, "icoeff", 256, inner_tag="l.0")
+        knl = lp.split_iname(knl, "itgt_offset", max_work_items, inner_tag="l.0")
+        knl = lp.split_iname(knl, "icoeff", max_work_items, inner_tag="l.0")
         knl = lp.privatize_temporaries_with_inames(knl,
             "itgt_offset_outer", "result_temp")
         knl = lp.duplicate_inames(knl, "itgt_offset_outer", "id:init_result")