Skip to content

Commit

Permalink
Fix the issue with cpp codegen, where it currently cannot handle inputs
Browse files Browse the repository at this point in the history
like:
```c++
cpp.reshape_strides(Range([(0, 4, 1), (0, 5, 1)]), None, None, [2, 3, 5])
```
and crashes with an index error.
  • Loading branch information
pratyai committed Oct 19, 2024
1 parent 975a065 commit a3d1de5
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 0 deletions.
1 change: 1 addition & 0 deletions dace/codegen/targets/cpp.py
Original file line number Diff line number Diff line change
Expand Up @@ -417,6 +417,7 @@ def reshape_strides(subset, strides, original_strides, copy_shape):
dims = len(copy_shape)

reduced_tile_sizes = [ts for ts, s in zip(subset.tile_sizes, original_copy_shape) if s != 1]
reduced_tile_sizes += [1] * (dims - len(reduced_tile_sizes)) # Pad the remainder with 1s to maintain dimensions.

reshaped_copy = copy_shape + [ts for ts in subset.tile_sizes if ts != 1]
reshaped_copy[:len(copy_shape)] = [s / ts for s, ts in zip(copy_shape, reduced_tile_sizes)]
Expand Down
75 changes: 75 additions & 0 deletions tests/codegen/targets/cpp_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
import unittest
from functools import reduce
from operator import mul

from dace.codegen.targets import cpp
from dace.subsets import Range


class ReshapeStrides(unittest.TestCase):
def test_multidim_array_all_dims_unit(self):
r = Range([(0, 0, 1), (0, 0, 1)])

# To smaller-sized shape
target_dims = [1]
self.assertEqual(r.num_elements_exact(), reduce(mul, target_dims))
reshaped, strides = cpp.reshape_strides(r, None, None, target_dims)
self.assertEqual(reshaped, [1])

# To equal-sized shape
target_dims = [1, 1]
self.assertEqual(r.num_elements_exact(), reduce(mul, target_dims))
reshaped, strides = cpp.reshape_strides(r, None, None, target_dims)
self.assertEqual(reshaped, [1, 1])

# To larger-sized shape
target_dims = [1, 1, 1]
self.assertEqual(r.num_elements_exact(), reduce(mul, target_dims))
reshaped, strides = cpp.reshape_strides(r, None, None, target_dims)
self.assertEqual(reshaped, [1, 1, 1])

def test_multidim_array_some_dims_unit(self):
r = Range([(0, 1, 1), (0, 0, 1)])

# To smaller-sized shape
target_dims = [2]
self.assertEqual(r.num_elements_exact(), reduce(mul, target_dims))
reshaped, strides = cpp.reshape_strides(r, None, None, target_dims)
self.assertEqual(reshaped, target_dims)

# To equal-sized shape
target_dims = [2, 1]
self.assertEqual(r.num_elements_exact(), reduce(mul, target_dims))
reshaped, strides = cpp.reshape_strides(r, None, None, target_dims)
self.assertEqual(reshaped, target_dims)

# To equal-sized shape
target_dims = [2, 1, 1]
self.assertEqual(r.num_elements_exact(), reduce(mul, target_dims))
reshaped, strides = cpp.reshape_strides(r, None, None, target_dims)
self.assertEqual(reshaped, target_dims)

def test_multidim_array_different_shape(self):
r = Range([(0, 4, 1), (0, 5, 1)])

# To smaller-sized shape
target_dims = [30]
self.assertEqual(r.num_elements_exact(), reduce(mul, target_dims))
reshaped, strides = cpp.reshape_strides(r, None, None, target_dims)
self.assertEqual(reshaped, target_dims)

# To equal-sized shape
target_dims = [15, 2]
self.assertEqual(r.num_elements_exact(), reduce(mul, target_dims))
reshaped, strides = cpp.reshape_strides(r, None, None, target_dims)
self.assertEqual(reshaped, target_dims)

# To equal-sized shape
target_dims = [3, 5, 2]
self.assertEqual(r.num_elements_exact(), reduce(mul, target_dims))
reshaped, strides = cpp.reshape_strides(r, None, None, target_dims)
self.assertEqual(reshaped, target_dims)


if __name__ == '__main__':
unittest.main()

0 comments on commit a3d1de5

Please sign in to comment.