Merge pull request #193 from adityapb/master

Use allocator for radix sort in ZOrderGPUNNPS
pypr · Mar 23, 2019 · ed701cc · ed701cc
2 parents 1e59cb3 + 7579f75
commit ed701cc
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 1 deletion.
diff --git a/pysph/base/gpu_nnps_base.pyx b/pysph/base/gpu_nnps_base.pyx
@@ -179,6 +179,7 @@ cdef class GPUNNPS(NNPSBase):
 
         self.backend = get_backend(backend)
         self.use_double = get_config().use_double
+        self.queue = get_queue()
         self.dtype = np.float64 if self.use_double else np.float32
         self.dtype_max = np.finfo(self.dtype).max
         self._last_domain_size = 0.0

diff --git a/pysph/base/z_order_gpu_nnps.pxd b/pysph/base/z_order_gpu_nnps.pxd
@@ -22,6 +22,7 @@ cdef class ZOrderGPUNNPS(GPUNNPS):
     cdef object overflow_cid_to_idx
     cdef object curr_cid
     cdef object max_cid_src
+    cdef object allocator
 
     cdef object helper
     cdef object radix_sort

diff --git a/pysph/base/z_order_gpu_nnps.pyx b/pysph/base/z_order_gpu_nnps.pyx
@@ -66,6 +66,10 @@ cdef class ZOrderGPUNNPS(GPUNNPS):
         self.cids = []
         self.cid_to_idx = []
 
+        self.allocator = cl.tools.MemoryPool(
+                cl.tools.ImmediateAllocator(self.queue)
+                )
+
         for i from 0<=i<self.narrays:
             pa_wrapper = <NNPSParticleArrayWrapper>self.pa_wrappers[i]
             num_particles = pa_wrapper.get_number_of_particles()
@@ -130,7 +134,7 @@ cdef class ZOrderGPUNNPS(GPUNNPS):
 
         (sorted_indices, sorted_keys), evnt = self.radix_sort(
             self.pids[pa_index].dev, self.pid_keys[pa_index].dev,
-            key_bits=max_num_bits
+            key_bits=max_num_bits, allocator=self.allocator
         )
         self.pids[pa_index].set_data(sorted_indices)
         self.pid_keys[pa_index].set_data(sorted_keys)