Skip to content

cost model mapper incompatible with pyfmmlib

In cost_model_compute_potential_insn, the first return value is wrangler.full_output_zeros, which in the case of pyfmmlib, is a numpy array.

However, in the normal compute mapper exec_compute_potential_insn, the first return value is always a CL array.

My proposed change:

diff --git a/pytential/qbx/__init__.py b/pytential/qbx/__init__.py
index fc92895..4292b18 100644
--- a/pytential/qbx/__init__.py
+++ b/pytential/qbx/__init__.py
@@ -581,7 +581,14 @@ class QBXLayerPotentialSource(LayerPotentialSourceBase):
             cost_model_result = self.cost_model(
                 geo_data, kernel, kernel_arguments, calibration_params
             )
-            return wrangler.full_output_zeros(), cost_model_result
+
+            from pytools.obj_array import with_object_array_or_scalar
+            output_placeholder = with_object_array_or_scalar(
+                wrangler.finalize_potentials,
+                wrangler.full_output_zeros()
+            )
+
+            return output_placeholder, cost_model_result
 
         return self._dispatch_compute_potential_insn(
             queue, insn, bound_expr, evaluate,
Edited by Hao Gao