From afc7be29ad04b141ccfa40254ece5fdc024da65c Mon Sep 17 00:00:00 2001 From: Pratyai Mazumder Date: Wed, 23 Oct 2024 11:46:28 +0200 Subject: [PATCH 1/8] Add a volume function to `Range` and `Indices` to count the total number of elements they have. --- dace/subsets.py | 10 +++++++++ tests/subsets_test.py | 51 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+) create mode 100644 tests/subsets_test.py diff --git a/dace/subsets.py b/dace/subsets.py index e7b6869678..ba3fc2bfad 100644 --- a/dace/subsets.py +++ b/dace/subsets.py @@ -1,4 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. +import operator + import dace.serialize from dace import data, symbolic, dtypes import re @@ -334,6 +336,10 @@ def size_exact(self): for (iMin, iMax, step), ts in zip(self.ranges, self.tile_sizes) ] + def volume_exact(self) -> int: + """ Returns the total number of elements in all dimenssions together. """ + return reduce(operator.mul, self.size_exact()) + def bounding_box_size(self): """ Returns the size of a bounding box around this range. """ return [ @@ -895,6 +901,10 @@ def size(self): def size_exact(self): return self.size() + def volume_exact(self) -> int: + """ Returns the total number of elements in all dimenssions together. """ + return reduce(operator.mul, self.size_exact()) + def min_element(self): return self.indices diff --git a/tests/subsets_test.py b/tests/subsets_test.py new file mode 100644 index 0000000000..1002eea597 --- /dev/null +++ b/tests/subsets_test.py @@ -0,0 +1,51 @@ +import unittest + +from sympy import ceiling + +from dace.subsets import Range, Indices + +import dace + + +class Volume(unittest.TestCase): + def test_range(self): + K, N, M = dace.symbol('K', positive=True), dace.symbol('N', positive=True), dace.symbol('M', positive=True) + + # A regular cube. + r = Range([(0, K - 1, 1), (0, N - 1, 1), (0, M - 1, 1)]) + self.assertEqual(K * N * M, r.volume_exact()) + + # A regular cube with offsets. + r = Range([(1, 1 + K - 1, 1), (2, 2 + N - 1, 1), (3, 3 + M - 1, 1)]) + self.assertEqual(K * N * M, r.volume_exact()) + + # A regular cube with strides. + r = Range([(0, K - 1, 2), (0, N - 1, 3), (0, M - 1, 4)]) + self.assertEqual(ceiling(K / 2) * ceiling(N / 3) * ceiling(M / 4), r.volume_exact()) + + # A regular cube with both offsets and strides. + r = Range([(1, 1 + K - 1, 2), (2, 2 + N - 1, 3), (3, 3 + M - 1, 4)]) + self.assertEqual(ceiling(K / 2) * ceiling(N / 3) * ceiling(M / 4), r.volume_exact()) + + # A 2D square on 3D coordinate system. + r = Range([(1, 1 + K - 1, 2), (2, 2, 3), (3, 3 + M - 1, 4)]) + self.assertEqual(ceiling(K / 2) * ceiling(M / 4), r.volume_exact()) + + # A 3D point. + r = Range([(1, 1, 2), (2, 2, 3), (3, 3, 4)]) + self.assertEqual(1, r.volume_exact()) + + def test_indices(self): + # Indices always have volume 1 no matter what, since they are just points. + ind = Indices([0, 1, 2]) + self.assertEqual(1, ind.volume_exact()) + + ind = Indices([1]) + self.assertEqual(1, ind.volume_exact()) + + ind = Indices([0, 2]) + self.assertEqual(1, ind.volume_exact()) + + +if __name__ == '__main__': + unittest.main() From 112ee3bb092347039aa6719d8c3ffd816b216f8b Mon Sep 17 00:00:00 2001 From: Pratyai Mazumder Date: Wed, 23 Oct 2024 14:30:26 +0200 Subject: [PATCH 2/8] Add a `SubrangeMapper` helper class. Its functionality === Equipped with a `src` and a `dst` range of equal volumes but possibly different shapes or even dimensions (i.e., there is an 1-to-1 correspondence between the elements of `src` and `dst`), maps a subrange of `src` to its counterpart in `dst`, if possible. Note that such subrange-to-subrange mapping may not always exist. --- dace/subsets.py | 143 ++++++++++++++++++++++++++++++++++-------- tests/subsets_test.py | 75 +++++++++++++++++++++- 2 files changed, 188 insertions(+), 30 deletions(-) diff --git a/dace/subsets.py b/dace/subsets.py index ba3fc2bfad..ae0c9b1a11 100644 --- a/dace/subsets.py +++ b/dace/subsets.py @@ -1,14 +1,16 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import operator - -import dace.serialize -from dace import data, symbolic, dtypes import re -import sympy as sp +import warnings from functools import reduce -import sympy.core.sympify from typing import List, Optional, Sequence, Set, Union -import warnings + +import sympy as sp +import sympy.core.sympify +from sympy import ceiling + +import dace.serialize +from dace import symbolic from dace.config import Config @@ -22,6 +24,7 @@ def nng(expr): except AttributeError: # No free_symbols in expr return expr + def bounding_box_cover_exact(subset_a, subset_b) -> bool: min_elements_a = subset_a.min_element() max_elements_a = subset_a.max_element() @@ -31,8 +34,8 @@ def bounding_box_cover_exact(subset_a, subset_b) -> bool: # Covering only make sense if the two subsets have the same number of dimensions. if len(min_elements_a) != len(min_elements_b): return ValueError( - f"A bounding box of dimensionality {len(min_elements_a)} cannot" - f" test covering a bounding box of dimensionality {len(min_elements_b)}." + f"A bounding box of dimensionality {len(min_elements_a)} cannot" + f" test covering a bounding box of dimensionality {len(min_elements_b)}." ) return all([(symbolic.simplify_ext(nng(rb)) <= symbolic.simplify_ext(nng(orb))) == True @@ -40,7 +43,8 @@ def bounding_box_cover_exact(subset_a, subset_b) -> bool: for rb, re, orb, ore in zip(min_elements_a, max_elements_a, min_elements_b, max_elements_b)]) -def bounding_box_symbolic_positive(subset_a, subset_b, approximation = False)-> bool: + +def bounding_box_symbolic_positive(subset_a, subset_b, approximation=False) -> bool: min_elements_a = subset_a.min_element_approx() if approximation else subset_a.min_element() max_elements_a = subset_a.max_element_approx() if approximation else subset_a.max_element() min_elements_b = subset_b.min_element_approx() if approximation else subset_b.min_element() @@ -49,8 +53,8 @@ def bounding_box_symbolic_positive(subset_a, subset_b, approximation = False)-> # Covering only make sense if the two subsets have the same number of dimensions. if len(min_elements_a) != len(min_elements_b): return ValueError( - f"A bounding box of dimensionality {len(min_elements_a)} cannot" - f" test covering a bounding box of dimensionality {len(min_elements_b)}." + f"A bounding box of dimensionality {len(min_elements_a)} cannot" + f" test covering a bounding box of dimensionality {len(min_elements_b)}." ) for rb, re, orb, ore in zip(min_elements_a, max_elements_a, @@ -72,6 +76,7 @@ def bounding_box_symbolic_positive(subset_a, subset_b, approximation = False)-> return False return True + class Subset(object): """ Defines a subset of a data descriptor. """ @@ -82,7 +87,7 @@ def covers(self, other): # Subsets of different dimensionality can never cover each other. if self.dims() != other.dims(): return ValueError( - f"A subset of dimensionality {self.dim()} cannot test covering a subset of dimensionality {other.dims()}" + f"A subset of dimensionality {self.dim()} cannot test covering a subset of dimensionality {other.dims()}" ) if not Config.get('optimizer', 'symbolic_positive'): @@ -101,20 +106,22 @@ def covers(self, other): return False return True - + def covers_precise(self, other): """ Returns True if self contains all the elements in other. """ # Subsets of different dimensionality can never cover each other. if self.dims() != other.dims(): return ValueError( - f"A subset of dimensionality {self.dim()} cannot test covering a subset of dimensionality {other.dims()}" + f"A subset of dimensionality {self.dim()} cannot test covering a subset of dimensionality {other.dims()}" ) # If self does not cover other with a bounding box union, return false. symbolic_positive = Config.get('optimizer', 'symbolic_positive') try: - bounding_box_cover = bounding_box_cover_exact(self, other) if symbolic_positive else bounding_box_symbolic_positive(self, other) + bounding_box_cover = bounding_box_cover_exact(self, + other) if symbolic_positive else bounding_box_symbolic_positive( + self, other) if not bounding_box_cover: return False except TypeError: @@ -153,14 +160,13 @@ def covers_precise(self, other): except: return False return True - # unknown type + # unknown type else: raise TypeError except TypeError: return False - def __repr__(self): return '%s (%s)' % (type(self).__name__, self.__str__()) @@ -231,6 +237,7 @@ def _tuple_to_symexpr(val): @dace.serialize.serializable class Range(Subset): """ Subset defined in terms of a fixed range. """ + def __init__(self, ranges): parsed_ranges = [] parsed_tiles = [] @@ -584,7 +591,7 @@ def from_string(string): value = symbolic.pystr_to_symbolic(uni_dim_tokens[0].strip()) ranges.append((value, value, 1)) continue - #return Range(ranges) + # return Range(ranges) # If dimension has more than 4 tokens, the range is invalid if len(uni_dim_tokens) > 4: raise SyntaxError("Invalid range: {}".format(multi_dim_tokens)) @@ -854,6 +861,7 @@ def intersects(self, other: 'Range'): class Indices(Subset): """ A subset of one element representing a single index in an N-dimensional data descriptor. """ + def __init__(self, indices): if indices is None or len(indices) == 0: raise TypeError('Expected an array of index expressions: got empty' ' array or None') @@ -880,7 +888,7 @@ def from_json(obj, context=None): raise TypeError("from_json of class \"Indices\" called on json " "with type %s (expected 'Indices')" % obj['type']) - #return Indices(symbolic.SymExpr(obj['indices'])) + # return Indices(symbolic.SymExpr(obj['indices'])) return Indices([*map(symbolic.pystr_to_symbolic, obj['indices'])]) def __hash__(self): @@ -1091,6 +1099,7 @@ def intersection(self, other: 'Indices'): return self return None + class SubsetUnion(Subset): """ Wrapper subset type that stores multiple Subsets in a list. @@ -1128,7 +1137,7 @@ def covers(self, other): return False else: return any(s.covers(other) for s in self.subset_list) - + def covers_precise(self, other): """ Returns True if this SubsetUnion covers another @@ -1154,7 +1163,7 @@ def __str__(self): string += " " string += subset.__str__() return string - + def dims(self): if not self.subset_list: return 0 @@ -1178,7 +1187,7 @@ def free_symbols(self) -> Set[str]: for subset in self.subset_list: result |= subset.free_symbols return result - + def replace(self, repl_dict): for subset in self.subset_list: subset.replace(repl_dict) @@ -1188,13 +1197,12 @@ def num_elements(self): min = 0 for subset in self.subset_list: try: - if subset.num_elements() < min or min ==0: + if subset.num_elements() < min or min == 0: min = subset.num_elements() except: continue - - return min + return min def _union_special_cases(arb: symbolic.SymbolicType, brb: symbolic.SymbolicType, are: symbolic.SymbolicType, @@ -1261,8 +1269,6 @@ def bounding_box_union(subset_a: Subset, subset_b: Subset) -> Range: return Range(result) - - def union(subset_a: Subset, subset_b: Subset) -> Subset: """ Compute the union of two Subset objects. If the subsets are not of the same type, degenerates to bounding-box @@ -1331,6 +1337,7 @@ def list_union(subset_a: Subset, subset_b: Subset) -> Subset: except TypeError: return None + def intersects(subset_a: Subset, subset_b: Subset) -> Union[bool, None]: """ Returns True if two subsets intersect, False if they do not, or @@ -1352,3 +1359,85 @@ def intersects(subset_a: Subset, subset_b: Subset) -> Union[bool, None]: return None except TypeError: # cannot determine truth value of Relational return None + + +class SubrangeMapper: + """ + Equipped with a `src` and a `dst` range of equal volumes but possibly different shapes or even dimensions (i.e., + there is an 1-to-1 correspondence between the elements of `src` and `dst`), maps a subrange of `src` to its + counterpart in `dst`, if possible. + + Note that such subrange-to-subrange mapping may not always exist. + """ + + def __init__(self, src: Range, dst: Range): + src, dst = self.canonical(src), self.canonical(dst) + assert src.volume_exact() == dst.volume_exact() + self.src, self.dst = src, dst + + @staticmethod + def canonical(r: Range) -> Range: + """ + Extends the (excluded) upper bound of each component of the ranges as much as possible, without affecting the + volume of the range. + """ + return Range([(b, b + s * ceiling((e - b + 1) / s) - 1, s) + for b, e, s in r.ndrange()]) + + def map(self, r: Range) -> Optional[Range]: + r = self.canonical(r) + # Ideally we also have `assert self.src.covers_precise(r)`. However, we cannot determine that for symbols. + assert self.src.dims() == r.dims() + out = [] + src_i, dst_i = 0, 0 + while src_i < self.src.dims(): + assert dst_i < self.dst.dims() + + src_j, dst_j = None, None + for sj in range(src_i + 1, self.src.dims() + 1): + for dj in range(dst_i + 1, self.dst.dims() + 1): + if Range(self.src.ranges[src_i:sj]).volume_exact() == Range( + self.dst.ranges[dst_i:dj]).volume_exact(): + src_j, dst_j = sj, dj + break + else: + continue + break + if src_j is None: + return None + + if Range(r.ranges[src_i: src_j]).volume_exact() == 1: + # If we are selecting just a single point in this segment, we can just pick the mapping of that point. + src_segment, dst_segment = Range(self.src.ranges[src_i: src_j]), Range(self.dst.ranges[dst_i: dst_j]) + # Compute the local 1D coordinate of the point on `src`. + loc = 0 + for (idx, _, _), (ridx, _, _), s in zip(reversed(src_segment.ranges), + reversed(r.ranges[src_i: src_j]), + reversed(src_segment.size())): + loc = loc * s + (ridx - idx) + # Translate that local 1D coordinate onto `dst`. + dst_coord = [] + for (idx, _, _), s in zip(dst_segment.ranges, dst_segment.size()): + dst_coord.append(loc % s + idx) + loc = loc // s + out.extend([(idx, idx, 1) for idx in dst_coord]) + elif self.src.ranges[src_i: src_j] == r.ranges[src_i: src_j]: + # If we are selecting the entirety of this segment, we can just pick the corresponding mapped segment in + # its entirety too. + out.extend(self.dst.ranges[dst_i:dst_j]) + elif src_j - src_i == 1 and dst_j - dst_i == 1: + # If the segment lengths on both sides are just 1, the mapping is easy to compute. + sb, se, ss = self.src.ranges[src_i] + db, de, ds = self.dst.ranges[dst_i] + b, e, s = r.ranges[src_i] + lb, le, ls = (b - sb) // ss, (e - se) // ss - 1, s // ss + tb, te, ts = db + lb * ds, de + (le + 1) * ds, ds * ls + out.append((tb, te, ts)) + else: + # TODO: Can we narrow down this case even more? That would be number theoretic problem. + # E.g., If we are reshaping [6, 5] to [2, 15], we are demanding that these dimensions must be wholly + # selected for now. + return None + + src_i, dst_i = src_j, dst_j + return Range(out) diff --git a/tests/subsets_test.py b/tests/subsets_test.py index 1002eea597..65a2f19a98 100644 --- a/tests/subsets_test.py +++ b/tests/subsets_test.py @@ -1,13 +1,19 @@ import unittest +from typing import Dict +import numpy as np from sympy import ceiling -from dace.subsets import Range, Indices - import dace +from dace.subsets import Range, Indices, SubrangeMapper +from dace.symbolic import simplify + + +def eval_range(r: Range, vals: Dict): + return Range([(simplify(b).subs(vals), simplify(e).subs(vals), simplify(s).subs(vals)) for b, e, s in r.ranges]) -class Volume(unittest.TestCase): +class VolumeTest(unittest.TestCase): def test_range(self): K, N, M = dace.symbol('K', positive=True), dace.symbol('N', positive=True), dace.symbol('M', positive=True) @@ -47,5 +53,68 @@ def test_indices(self): self.assertEqual(1, ind.volume_exact()) +class RangeRemapperTest(unittest.TestCase): + def test_mapping_without_symbols(self): + K, N, M = 6, 7, 8 + + # A regular cube. + src = Range([(0, K - 1, 1), (0, N - 1, 1), (0, M - 1, 1)]) + # A regular cube with offsets. + dst = Range([(1, 1 + K - 1, 1), (2, 2 + N - 1, 1), (3, 3 + M - 1, 1)]) + # A Mapper + sm = SubrangeMapper(src, dst) + + # Pick the entire range. + self.assertEqual(dst, sm.map(src)) + # Pick a point 0, 0, 0. + self.assertEqual(Range([(1, 1, 1), (2, 2, 1), (3, 3, 1)]), + sm.map(Range([(0, 0, 1), (0, 0, 1), (0, 0, 1)]))) + # Pick a point K//2, N//2, M//2. + self.assertEqual(Range([(1 + K // 2, 1 + K // 2, 1), (2 + N // 2, 2 + N // 2, 1), (3 + M // 2, 3 + M // 2, 1)]), + sm.map(Range([(K // 2, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1)]))) + # Pick a point K-1, N-1, M-1. + self.assertEqual(Range([(1 + K - 1, 1 + K - 1, 1), (2 + N - 1, 2 + N - 1, 1), (3 + M - 1, 3 + M - 1, 1)]), + sm.map(Range([(K - 1, K - 1, 1), (N - 1, N - 1, 1), (M - 1, M - 1, 1)]))) + # Pick a quadrant. + self.assertEqual(Range([(1, 1 + K // 2, 1), (2, 2 + N // 2, 1), (3, 3 + M // 2, 1)]), + sm.map(Range([(0, K // 2, 1), (0, N // 2, 1), (0, M // 2, 1)]))) + + def test_mapping_with_only_offsets(self): + K, N, M = dace.symbol('K', positive=True), dace.symbol('N', positive=True), dace.symbol('M', positive=True) + + # A regular cube. + src = Range([(0, K - 1, 1), (0, N - 1, 1), (0, M - 1, 1)]) + # A regular cube with offsets. + dst = Range([(1, 1 + K - 1, 1), (2, 2 + N - 1, 1), (3, 3 + M - 1, 1)]) + # A Mapper + sm = SubrangeMapper(src, dst) + + # Pick the entire range. + self.assertEqual(dst, sm.map(src)) + + # NOTE: I couldn't make SymPy understand that `(K//2) % K == (K//2)` always holds for postive integers `K`. + # Hence, the numerical approach. + argslist = [{'K': k, 'N': n, 'M': m} for k, n, m in zip(np.random.randint(1, 100, size=20), + np.random.randint(1, 100, size=20), + np.random.randint(1, 100, size=20))] + for args in argslist: + # Pick a point K//2, N//2, M//2. + want = eval_range( + Range([(1 + K // 2, 1 + K // 2, 1), (2 + N // 2, 2 + N // 2, 1), (3 + M // 2, 3 + M // 2, 1)]), + args) + got = eval_range( + sm.map(Range([(K // 2, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1)])), + args) + self.assertEqual(want, got) + # Pick a quadrant. + want = eval_range( + Range([(1, 1 + K // 2, 1), (2, 2 + N // 2, 1), (3, 3 + M // 2, 1)]), + args) + got = eval_range( + sm.map(Range([(0, K // 2, 1), (0, N // 2, 1), (0, M // 2, 1)])), + args) + self.assertEqual(want, got) + + if __name__ == '__main__': unittest.main() From f450c8ef59764816baa335aa8352e412487b160c Mon Sep 17 00:00:00 2001 From: Pratyai Mazumder Date: Wed, 23 Oct 2024 15:08:26 +0200 Subject: [PATCH 3/8] Cover more cases with tests. --- dace/subsets.py | 7 ++++--- tests/subsets_test.py | 43 ++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/dace/subsets.py b/dace/subsets.py index ae0c9b1a11..2f1bfb0d1b 100644 --- a/dace/subsets.py +++ b/dace/subsets.py @@ -1406,9 +1406,10 @@ def map(self, r: Range) -> Optional[Range]: if src_j is None: return None - if Range(r.ranges[src_i: src_j]).volume_exact() == 1: - # If we are selecting just a single point in this segment, we can just pick the mapping of that point. - src_segment, dst_segment = Range(self.src.ranges[src_i: src_j]), Range(self.dst.ranges[dst_i: dst_j]) + # If we are selecting just a single point in this segment, we can just pick the mapping of that point. + src_segment, dst_segment, r_segment = Range(self.src.ranges[src_i: src_j]), Range( + self.dst.ranges[dst_i: dst_j]), Range(r.ranges[src_i: src_j]) + if r_segment.volume_exact() == 1: # Compute the local 1D coordinate of the point on `src`. loc = 0 for (idx, _, _), (ridx, _, _), s in zip(reversed(src_segment.ranges), diff --git a/tests/subsets_test.py b/tests/subsets_test.py index 65a2f19a98..4af4449c72 100644 --- a/tests/subsets_test.py +++ b/tests/subsets_test.py @@ -79,7 +79,7 @@ def test_mapping_without_symbols(self): self.assertEqual(Range([(1, 1 + K // 2, 1), (2, 2 + N // 2, 1), (3, 3 + M // 2, 1)]), sm.map(Range([(0, K // 2, 1), (0, N // 2, 1), (0, M // 2, 1)]))) - def test_mapping_with_only_offsets(self): + def test_mapping_with_symbols(self): K, N, M = dace.symbol('K', positive=True), dace.symbol('N', positive=True), dace.symbol('M', positive=True) # A regular cube. @@ -115,6 +115,47 @@ def test_mapping_with_only_offsets(self): args) self.assertEqual(want, got) + def test_mapping_with_reshaping(self): + K, N, M = dace.symbol('K', positive=True), dace.symbol('N', positive=True), dace.symbol('M', positive=True) + + # A regular cube. + src = Range([(0, K - 1, 1), (0, N - 1, 1), (0, M - 1, 1)]) + # A regular cube with different shape. + dst = Range([(0, K - 1, 1), (0, N * M - 1, 1)]) + # A Mapper + sm = SubrangeMapper(src, dst) + + # Pick the entire range. + self.assertEqual(dst, sm.map(src)) + + # NOTE: I couldn't make SymPy understand that `(K//2) % K == (K//2)` always holds for postive integers `K`. + # Hence, the numerical approach. + argslist = [{'K': k, 'N': n, 'M': m} for k, n, m in zip(np.random.randint(1, 10, size=20), + np.random.randint(1, 10, size=20), + np.random.randint(1, 10, size=20))] + # Pick a point K//2, N//2, M//2. + for args in argslist: + want = eval_range( + Range([(K // 2, K // 2, 1), ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1)]), + args) + got = eval_range( + sm.map(Range([(K // 2, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1)])), + args) + self.assertEqual(want, got) + # Pick a quadrant. + for args in argslist: + # But its mapping cannot be expressed as a simple range with offset and stride. + self.assertIsNone(sm.map(Range([(0, K // 2, 1), (0, N // 2, 1), (0, M // 2, 1)]))) + # Pick only points in problematic quadrants, but larger subsets elsewhere. + for args in argslist: + want = eval_range( + Range([(0, K // 2, 1), ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1)]), + args) + got = eval_range( + sm.map(Range([(0, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1)])), + args) + self.assertEqual(want, got) + if __name__ == '__main__': unittest.main() From 01b9e7073e12f188082256815bb78990b32c4a93 Mon Sep 17 00:00:00 2001 From: Pratyai Mazumder Date: Wed, 23 Oct 2024 15:24:04 +0200 Subject: [PATCH 4/8] Add some explanatory comments. --- dace/subsets.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/dace/subsets.py b/dace/subsets.py index 2f1bfb0d1b..aa3a3269e7 100644 --- a/dace/subsets.py +++ b/dace/subsets.py @@ -1393,6 +1393,9 @@ def map(self, r: Range) -> Optional[Range]: while src_i < self.src.dims(): assert dst_i < self.dst.dims() + # Find the next smallest segments of `src` and `dst` whose volumes matches (and therefore can possibly have + # a mapping). + # TODO: It's possible to do this in a O(max(|src|, |dst|)) loop instead of O(|src| * |dst|). src_j, dst_j = None, None for sj in range(src_i + 1, self.src.dims() + 1): for dj in range(dst_i + 1, self.dst.dims() + 1): @@ -1404,12 +1407,14 @@ def map(self, r: Range) -> Optional[Range]: continue break if src_j is None: + # Somehow, we couldn't find a matching segment. This should have been caught earlier. return None - # If we are selecting just a single point in this segment, we can just pick the mapping of that point. - src_segment, dst_segment, r_segment = Range(self.src.ranges[src_i: src_j]), Range( - self.dst.ranges[dst_i: dst_j]), Range(r.ranges[src_i: src_j]) + src_segment = Range(self.src.ranges[src_i: src_j]) + dst_segment = Range(self.dst.ranges[dst_i: dst_j]) + r_segment = Range(r.ranges[src_i: src_j]) if r_segment.volume_exact() == 1: + # If we are selecting just a single point in this segment, we can just pick the mapping of that point. # Compute the local 1D coordinate of the point on `src`. loc = 0 for (idx, _, _), (ridx, _, _), s in zip(reversed(src_segment.ranges), @@ -1427,7 +1432,7 @@ def map(self, r: Range) -> Optional[Range]: # its entirety too. out.extend(self.dst.ranges[dst_i:dst_j]) elif src_j - src_i == 1 and dst_j - dst_i == 1: - # If the segment lengths on both sides are just 1, the mapping is easy to compute. + # If the segment lengths on both sides are just 1, the mapping is easy to compute -- it's just a shift. sb, se, ss = self.src.ranges[src_i] db, de, ds = self.dst.ranges[dst_i] b, e, s = r.ranges[src_i] From 3194cc40ab16cfbcdd940f86f79642affcba7084 Mon Sep 17 00:00:00 2001 From: Pratyai Mazumder Date: Thu, 24 Oct 2024 17:58:08 +0200 Subject: [PATCH 5/8] Missed the case of if we run of list only one side. Handle that, add more tests to cover everything. --- dace/subsets.py | 12 ++++++-- tests/subsets_test.py | 65 +++++++++++++++++++++++++++++++++++-------- 2 files changed, 63 insertions(+), 14 deletions(-) diff --git a/dace/subsets.py b/dace/subsets.py index aa3a3269e7..be8946c945 100644 --- a/dace/subsets.py +++ b/dace/subsets.py @@ -1390,8 +1390,8 @@ def map(self, r: Range) -> Optional[Range]: assert self.src.dims() == r.dims() out = [] src_i, dst_i = 0, 0 - while src_i < self.src.dims(): - assert dst_i < self.dst.dims() + while src_i < self.src.dims() and dst_i < self.dst.dims(): + # If we run out only on one side, handle that case after the loop. # Find the next smallest segments of `src` and `dst` whose volumes matches (and therefore can possibly have # a mapping). @@ -1446,4 +1446,12 @@ def map(self, r: Range) -> Optional[Range]: return None src_i, dst_i = src_j, dst_j + if src_i < self.src.dims(): + src_segment = Range(self.src.ranges[src_i: self.src.dims()]) + assert src_segment.volume_exact() == 1 + if dst_i < self.dst.dims(): + # Take the remaining dst segment which must have a volume of 1 by now. + dst_segment = Range(self.dst.ranges[dst_i: self.dst.dims()]) + assert dst_segment.volume_exact() == 1 + out.extend(dst_segment.ranges) return Range(out) diff --git a/tests/subsets_test.py b/tests/subsets_test.py index 4af4449c72..cc8a304400 100644 --- a/tests/subsets_test.py +++ b/tests/subsets_test.py @@ -124,9 +124,11 @@ def test_mapping_with_reshaping(self): dst = Range([(0, K - 1, 1), (0, N * M - 1, 1)]) # A Mapper sm = SubrangeMapper(src, dst) + sm_inv = SubrangeMapper(dst, src) # Pick the entire range. self.assertEqual(dst, sm.map(src)) + self.assertEqual(src, sm_inv.map(dst)) # NOTE: I couldn't make SymPy understand that `(K//2) % K == (K//2)` always holds for postive integers `K`. # Hence, the numerical approach. @@ -135,12 +137,11 @@ def test_mapping_with_reshaping(self): np.random.randint(1, 10, size=20))] # Pick a point K//2, N//2, M//2. for args in argslist: - want = eval_range( - Range([(K // 2, K // 2, 1), ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1)]), - args) - got = eval_range( - sm.map(Range([(K // 2, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1)])), - args) + orig = Range([(K // 2, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1)]) + orig_maps_to = Range([(K // 2, K // 2, 1), ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1)]) + want, got = eval_range(orig_maps_to, args), eval_range(sm.map(orig), args) + self.assertEqual(want, got) + want, got = eval_range(orig, args), eval_range(sm_inv.map(orig_maps_to), args) self.assertEqual(want, got) # Pick a quadrant. for args in argslist: @@ -148,12 +149,52 @@ def test_mapping_with_reshaping(self): self.assertIsNone(sm.map(Range([(0, K // 2, 1), (0, N // 2, 1), (0, M // 2, 1)]))) # Pick only points in problematic quadrants, but larger subsets elsewhere. for args in argslist: - want = eval_range( - Range([(0, K // 2, 1), ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1)]), - args) - got = eval_range( - sm.map(Range([(0, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1)])), - args) + orig = Range([(0, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1)]) + orig_maps_to = Range([(0, K // 2, 1), ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1)]) + want, got = eval_range(orig_maps_to, args), eval_range(sm.map(orig), args) + self.assertEqual(want, got) + want, got = eval_range(orig, args), eval_range(sm_inv.map(orig_maps_to), args) + self.assertEqual(want, got) + + def test_mapping_with_reshaping_unit_dims(self): + K, N, M = dace.symbol('K', positive=True), dace.symbol('N', positive=True), dace.symbol('M', positive=True) + + # A regular cube. + src = Range([(0, K - 1, 1), (0, N - 1, 1), (0, M - 1, 1), (0, 0, 1)]) + # A regular cube with different shape. + dst = Range([(0, K - 1, 1), (0, N * M - 1, 1), (0, 0, 1), (0, 0, 1)]) + # A Mapper + sm = SubrangeMapper(src, dst) + sm_inv = SubrangeMapper(dst, src) + + # Pick the entire range. + self.assertEqual(dst, sm.map(src)) + self.assertEqual(src, sm_inv.map(dst)) + + # NOTE: I couldn't make SymPy understand that `(K//2) % K == (K//2)` always holds for postive integers `K`. + # Hence, the numerical approach. + argslist = [{'K': k, 'N': n, 'M': m} for k, n, m in zip(np.random.randint(1, 10, size=20), + np.random.randint(1, 10, size=20), + np.random.randint(1, 10, size=20))] + # Pick a point K//2, N//2, M//2. + for args in argslist: + orig = Range([(K // 2, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1), (0, 0, 1)]) + orig_maps_to = Range([(K // 2, K // 2, 1), ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1), (0, 0, 1), (0, 0, 1)]) + want, got = eval_range(orig_maps_to, args), eval_range(sm.map(orig), args) + self.assertEqual(want, got) + want, got = eval_range(orig, args), eval_range(sm_inv.map(orig_maps_to), args) + self.assertEqual(want, got) + # Pick a quadrant. + for args in argslist: + # But its mapping cannot be expressed as a simple range with offset and stride. + self.assertIsNone(sm.map(Range([(0, K // 2, 1), (0, N // 2, 1), (0, M // 2, 1), (0, 0, 1)]))) + # Pick only points in problematic quadrants, but larger subsets elsewhere. + for args in argslist: + orig = Range([(0, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1), (0, 0, 1)]) + orig_maps_to = Range([(0, K // 2, 1), ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1), (0, 0, 1), (0, 0, 1)]) + want, got = eval_range(orig_maps_to, args), eval_range(sm.map(orig), args) + self.assertEqual(want, got) + want, got = eval_range(orig, args), eval_range(sm_inv.map(orig_maps_to), args) self.assertEqual(want, got) From 0e1cdf54f61576eca27a30f88e25c834b3f29aa0 Mon Sep 17 00:00:00 2001 From: Pratyai Mazumder Date: Fri, 25 Oct 2024 15:09:42 +0200 Subject: [PATCH 6/8] Add unit dimensions even in the middle. And a typo fix. And replace `operator.mul` with a lambda function. --- dace/subsets.py | 9 ++++----- tests/subsets_test.py | 12 +++++++++--- 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/dace/subsets.py b/dace/subsets.py index be8946c945..4f4c8905b8 100644 --- a/dace/subsets.py +++ b/dace/subsets.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import operator import re import warnings from functools import reduce @@ -344,8 +343,8 @@ def size_exact(self): ] def volume_exact(self) -> int: - """ Returns the total number of elements in all dimenssions together. """ - return reduce(operator.mul, self.size_exact()) + """ Returns the total number of elements in all dimensions together. """ + return reduce(lambda a, b: a * b, self.size_exact()) def bounding_box_size(self): """ Returns the size of a bounding box around this range. """ @@ -910,8 +909,8 @@ def size_exact(self): return self.size() def volume_exact(self) -> int: - """ Returns the total number of elements in all dimenssions together. """ - return reduce(operator.mul, self.size_exact()) + """ Returns the total number of elements in all dimensions together. """ + return reduce(lambda a, b: a * b, self.size_exact()) def min_element(self): return self.indices diff --git a/tests/subsets_test.py b/tests/subsets_test.py index cc8a304400..cbdcfca9f3 100644 --- a/tests/subsets_test.py +++ b/tests/subsets_test.py @@ -162,7 +162,7 @@ def test_mapping_with_reshaping_unit_dims(self): # A regular cube. src = Range([(0, K - 1, 1), (0, N - 1, 1), (0, M - 1, 1), (0, 0, 1)]) # A regular cube with different shape. - dst = Range([(0, K - 1, 1), (0, N * M - 1, 1), (0, 0, 1), (0, 0, 1)]) + dst = Range([(0, K - 1, 1), (0, 0, 1), (0, N * M - 1, 1), (0, 0, 1), (0, 0, 1)]) # A Mapper sm = SubrangeMapper(src, dst) sm_inv = SubrangeMapper(dst, src) @@ -179,7 +179,10 @@ def test_mapping_with_reshaping_unit_dims(self): # Pick a point K//2, N//2, M//2. for args in argslist: orig = Range([(K // 2, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1), (0, 0, 1)]) - orig_maps_to = Range([(K // 2, K // 2, 1), ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1), (0, 0, 1), (0, 0, 1)]) + orig_maps_to = Range([(K // 2, K // 2, 1), + (0, 0, 1), + ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1), + (0, 0, 1), (0, 0, 1)]) want, got = eval_range(orig_maps_to, args), eval_range(sm.map(orig), args) self.assertEqual(want, got) want, got = eval_range(orig, args), eval_range(sm_inv.map(orig_maps_to), args) @@ -191,7 +194,10 @@ def test_mapping_with_reshaping_unit_dims(self): # Pick only points in problematic quadrants, but larger subsets elsewhere. for args in argslist: orig = Range([(0, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1), (0, 0, 1)]) - orig_maps_to = Range([(0, K // 2, 1), ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1), (0, 0, 1), (0, 0, 1)]) + orig_maps_to = Range([(0, K // 2, 1), + (0, 0, 1), + ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1), + (0, 0, 1), (0, 0, 1)]) want, got = eval_range(orig_maps_to, args), eval_range(sm.map(orig), args) self.assertEqual(want, got) want, got = eval_range(orig, args), eval_range(sm_inv.map(orig_maps_to), args) From 6f77900c5fea17a18ac498f1ab81a20647a4acb4 Mon Sep 17 00:00:00 2001 From: Pratyai Mazumder Date: Fri, 25 Oct 2024 15:19:56 +0200 Subject: [PATCH 7/8] Stick to `pytest`. --- tests/subsets_test.py | 381 +++++++++++++++++++++--------------------- 1 file changed, 193 insertions(+), 188 deletions(-) diff --git a/tests/subsets_test.py b/tests/subsets_test.py index cbdcfca9f3..735a971068 100644 --- a/tests/subsets_test.py +++ b/tests/subsets_test.py @@ -13,196 +13,201 @@ def eval_range(r: Range, vals: Dict): return Range([(simplify(b).subs(vals), simplify(e).subs(vals), simplify(s).subs(vals)) for b, e, s in r.ranges]) -class VolumeTest(unittest.TestCase): - def test_range(self): - K, N, M = dace.symbol('K', positive=True), dace.symbol('N', positive=True), dace.symbol('M', positive=True) - - # A regular cube. - r = Range([(0, K - 1, 1), (0, N - 1, 1), (0, M - 1, 1)]) - self.assertEqual(K * N * M, r.volume_exact()) - - # A regular cube with offsets. - r = Range([(1, 1 + K - 1, 1), (2, 2 + N - 1, 1), (3, 3 + M - 1, 1)]) - self.assertEqual(K * N * M, r.volume_exact()) - - # A regular cube with strides. - r = Range([(0, K - 1, 2), (0, N - 1, 3), (0, M - 1, 4)]) - self.assertEqual(ceiling(K / 2) * ceiling(N / 3) * ceiling(M / 4), r.volume_exact()) - - # A regular cube with both offsets and strides. - r = Range([(1, 1 + K - 1, 2), (2, 2 + N - 1, 3), (3, 3 + M - 1, 4)]) - self.assertEqual(ceiling(K / 2) * ceiling(N / 3) * ceiling(M / 4), r.volume_exact()) - - # A 2D square on 3D coordinate system. - r = Range([(1, 1 + K - 1, 2), (2, 2, 3), (3, 3 + M - 1, 4)]) - self.assertEqual(ceiling(K / 2) * ceiling(M / 4), r.volume_exact()) - - # A 3D point. - r = Range([(1, 1, 2), (2, 2, 3), (3, 3, 4)]) - self.assertEqual(1, r.volume_exact()) - - def test_indices(self): - # Indices always have volume 1 no matter what, since they are just points. - ind = Indices([0, 1, 2]) - self.assertEqual(1, ind.volume_exact()) - - ind = Indices([1]) - self.assertEqual(1, ind.volume_exact()) - - ind = Indices([0, 2]) - self.assertEqual(1, ind.volume_exact()) - - -class RangeRemapperTest(unittest.TestCase): - def test_mapping_without_symbols(self): - K, N, M = 6, 7, 8 - - # A regular cube. - src = Range([(0, K - 1, 1), (0, N - 1, 1), (0, M - 1, 1)]) - # A regular cube with offsets. - dst = Range([(1, 1 + K - 1, 1), (2, 2 + N - 1, 1), (3, 3 + M - 1, 1)]) - # A Mapper - sm = SubrangeMapper(src, dst) - - # Pick the entire range. - self.assertEqual(dst, sm.map(src)) - # Pick a point 0, 0, 0. - self.assertEqual(Range([(1, 1, 1), (2, 2, 1), (3, 3, 1)]), - sm.map(Range([(0, 0, 1), (0, 0, 1), (0, 0, 1)]))) +def test_volume_of_range(): + K, N, M = dace.symbol('K', positive=True), dace.symbol('N', positive=True), dace.symbol('M', positive=True) + + # A regular cube. + r = Range([(0, K - 1, 1), (0, N - 1, 1), (0, M - 1, 1)]) + assert K * N * M == r.volume_exact() + + # A regular cube with offsets. + r = Range([(1, 1 + K - 1, 1), (2, 2 + N - 1, 1), (3, 3 + M - 1, 1)]) + assert K * N * M == r.volume_exact() + + # A regular cube with strides. + r = Range([(0, K - 1, 2), (0, N - 1, 3), (0, M - 1, 4)]) + assert ceiling(K / 2) * ceiling(N / 3) * ceiling(M / 4) == r.volume_exact() + + # A regular cube with both offsets and strides. + r = Range([(1, 1 + K - 1, 2), (2, 2 + N - 1, 3), (3, 3 + M - 1, 4)]) + assert ceiling(K / 2) * ceiling(N / 3) * ceiling(M / 4) == r.volume_exact() + + # A 2D square on 3D coordinate system. + r = Range([(1, 1 + K - 1, 2), (2, 2, 3), (3, 3 + M - 1, 4)]) + assert ceiling(K / 2) * ceiling(M / 4) == r.volume_exact() + + # A 3D point. + r = Range([(1, 1, 2), (2, 2, 3), (3, 3, 4)]) + assert 1 == r.volume_exact() + + +def test_volume_of_indices(): + # Indices always have volume 1 no matter what, since they are just points. + ind = Indices([0, 1, 2]) + assert 1 == ind.volume_exact() + + ind = Indices([1]) + assert 1 == ind.volume_exact() + + ind = Indices([0, 2]) + assert 1 == ind.volume_exact() + + +def test_subrange_mapping_no_symbols(): + K, N, M = 6, 7, 8 + + # A regular cube. + src = Range([(0, K - 1, 1), (0, N - 1, 1), (0, M - 1, 1)]) + # A regular cube with offsets. + dst = Range([(1, 1 + K - 1, 1), (2, 2 + N - 1, 1), (3, 3 + M - 1, 1)]) + # A Mapper + sm = SubrangeMapper(src, dst) + + # Pick the entire range. + assert dst == sm.map(src) + # Pick a point 0, 0, 0. + assert (Range([(1, 1, 1), (2, 2, 1), (3, 3, 1)]) + == sm.map(Range([(0, 0, 1), (0, 0, 1), (0, 0, 1)]))) + # Pick a point K//2, N//2, M//2. + assert (Range([(1 + K // 2, 1 + K // 2, 1), (2 + N // 2, 2 + N // 2, 1), (3 + M // 2, 3 + M // 2, 1)]) + == sm.map(Range([(K // 2, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1)]))) + # Pick a point K-1, N-1, M-1. + assert (Range([(1 + K - 1, 1 + K - 1, 1), (2 + N - 1, 2 + N - 1, 1), (3 + M - 1, 3 + M - 1, 1)]) + == sm.map(Range([(K - 1, K - 1, 1), (N - 1, N - 1, 1), (M - 1, M - 1, 1)]))) + # Pick a quadrant. + assert (Range([(1, 1 + K // 2, 1), (2, 2 + N // 2, 1), (3, 3 + M // 2, 1)]) + == sm.map(Range([(0, K // 2, 1), (0, N // 2, 1), (0, M // 2, 1)]))) + + +def test_subrange_mapping__with_symbols(): + K, N, M = dace.symbol('K', positive=True), dace.symbol('N', positive=True), dace.symbol('M', positive=True) + + # A regular cube. + src = Range([(0, K - 1, 1), (0, N - 1, 1), (0, M - 1, 1)]) + # A regular cube with offsets. + dst = Range([(1, 1 + K - 1, 1), (2, 2 + N - 1, 1), (3, 3 + M - 1, 1)]) + # A Mapper + sm = SubrangeMapper(src, dst) + + # Pick the entire range. + assert dst == sm.map(src) + + # NOTE: I couldn't make SymPy understand that `(K//2) % K == (K//2)` always holds for postive integers `K`. + # Hence, the numerical approach. + argslist = [{'K': k, 'N': n, 'M': m} for k, n, m in zip(np.random.randint(1, 100, size=20), + np.random.randint(1, 100, size=20), + np.random.randint(1, 100, size=20))] + for args in argslist: # Pick a point K//2, N//2, M//2. - self.assertEqual(Range([(1 + K // 2, 1 + K // 2, 1), (2 + N // 2, 2 + N // 2, 1), (3 + M // 2, 3 + M // 2, 1)]), - sm.map(Range([(K // 2, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1)]))) - # Pick a point K-1, N-1, M-1. - self.assertEqual(Range([(1 + K - 1, 1 + K - 1, 1), (2 + N - 1, 2 + N - 1, 1), (3 + M - 1, 3 + M - 1, 1)]), - sm.map(Range([(K - 1, K - 1, 1), (N - 1, N - 1, 1), (M - 1, M - 1, 1)]))) + want = eval_range( + Range([(1 + K // 2, 1 + K // 2, 1), (2 + N // 2, 2 + N // 2, 1), (3 + M // 2, 3 + M // 2, 1)]), + args) + got = eval_range( + sm.map(Range([(K // 2, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1)])), + args) + assert want == got # Pick a quadrant. - self.assertEqual(Range([(1, 1 + K // 2, 1), (2, 2 + N // 2, 1), (3, 3 + M // 2, 1)]), - sm.map(Range([(0, K // 2, 1), (0, N // 2, 1), (0, M // 2, 1)]))) - - def test_mapping_with_symbols(self): - K, N, M = dace.symbol('K', positive=True), dace.symbol('N', positive=True), dace.symbol('M', positive=True) - - # A regular cube. - src = Range([(0, K - 1, 1), (0, N - 1, 1), (0, M - 1, 1)]) - # A regular cube with offsets. - dst = Range([(1, 1 + K - 1, 1), (2, 2 + N - 1, 1), (3, 3 + M - 1, 1)]) - # A Mapper - sm = SubrangeMapper(src, dst) - - # Pick the entire range. - self.assertEqual(dst, sm.map(src)) - - # NOTE: I couldn't make SymPy understand that `(K//2) % K == (K//2)` always holds for postive integers `K`. - # Hence, the numerical approach. - argslist = [{'K': k, 'N': n, 'M': m} for k, n, m in zip(np.random.randint(1, 100, size=20), - np.random.randint(1, 100, size=20), - np.random.randint(1, 100, size=20))] - for args in argslist: - # Pick a point K//2, N//2, M//2. - want = eval_range( - Range([(1 + K // 2, 1 + K // 2, 1), (2 + N // 2, 2 + N // 2, 1), (3 + M // 2, 3 + M // 2, 1)]), - args) - got = eval_range( - sm.map(Range([(K // 2, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1)])), - args) - self.assertEqual(want, got) - # Pick a quadrant. - want = eval_range( - Range([(1, 1 + K // 2, 1), (2, 2 + N // 2, 1), (3, 3 + M // 2, 1)]), - args) - got = eval_range( - sm.map(Range([(0, K // 2, 1), (0, N // 2, 1), (0, M // 2, 1)])), - args) - self.assertEqual(want, got) - - def test_mapping_with_reshaping(self): - K, N, M = dace.symbol('K', positive=True), dace.symbol('N', positive=True), dace.symbol('M', positive=True) - - # A regular cube. - src = Range([(0, K - 1, 1), (0, N - 1, 1), (0, M - 1, 1)]) - # A regular cube with different shape. - dst = Range([(0, K - 1, 1), (0, N * M - 1, 1)]) - # A Mapper - sm = SubrangeMapper(src, dst) - sm_inv = SubrangeMapper(dst, src) - - # Pick the entire range. - self.assertEqual(dst, sm.map(src)) - self.assertEqual(src, sm_inv.map(dst)) - - # NOTE: I couldn't make SymPy understand that `(K//2) % K == (K//2)` always holds for postive integers `K`. - # Hence, the numerical approach. - argslist = [{'K': k, 'N': n, 'M': m} for k, n, m in zip(np.random.randint(1, 10, size=20), - np.random.randint(1, 10, size=20), - np.random.randint(1, 10, size=20))] - # Pick a point K//2, N//2, M//2. - for args in argslist: - orig = Range([(K // 2, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1)]) - orig_maps_to = Range([(K // 2, K // 2, 1), ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1)]) - want, got = eval_range(orig_maps_to, args), eval_range(sm.map(orig), args) - self.assertEqual(want, got) - want, got = eval_range(orig, args), eval_range(sm_inv.map(orig_maps_to), args) - self.assertEqual(want, got) - # Pick a quadrant. - for args in argslist: - # But its mapping cannot be expressed as a simple range with offset and stride. - self.assertIsNone(sm.map(Range([(0, K // 2, 1), (0, N // 2, 1), (0, M // 2, 1)]))) - # Pick only points in problematic quadrants, but larger subsets elsewhere. - for args in argslist: - orig = Range([(0, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1)]) - orig_maps_to = Range([(0, K // 2, 1), ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1)]) - want, got = eval_range(orig_maps_to, args), eval_range(sm.map(orig), args) - self.assertEqual(want, got) - want, got = eval_range(orig, args), eval_range(sm_inv.map(orig_maps_to), args) - self.assertEqual(want, got) - - def test_mapping_with_reshaping_unit_dims(self): - K, N, M = dace.symbol('K', positive=True), dace.symbol('N', positive=True), dace.symbol('M', positive=True) - - # A regular cube. - src = Range([(0, K - 1, 1), (0, N - 1, 1), (0, M - 1, 1), (0, 0, 1)]) - # A regular cube with different shape. - dst = Range([(0, K - 1, 1), (0, 0, 1), (0, N * M - 1, 1), (0, 0, 1), (0, 0, 1)]) - # A Mapper - sm = SubrangeMapper(src, dst) - sm_inv = SubrangeMapper(dst, src) - - # Pick the entire range. - self.assertEqual(dst, sm.map(src)) - self.assertEqual(src, sm_inv.map(dst)) - - # NOTE: I couldn't make SymPy understand that `(K//2) % K == (K//2)` always holds for postive integers `K`. - # Hence, the numerical approach. - argslist = [{'K': k, 'N': n, 'M': m} for k, n, m in zip(np.random.randint(1, 10, size=20), - np.random.randint(1, 10, size=20), - np.random.randint(1, 10, size=20))] - # Pick a point K//2, N//2, M//2. - for args in argslist: - orig = Range([(K // 2, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1), (0, 0, 1)]) - orig_maps_to = Range([(K // 2, K // 2, 1), - (0, 0, 1), - ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1), - (0, 0, 1), (0, 0, 1)]) - want, got = eval_range(orig_maps_to, args), eval_range(sm.map(orig), args) - self.assertEqual(want, got) - want, got = eval_range(orig, args), eval_range(sm_inv.map(orig_maps_to), args) - self.assertEqual(want, got) - # Pick a quadrant. - for args in argslist: - # But its mapping cannot be expressed as a simple range with offset and stride. - self.assertIsNone(sm.map(Range([(0, K // 2, 1), (0, N // 2, 1), (0, M // 2, 1), (0, 0, 1)]))) - # Pick only points in problematic quadrants, but larger subsets elsewhere. - for args in argslist: - orig = Range([(0, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1), (0, 0, 1)]) - orig_maps_to = Range([(0, K // 2, 1), - (0, 0, 1), - ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1), - (0, 0, 1), (0, 0, 1)]) - want, got = eval_range(orig_maps_to, args), eval_range(sm.map(orig), args) - self.assertEqual(want, got) - want, got = eval_range(orig, args), eval_range(sm_inv.map(orig_maps_to), args) - self.assertEqual(want, got) + want = eval_range( + Range([(1, 1 + K // 2, 1), (2, 2 + N // 2, 1), (3, 3 + M // 2, 1)]), + args) + got = eval_range( + sm.map(Range([(0, K // 2, 1), (0, N // 2, 1), (0, M // 2, 1)])), + args) + assert want == got + + +def test_subrange_mapping__with_reshaping(): + K, N, M = dace.symbol('K', positive=True), dace.symbol('N', positive=True), dace.symbol('M', positive=True) + + # A regular cube. + src = Range([(0, K - 1, 1), (0, N - 1, 1), (0, M - 1, 1)]) + # A regular cube with different shape. + dst = Range([(0, K - 1, 1), (0, N * M - 1, 1)]) + # A Mapper + sm = SubrangeMapper(src, dst) + sm_inv = SubrangeMapper(dst, src) + + # Pick the entire range. + assert dst == sm.map(src) + assert src == sm_inv.map(dst) + + # NOTE: I couldn't make SymPy understand that `(K//2) % K == (K//2)` always holds for postive integers `K`. + # Hence, the numerical approach. + argslist = [{'K': k, 'N': n, 'M': m} for k, n, m in zip(np.random.randint(1, 10, size=20), + np.random.randint(1, 10, size=20), + np.random.randint(1, 10, size=20))] + # Pick a point K//2, N//2, M//2. + for args in argslist: + orig = Range([(K // 2, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1)]) + orig_maps_to = Range([(K // 2, K // 2, 1), ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1)]) + want, got = eval_range(orig_maps_to, args), eval_range(sm.map(orig), args) + assert want == got + want, got = eval_range(orig, args), eval_range(sm_inv.map(orig_maps_to), args) + assert want == got + # Pick a quadrant. + # But its mapping cannot be expressed as a simple range with offset and stride. + assert sm.map(Range([(0, K // 2, 1), (0, N // 2, 1), (0, M // 2, 1)])) is None + # Pick only points in problematic quadrants, but larger subsets elsewhere. + for args in argslist: + orig = Range([(0, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1)]) + orig_maps_to = Range([(0, K // 2, 1), ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1)]) + want, got = eval_range(orig_maps_to, args), eval_range(sm.map(orig), args) + assert want == got + want, got = eval_range(orig, args), eval_range(sm_inv.map(orig_maps_to), args) + assert want == got + + +def test_subrange_mapping__with_reshaping_unit_dims(): + K, N, M = dace.symbol('K', positive=True), dace.symbol('N', positive=True), dace.symbol('M', positive=True) + + # A regular cube. + src = Range([(0, K - 1, 1), (0, N - 1, 1), (0, M - 1, 1), (0, 0, 1)]) + # A regular cube with different shape. + dst = Range([(0, K - 1, 1), (0, 0, 1), (0, N * M - 1, 1), (0, 0, 1), (0, 0, 1)]) + # A Mapper + sm = SubrangeMapper(src, dst) + sm_inv = SubrangeMapper(dst, src) + + # Pick the entire range. + assert dst == sm.map(src) + assert src == sm_inv.map(dst) + + # NOTE: I couldn't make SymPy understand that `(K//2) % K == (K//2)` always holds for postive integers `K`. + # Hence, the numerical approach. + argslist = [{'K': k, 'N': n, 'M': m} for k, n, m in zip(np.random.randint(1, 10, size=20), + np.random.randint(1, 10, size=20), + np.random.randint(1, 10, size=20))] + # Pick a point K//2, N//2, M//2. + for args in argslist: + orig = Range([(K // 2, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1), (0, 0, 1)]) + orig_maps_to = Range([(K // 2, K // 2, 1), + (0, 0, 1), + ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1), + (0, 0, 1), (0, 0, 1)]) + want, got = eval_range(orig_maps_to, args), eval_range(sm.map(orig), args) + assert want == got + want, got = eval_range(orig, args), eval_range(sm_inv.map(orig_maps_to), args) + assert want == got + # Pick a quadrant. + # But its mapping cannot be expressed as a simple range with offset and stride. + assert sm.map(Range([(0, K // 2, 1), (0, N // 2, 1), (0, M // 2, 1), (0, 0, 1)])) is None + # Pick only points in problematic quadrants, but larger subsets elsewhere. + for args in argslist: + orig = Range([(0, K // 2, 1), (N // 2, N // 2, 1), (M // 2, M // 2, 1), (0, 0, 1)]) + orig_maps_to = Range([(0, K // 2, 1), + (0, 0, 1), + ((N // 2) + (M // 2) * N, (N // 2) + (M // 2) * N, 1), + (0, 0, 1), (0, 0, 1)]) + want, got = eval_range(orig_maps_to, args), eval_range(sm.map(orig), args) + assert want == got + want, got = eval_range(orig, args), eval_range(sm_inv.map(orig_maps_to), args) + assert want == got if __name__ == '__main__': - unittest.main() + test_volume_of_range() + test_volume_of_indices() + test_subrange_mapping_no_symbols() + test_subrange_mapping__with_symbols() + test_subrange_mapping__with_reshaping() + test_subrange_mapping__with_reshaping_unit_dims() From 8dd73d53ae52c7c45f847be4796562dab8ff117a Mon Sep 17 00:00:00 2001 From: Pratyai Mazumder Date: Thu, 24 Oct 2024 18:34:50 +0200 Subject: [PATCH 8/8] JUST CHECKING CI --- dace/libraries/standard/nodes/reduce.py | 3 +- .../dataflow/redundant_array.py | 236 +++++------------- 2 files changed, 62 insertions(+), 177 deletions(-) diff --git a/dace/libraries/standard/nodes/reduce.py b/dace/libraries/standard/nodes/reduce.py index fa231c07f2..68cb45e5a7 100644 --- a/dace/libraries/standard/nodes/reduce.py +++ b/dace/libraries/standard/nodes/reduce.py @@ -103,7 +103,8 @@ def expansion(node: 'Reduce', state: SDFGState, sdfg: SDFG): 'reduce_init', {'_o%d' % i: '0:%s' % symstr(d) for i, d in enumerate(outedge.data.subset.size())}, {}, '__out = %s' % node.identity, - {'__out': dace.Memlet.simple('_out', ','.join(['_o%d' % i for i in range(output_dims)]))}, + # {'__out': dace.Memlet.simple('_out', ','.join(['_o%d' % i for i in range(output_dims)]))}, + {'__out': dace.Memlet.simple('_out', ','.join(['_o%d' % i for i in osqdim]))}, external_edges=True) else: nstate = nsdfg.add_state() diff --git a/dace/transformation/dataflow/redundant_array.py b/dace/transformation/dataflow/redundant_array.py index 7b241ff9cd..b0b9ecc843 100644 --- a/dace/transformation/dataflow/redundant_array.py +++ b/dace/transformation/dataflow/redundant_array.py @@ -7,9 +7,10 @@ from typing import Dict, List, Optional, Tuple import networkx as nx +from dace.subsets import Range, Indices, SubrangeMapper from networkx.exception import NetworkXError, NodeNotFound -from dace import data, dtypes +from dace import data, dtypes, Memlet from dace import memlet as mm from dace import subsets, symbolic from dace.config import Config @@ -24,7 +25,7 @@ def _validate_subsets(edge: graph.MultiConnectorEdge, arrays: Dict[str, data.Data], src_name: str = None, - dst_name: str = None) -> Tuple[subsets.Subset]: + dst_name: str = None) -> Tuple[subsets.Subset, subsets.Subset]: """ Extracts and validates src and dst subsets from the edge. """ # Find src and dst names @@ -499,182 +500,65 @@ def _is_reshaping_memlet( return True - def apply(self, graph, sdfg): - in_array = self.in_array - out_array = self.out_array - in_desc = sdfg.arrays[in_array.data] - out_desc = sdfg.arrays[out_array.data] + def apply(self, graph: SDFGState, sdfg: SDFG): + # The pattern is A ---> B, and we want to remove A + A, B = self.in_array, self.out_array # 1. Get edge e1 and extract subsets for arrays A and B - e1 = graph.edges_between(in_array, out_array)[0] - a1_subset, b_subset = _validate_subsets(e1, sdfg.arrays) - - # View connected to a view: simple case - if (isinstance(in_desc, data.View) and isinstance(out_desc, data.View)): - simple_case = True - for e in graph.in_edges(in_array): - if e.data.dst_subset is not None and a1_subset != e.data.dst_subset: - simple_case = False - break - if simple_case: - for e in graph.in_edges(in_array): - for e2 in graph.memlet_tree(e): - if e2 is e: - continue - if e2.data.data == in_array.data: - e2.data.data = out_array.data - new_memlet = copy.deepcopy(e.data) - if new_memlet.data == in_array.data: - new_memlet.data = out_array.data - new_memlet.dst_subset = b_subset - graph.add_edge(e.src, e.src_conn, out_array, e.dst_conn, new_memlet) - graph.remove_node(in_array) - try: - if in_array.data in sdfg.arrays: - sdfg.remove_data(in_array.data) - except ValueError: # Used somewhere else - pass - return - - # Find extraneous A or B subset dimensions - a_dims_to_pop = [] - b_dims_to_pop = [] - bset = b_subset - popped = [] - if a1_subset and b_subset and a1_subset.dims() != b_subset.dims(): - a_size = a1_subset.size_exact() - b_size = b_subset.size_exact() - if a1_subset.dims() > b_subset.dims(): - a_dims_to_pop = find_dims_to_pop(a_size, b_size) - else: - b_dims_to_pop = find_dims_to_pop(b_size, a_size) - bset, popped = pop_dims(b_subset, b_dims_to_pop) - - from dace.libraries.standard import Reduce - reduction = False - for e in graph.in_edges(in_array): - if isinstance(e.src, Reduce) or (isinstance(e.src, nodes.NestedSDFG) - and len(in_desc.shape) != len(out_desc.shape)): - reduction = True - - # If: - # 1. A reduce node is involved; or - # 2. A NestedSDFG node is involved and the arrays have different dimensionality; or - # 3. The memlet does not cover the removed array; or - # 4. Dimensions are mismatching (all dimensions are popped); - # create a view. - if ( - reduction - or len(a_dims_to_pop) == len(in_desc.shape) - or any(m != a for m, a in zip(a1_subset.size(), in_desc.shape)) - ): - self._make_view(sdfg, graph, in_array, out_array, e1, b_subset, b_dims_to_pop) - return in_array - - # TODO: Fix me. - # As described in [issue 1595](https://github.com/spcl/dace/issues/1595) the - # transformation is unable to handle certain cases of reshaping Memlets - # correctly and fixing this case has proven rather difficult. In a first - # attempt the case of reshaping Memlets was forbidden (in the - # `can_be_applied()` method), however, this caused other (useful) cases to - # fail. For that reason such Memlets are transformed to Views. - # This is a fix and it should be addressed. - if self._is_reshaping_memlet(graph=graph, edge=e1): - self._make_view(sdfg, graph, in_array, out_array, e1, b_subset, b_dims_to_pop) - return in_array - - # Validate that subsets are composable. If not, make a view - try: - for e2 in graph.in_edges(in_array): - path = graph.memlet_tree(e2) - wcr = e1.data.wcr - wcr_nonatomic = e1.data.wcr_nonatomic - for e3 in path: - # 2-a. Extract subsets for array B and others - other_subset, a3_subset = _validate_subsets(e3, sdfg.arrays, dst_name=in_array.data) - # 2-b. Modify memlet to match array B. - dname = out_array.data - src_is_data = False - a3_subset.offset(a1_subset, negative=True) - - if a3_subset and a_dims_to_pop: - aset, _ = pop_dims(a3_subset, a_dims_to_pop) - else: - aset = a3_subset - - compose_and_push_back(bset, aset, b_dims_to_pop, popped) - except (ValueError, NotImplementedError): - self._make_view(sdfg, graph, in_array, out_array, e1, b_subset, b_dims_to_pop) - print(f"CREATED VIEW(2): {in_array}") - return in_array - - # 2. Iterate over the e2 edges and traverse the memlet tree - for e2 in graph.in_edges(in_array): - path = graph.memlet_tree(e2) - wcr = e1.data.wcr - wcr_nonatomic = e1.data.wcr_nonatomic - for e3 in path: - # 2-a. Extract subsets for array B and others - other_subset, a3_subset = _validate_subsets(e3, sdfg.arrays, dst_name=in_array.data) - # 2-b. Modify memlet to match array B. - dname = out_array.data - src_is_data = False - a3_subset.offset(a1_subset, negative=True) - - if a3_subset and a_dims_to_pop: - aset, _ = pop_dims(a3_subset, a_dims_to_pop) - else: - aset = a3_subset - - dst_subset = compose_and_push_back(bset, aset, b_dims_to_pop, popped) - # NOTE: This fixes the following case: - # Tasklet ----> A[subset] ----> ... -----> A - # Tasklet is not data, so it doesn't have an other subset. - if isinstance(e3.src, nodes.AccessNode): - if e3.src.data == out_array.data: - dname = e3.src.data - src_is_data = True - src_subset = other_subset - else: - src_subset = None - - subset = src_subset if src_is_data else dst_subset - other_subset = dst_subset if src_is_data else src_subset - e3.data.data = dname - e3.data.subset = subset - e3.data.other_subset = other_subset - wcr = wcr or e3.data.wcr - wcr_nonatomic = wcr_nonatomic or e3.data.wcr_nonatomic - e3.data.wcr = wcr - e3.data.wcr_nonatomic = wcr_nonatomic - - # 2-c. Remove edge and add new one - graph.remove_edge(e2) - e2.data.wcr = wcr - e2.data.wcr_nonatomic = wcr_nonatomic - graph.add_edge(e2.src, e2.src_conn, out_array, e2.dst_conn, e2.data) - - # 2-d. Fix strides in nested SDFGs - if in_desc.strides != out_desc.strides: - sources = [] - if path.downwards: - sources = [path.root().edge] - else: - sources = [e for e in path.leaves()] - for source_edge in sources: - if not isinstance(source_edge.src, nodes.NestedSDFG): - continue - conn = source_edge.src_conn - inner_desc = source_edge.src.sdfg.arrays[conn] - inner_desc.strides = out_desc.strides - - # Finally, remove in_array node - graph.remove_node(in_array) - try: - if in_array.data in sdfg.arrays: - sdfg.remove_data(in_array.data) - except ValueError: # Already in use (e.g., with Views) - pass + e_ab = graph.edges_between(A, B) + assert len(e_ab) == 1 + e_ab = e_ab[0] + print(e_ab) + a_subset, b_subset = _validate_subsets(e_ab, sdfg.arrays) + # Other cases should have been handled in `can_be_applied()`. + assert isinstance(a_subset, Range) or isinstance(a_subset, Indices) + assert isinstance(b_subset, Range) or isinstance(b_subset, Indices) + # And this should be self-evident. + assert a_subset.volume_exact() == b_subset.volume_exact() + + for ie in graph.in_edges(A): + # The pattern is now: C -(ie)-> A ---> B + path = graph.memlet_tree(ie) + for pe in path: + # The pattern is now: C -(pe)-> C1 ---> ... ---> A ---> B + print('PE:', pe) + c_subset, a0_subset = _validate_subsets(pe, sdfg.arrays, dst_name=A.data) + print('c, a0:', c_subset, a0_subset) + if a0_subset is None: + continue + + # Other cases should have been handled already in `can_be_applied()`. + assert c_subset is None or isinstance(c_subset, Range) or isinstance(c_subset, Indices) + if c_subset is not None: + assert c_subset.volume_exact() == a0_subset.volume_exact() + assert isinstance(a0_subset, Range) or isinstance(a0_subset, Indices) + print('SUBS:', a_subset, '|', a0_subset) + assert a_subset.dims() == a0_subset.dims() + # assert a_subset.covers_precise(a0_subset) + # assert all(b0 >= b and (b0 - b) % s == 0 and s0 % s == 0 + # for (b, e, s), (b0, e0, s0) in zip(a_subset.ndrange(), a0_subset.ndrange())) + + # Find out where `a0_subset` maps to, given that `a_subset` precisely maps to `b_subset`. + # `reshapr` describes how `a_subset` maps to `b_subset`. + reshapr = SubrangeMapper(a_subset, b_subset) + # `b0_subset` is the mapping for `a0_subset`. + b0_subset = reshapr.map(a0_subset) + print(a_subset, b_subset) + print(a0_subset, b0_subset) + assert isinstance(b0_subset, Range) or isinstance(b0_subset, Indices) + assert b0_subset.volume_exact() == a0_subset.volume_exact() + + # Now we can replace the path: C -(pe)-> C1 ---> ... ---> A ---> B + # with an equivalent path: C -(pe)-> C1 ---> ... ---> B + dst, dst_conn = (B, None) if pe.dst is A else (pe.dst, pe.dst_conn) + print('dst:', pe.src, dst, dst_conn) + print('mem:', B.data, b0_subset, c_subset) + e = graph.add_edge(pe.src, pe.src_conn, dst, dst_conn, + memlet=Memlet(data=B.data, subset=b0_subset, other_subset=c_subset)) + print('e:', e) + graph.remove_edge(pe) + graph.remove_node(A) + sdfg.remove_data(A.data) class RedundantSecondArray(pm.SingleStateTransformation):