Merge pull request #222 from prabhuramachandran/fix-align-gpu-bug

BUG: Align was failing with strided props.

pysph/base/device_helper.py

@@ -62,6 +62,7 @@ def __init__(self, particle_array, backend=None):
         self._data = {}
         self.properties = []
         self.constants = []
+        self._strided_indices_knl = None
 
         for prop, ary in pa.properties.items():
             self.add_prop(prop, ary)
@@ -109,11 +110,37 @@ def align(self, indices):
         '''
         if not isinstance(indices, dict):
             indices = {1: indices}
+            self._make_strided_indices(indices)
+
         for prop in self.properties:
             stride = self._particle_array.stride.get(prop, 1)
             self._data[prop] = self._data[prop].align(indices.get(stride))
             setattr(self, prop, self._data[prop])
 
+    def _make_strided_indices(self, indices_dict):
+        '''Takes the indices in a dict assuming that the indices are for a
+        stride of 1 and makes suitable indices for other possible stride
+        values.
+        '''
+        indices = indices_dict[1]
+        n = indices.length
+        if not self._strided_indices_knl:
+            self._setup_strided_indices_kernel()
+        for stride in set(self._particle_array.stride.values()):
+            dest = array.empty(n*stride, dtype=np.int32, backend=self.backend)
+            self._strided_indices_knl(indices, dest, stride)
+            indices_dict[stride] = dest
+
+    def _setup_strided_indices_kernel(self):
+        @annotate(int='i, stride', gintp='indices, dest')
+        def set_indices(i, indices, dest, stride):
+            j = declare('int')
+            for j in range(stride):
+                dest[i*stride + j] = indices[i]*stride + j
+
+        knl = Elementwise(set_indices, backend=self.backend)
+        self._strided_indices_knl = knl
+
     def add_prop(self, name, carray):
         """Add a new property given the name and carray, note
         that this assumes that this property is already added to the
@@ -161,8 +188,7 @@ def max(self, arg):
     def update_minmax_cl(self, props, only_min=False, only_max=False):
         ary_list = [getattr(self, prop) for prop in props]
         array.update_minmax_gpu(ary_list, only_min=only_min,
-                               only_max=only_max, backend=self.backend)
-
+                                only_max=only_max, backend=self.backend)
 
     def update_min_max(self, props=None):
         """Update the min,max values of all properties """

pysph/base/tests/test_device_helper.py

@@ -234,12 +234,15 @@ def test_align(self, backend):
         self.setup()
         # Given
         pa = self.pa
+        pa.add_property('force', stride=3)
         h = DeviceHelper(pa, backend=backend)
 
         # When
         pa.set_device_helper(h)
-        h.resize(5)
+        n = 5
+        h.resize(n)
         h.x.set(np.array([2.0, 3.0, 4.0, 5.0, 6.0], h.x.dtype))
+        h.force.set(np.arange(n*3, dtype=h.force.dtype))
 
         indices = array.arange(4, -1, -1, dtype=np.int32,
                                backend=backend)
@@ -248,6 +251,10 @@ def test_align(self, backend):
 
         # Then
         assert np.all(h.x.get() == np.array([6., 5., 4., 3., 2.]))
+        x = np.arange(n*3)
+        x.shape = (n, 3)
+        expect = x[::-1, :].ravel()
+        assert np.all(h.force.get() == expect)
 
     @test_all_backends
     def test_align_particles(self, backend):