diff --git a/pymbolic/imperative/transform.py b/pymbolic/imperative/transform.py index a52d44630b44266274f8703502154941796e7f15..c380a49b83c738cae123382fecbb9fc2474eea0c 100644 --- a/pymbolic/imperative/transform.py +++ b/pymbolic/imperative/transform.py @@ -23,14 +23,32 @@ THE SOFTWARE. """ +# {{{ fuse overlapping instruction streams + +def fuse_instruction_streams_with_overlapping_ids(instructions_a, instructions_b, + allowed_duplicates=[]): + id_a = set([insna.id for insna in instructions_a]) + # filter b instructions + uniques_b = [x for x in instructions_b + if x.id in id_a + and x.id not in allowed_duplicates] + + return fuse_instruction_streams_with_unique_ids(instructions_a, uniques_b, + allow_b_depend_on_a=True) + +# }}} + + # {{{ fuse instruction streams -def fuse_instruction_streams_with_unique_ids(instructions_a, instructions_b): +def fuse_instruction_streams_with_unique_ids(instructions_a, instructions_b, + allow_b_depend_on_a=False): new_instructions = list(instructions_a) from pytools import UniqueNameGenerator insn_id_gen = UniqueNameGenerator( set([insna.id for insna in new_instructions])) + a_ids = set([insna.id for insna in instructions_a]) b_unique_instructions = [] old_b_id_to_new_b_id = {} for insnb in instructions_b: @@ -42,11 +60,23 @@ def fuse_instruction_streams_with_unique_ids(instructions_a, instructions_b): insnb.copy(id=new_id)) for insnb in b_unique_instructions: + b_deps = set() + for dep_id in insnb.depends_on: + # First, see if the dependency is in the updated `instructions_b` + # ids. + new_dep_id = old_b_id_to_new_b_id.get(dep_id) + # Next if `allow_b_depend_on_a` is True, check for the dendency in the + # `a_ids` + if allow_b_depend_on_a and new_dep_id is None: + new_dep_id = dep_id if dep_id in a_ids else None + # If the dependency is not found in `instructions_a` or the updated + # `instructions_b`, raise an AssertionError. + assert new_dep_id is not None, ('Instruction {0} in stream b ' + 'missing dependency {1}'.format(insnb.id, dep_id)) + b_deps.add(new_dep_id) + new_instructions.append( - insnb.copy( - depends_on=frozenset( - old_b_id_to_new_b_id[dep_id] - for dep_id in insnb.depends_on))) + insnb.copy(depends_on=frozenset(b_deps))) return new_instructions, old_b_id_to_new_b_id diff --git a/test/test_pymbolic.py b/test/test_pymbolic.py index fe7c8a395b1ad433395d9e0771f6b833ca74884f..252f871d7416d148c39074be0b661a2d3deb5ed7 100644 --- a/test/test_pymbolic.py +++ b/test/test_pymbolic.py @@ -487,6 +487,102 @@ def test_long_sympy_mapping(): SympyToPymbolicMapper()(sp.sympify(int(10))) +# dummy instruction +class DummyInsn(object): + def __init__(self, idv, depends_on=frozenset()): + self.id = idv + self.depends_on = depends_on + + def copy(self, id=None, depends_on=None): + if id is not None: + self.id = id + if depends_on is not None: + self.depends_on = depends_on + return DummyInsn(self.id, self.depends_on.copy()) + + +def test_unique_fusion(): + from pymbolic.imperative.transform import \ + fuse_instruction_streams_with_unique_ids + + # create instructions 1-5 with depends + insns_a = [DummyInsn(str(i), frozenset([str(x) for x in range(i)])) + for i in range(5)] + + # first try feeding it 'a' twice + new_insns, b_map = fuse_instruction_streams_with_unique_ids( + insns_a, [i.copy() for i in insns_a]) + new_insn_ids = set([x.id for x in new_insns]) + + # should have a duplicate of each instruction + assert all(x.id in b_map for x in insns_a) + assert all(x.id + '_0' in new_insn_ids for x in insns_a) + + # next create independent instructions 5-10 + insns_b = [DummyInsn(str(i), frozenset([str(x) for x in range(5, i)])) + for i in range(5, 10)] + + new_insns, b_map = fuse_instruction_streams_with_unique_ids( + insns_a, insns_b) + new_insn_ids = set([x.id for x in new_insns]) + + # should have both sets of instructions + assert all(x.id in new_insn_ids for x in + insns_a + insns_b) + assert all(k == v for k, v in b_map.items()) + + # next test b depends on a without flag + insns_b[0].depends_on = frozenset(['0']) + + with pytest.raises(AssertionError): + fuse_instruction_streams_with_unique_ids( + insns_a, insns_b) + + # and with flag + new_insns, b_map = fuse_instruction_streams_with_unique_ids( + insns_a, insns_b, allow_b_depend_on_a=True) + + # should have both sets of instructions + assert all(x.id in new_insn_ids for x in + insns_a + insns_b) + assert all(k == v for k, v in b_map.items()) + + +def test_overlapping_fusion(): + from pymbolic.imperative.transform import \ + fuse_instruction_streams_with_overlapping_ids + + # create instructions 1-5 with depends + insns_a = [DummyInsn(str(i)) + for i in range(5)] + a_copy = [i.copy() for i in insns_a] + a_ids = [i.id for i in insns_a] + + # first try feeding it 'a' twice + new_insns, b_map = fuse_instruction_streams_with_overlapping_ids( + insns_a, a_copy, a_ids) + + # should have just the original instructions + assert len(new_insns) == len(insns_a) + assert all(new_insns[i] == insns_a[i] for i in range(len(new_insns))) + assert len(b_map) == 0 + + # now try with no allowed overlaps + new_insns, b_map = fuse_instruction_streams_with_overlapping_ids( + insns_a, a_copy, []) + # check all a in new instructions + new_ids = set([x.id for x in new_insns]) + # and all a in maps + assert all(insn_a.id in new_ids for insn_a in insns_a) + assert all(insn_a.id in b_map for insn_a in insns_a) + + # try with b depends on a, and the instruction overlapping + insns_b = a_copy[:] + insns_b[1].depends_on = frozenset([insns_a[0].id]) + new_insns, b_map = fuse_instruction_streams_with_overlapping_ids( + insns_a, a_copy, [insns_a[0].id]) + + if __name__ == "__main__": import sys if len(sys.argv) > 1: