From 55d7cd89347c25bc6fb65110f427f68dab9f8c56 Mon Sep 17 00:00:00 2001
From: Isuru Fernando <isuruf@gmail.com>
Date: Thu, 6 Apr 2023 12:22:39 -0700
Subject: [PATCH] tag with vec only when the size is supported by pyopencl
---
sumpy/p2p.py | 4 +++-
1 file changed, 3 insertions(+), 1 deletion(-)
diff --git a/sumpy/p2p.py b/sumpy/p2p.py
index dc8368da..eba84537 100644
--- a/sumpy/p2p.py
+++ b/sumpy/p2p.py
@@ -666,7 +666,9 @@ class P2PFromCSR(P2PBase):
# optimization led to a 8% speedup in the performance.
knl = lp.concatenate_arrays(knl,
["local_isrc", "local_isrc_strength"], "local_isrc")
- knl = lp.tag_array_axes(knl, "local_isrc", "vec,C")
+ count = self.strength_count + self.dim
+ if count in [2, 3, 4, 8, 16]:
+ knl = lp.tag_array_axes(knl, "local_isrc", "vec,C")
knl = lp.add_inames_for_unused_hw_axes(knl)
# knl = lp.set_options(knl, write_code=True)
--
GitLab