generate_loop_schedule_v2 #350

kaushikcfd · 2021-04-28T22:45:20Z

Implementation for finding loop nest around map in O(N.k), 'N' being the number of inames and 'k' being the max. loop depth.

For comparison, let's consider the kernel in #288: on main this map in computed in 5 minutes and this branch takes 30 0.4 seconds.

setup.py

inducer

Some thoughts from a quick look here.

loopy/schedule/tools.py

inducer · 2021-05-01T23:21:03Z

loopy/schedule/__init__.py


-    iname_to_insns = kernel.iname_to_insns()
+    loop_nest_around_map = defaultdict(frozenset)


In my view, the loop_nest_around_map should go. The loop nest tree is a strictly more informative data structure, and it is not of size O(n^2). I think we should get rid of it. I looked through the uses, and it shouldn't be hard to replace, and it'll fix another bottleneck in linearization.

In the new version, I didn't touch the old implementation (v1-scheduler) to keep things simpler in this PR. So, for such cases $O(n^2)$ stays.

I can handle it as a separate PR.

inducer · 2021-05-01T23:22:25Z

loopy/schedule/tools.py

+    Updates *tree* to make *inames_to_pull_out* a loop nesting level in
+    *loop_nests*
+
+    :returns: a :class:`tuple` ``(outer_loop_nest, inner_loop_nest)``, where


Describe the types of the arguments.

Thanks, done.

inducer · 2021-05-01T23:26:14Z

loopy/schedule/tools.py

+    Updates *tree* to make *inames_to_pull_out* a loop nesting level in
+    *loop_nests*


Describe which is nested "inside" and which is "outside".

inducer · 2021-05-26T23:04:45Z

Does #372 supersede this?

kaushikcfd · 2021-05-26T23:14:33Z

@inducer: Not really. #372 includes commits from this branch, so that I could do some tests on my test problems. But I propose we merge them separately. (I've updated the description of #372 to record that)

inducer · 2022-06-09T20:14:08Z

I was just rebasing this. You beat me to it! :)

inducer

Here's a first look. I haven't really gotten very far into the core algorithm, because I'm not sure I understood _pull_out_loop_nest_tree, which seems like it's a core operation.

inducer · 2022-06-09T20:21:27Z

loopy/tools.py

@@ -930,6 +933,225 @@ def _get_persistent_hashable_arg(arg):

    return wrapper

+
+# {{{ tree data structure


This is big enough to be its own file.

#694 puts the Tree class into its own file (loopy/schedule/tree.py).

inducer · 2022-06-09T20:22:21Z

loopy/tools.py

@@ -930,6 +933,225 @@ def _get_persistent_hashable_arg(arg):

    return wrapper

+
+# {{{ tree data structure


Given that this has type annotations, mypy should like them (and there should be CI saying as much).

As far as I can see, mypy checks and passes the annotations.

inducer · 2022-06-09T20:23:15Z

loopy/tools.py

+@dataclass(frozen=True)
+class Tree(Generic[T]):
+    """
+    An immutable tree implementation.


Describe the role of the type variable T. Specifically describe that there's one Tree object, but many nodes T. Maybe rename T to NodeT?

c53a9af hopefully clarifies this (and renames T to NodeT).

inducer · 2022-06-09T20:24:57Z

loopy/tools.py

+
+
+@dataclass(frozen=True)
+class Tree(Generic[T]):


Is this mature enough to be in pytools?

inducer · 2022-06-09T20:26:53Z

loopy/tools.py

+    _parent_to_children: "PMap[T, FrozenSet[T]]"
+    _child_to_parent: "PMap[T, Optional[T]]"


In terms of data structure, this is a forest (i.e. it allows multiple trees). It wouldn't take much to allow that fully... YAGNI though.

inducer · 2022-06-09T21:20:11Z

loopy/schedule/tools.py

+    :returns: a :class:`tuple` ``(new_tree, outer_loop_nest, inner_loop_nest)``,
+        where outer_loop_nest is the identifier for the new outer and inner
+        loop nests so that *inames_to_pull_out* is a valid nesting.


Make a dataclass for the result?

inducer · 2022-06-09T21:25:34Z

loopy/schedule/tools.py

+        :class:`frozenset` of inames representing a loop nest. For example a
+        tree might look like:
+
+    :arg loop_nests: A collection of nodes in *tree* that cover


collection isn't an ABC. Do you mean sequence?

inducer · 2022-06-09T21:34:19Z

loopy/schedule/tools.py

+
+    loop_nests = sorted(loop_nests, key=lambda nest: tree.depth(nest))
+
+    for outer, inner in zip(loop_nests[:-1], loop_nests[1:]):


It seems that all you'd need to pass here is the innermost node in the loop nest tree, and tree.ancestors (at least the proposed changed version) would cheaply give you this without the need to enforce these invariants.

inducer · 2022-06-09T21:40:58Z

loopy/schedule/tools.py

+    obtained after constraining *loop_nest_tree* with the constraints enforced
+    by *priorities*.
+
+    :arg strict_priorities: Expresses strict nesting constraints similar to


"similar to" -> "using the same data structure as"?

inducer · 2022-06-09T21:41:02Z

loopy/schedule/tools.py

+        strictly i.e. if these conditions cannot be met a
+        :class:`loopy.diagnostic.LoopyError` is raised.
+
+    :arg relaxed_priorities: Expresses strict nesting constraints similar to


"similar to" -> "using the same data structure as"?

inducer

Thanks! Finally made it all the way through. There's a bit to digest here, but as far as I remember, it's mostly superficial. I didn't find any flaws in the thinking. Nice job!

inducer · 2022-06-09T22:46:29Z

loopy/schedule/tools.py

+    Returns a copy of *tree* that realizes *inames_to_pull_out* as loop
+    nesting.


Returns a version of the loop nest tree tree so that every node in the tree is either a subset of outermost_inames or has an empty intersection with outermost_inames.

This routine modifies at most one node of the tree. All its ancestors must satisfy ancestor <= outermost_inames. For the first node not satisfying this relationship, if node & outermost_inames is empty, no modification is made.
Otherwise, if node & outermost_inames < node, that node is split so as to separate outermost_inames in their own node.

inducer · 2022-06-09T22:49:35Z

loopy/schedule/tools.py

+    # }}}
+
+    innermost_loop_nest = loop_nests[-1]
+    new_outer_loop_nest = inames_to_pull_out - reduce(frozenset.union,


new_outermost_node

inducer · 2022-06-09T22:49:55Z

loopy/schedule/tools.py

+    new_outer_loop_nest = inames_to_pull_out - reduce(frozenset.union,
+                                                      loop_nests[:-1],
+                                                      frozenset())
+    new_inner_loop_nest = innermost_loop_nest - inames_to_pull_out


new_inner_node

inducer · 2022-06-09T22:50:14Z

loopy/schedule/tools.py

+
+    # }}}
+
+    innermost_loop_nest = loop_nests[-1]


innermost_node

inducer · 2022-06-09T23:01:23Z

loopy/schedule/tools.py

+                        # inner iname and outer iname are indirect family members
+                        # => must be realized via dependencies in the linearization
+                        # phase, not implemented in v2-scheduler yet.
+                        from loopy.schedule import V2SchedulerNotImplementedException


s/scheduler/linearizer/? (here and throughout)

inducer · 2022-06-09T23:51:27Z

loopy/schedule/tools.py

+        for iname in node:
+            iname_to_tree_node_id[iname] = node
+
+    return pmap(iname_to_tree_node_id)


This might be better as an immutables.Map.

inducer · 2022-06-09T23:56:54Z

loopy/schedule/__init__.py

+                  if kernel.iname_tags_of_type(iname, VectorizeTag)}
+    parallel_inames = (concurrent_inames - ilp_inames - vec_inames)
+
+    # {{{ can v2 scheduler handle??


...handle this kernel?

inducer · 2022-06-09T23:58:55Z

loopy/schedule/__init__.py

+    pass
+
+
+def generate_loop_schedules_v2(kernel):


This is nice! I like this quite a bit.

inducer · 2022-06-10T00:02:08Z

loopy/schedule/__init__.py

+        if outer_loop == "":
+            continue
+
+        for child in loop_nest_tree.children(outer_loop):


Add a comment that this only works because loop_nest_tree contains a total order of loops. Otherwise it might windup with mismatched enter/leaves.

inducer · 2022-06-10T00:04:57Z

loopy/schedule/__init__.py

+            # for i
+            #   insn
+            # end i
+            # 'insn' *must* come b/w 'for i' and 'end i'


b/w = between? If so, spell out.

inducer · 2022-06-10T00:07:06Z

Unsubscribing... @-mention or request review once it's ready for a look or needs attention.

But I'm eager to get this in soon!

…me LoopKernel

kaushikcfd requested a review from inducer April 28, 2021 23:20

kaushikcfd commented Apr 29, 2021

View reviewed changes

setup.py Outdated Show resolved Hide resolved

inducer reviewed May 1, 2021

View reviewed changes

kaushikcfd marked this pull request as draft May 4, 2021 15:08

kaushikcfd changed the title ~~Better loop around nest map~~ generate_loop_schedule_v2 May 4, 2021

inducer mentioned this pull request May 11, 2021

Add loop nest constraints v2 to better loop around nest map #360

Draft

kaushikcfd force-pushed the better_loop_around_nest_map branch 2 times, most recently from 96856f8 to 0203d09 Compare May 11, 2021 13:44

kaushikcfd mentioned this pull request May 26, 2021

Move to an AST-like IR before entering codegen #372

Draft

2 tasks

kaushikcfd force-pushed the better_loop_around_nest_map branch from e14a03d to 2d526dc Compare June 8, 2021 18:51

kaushikcfd marked this pull request as ready for review June 10, 2021 06:31

kaushikcfd requested a review from inducer June 11, 2021 22:00

kaushikcfd force-pushed the better_loop_around_nest_map branch from e8b03c8 to 72334e8 Compare August 16, 2021 17:23

kaushikcfd force-pushed the better_loop_around_nest_map branch from 6945d16 to d74da04 Compare September 28, 2021 15:16

kaushikcfd mentioned this pull request Oct 4, 2021

Implements a Loop Fusion Transformation #493

Draft

2 tasks

kaushikcfd force-pushed the better_loop_around_nest_map branch from d74da04 to 1aef04f Compare October 8, 2021 21:45

kaushikcfd force-pushed the better_loop_around_nest_map branch from 1aef04f to 581e82e Compare October 22, 2021 16:20

kaushikcfd force-pushed the better_loop_around_nest_map branch 5 times, most recently from e20c6ca to 931c6e9 Compare December 10, 2021 04:50

kaushikcfd force-pushed the better_loop_around_nest_map branch from 931c6e9 to 36a1a2f Compare January 12, 2022 06:59

kaushikcfd force-pushed the better_loop_around_nest_map branch from f10f24c to c1f8e0d Compare January 28, 2022 15:53

kaushikcfd force-pushed the better_loop_around_nest_map branch from c1f8e0d to 8c42d10 Compare February 18, 2022 05:52

kaushikcfd force-pushed the better_loop_around_nest_map branch from 8c42d10 to 15cea20 Compare March 10, 2022 15:16

inducer mentioned this pull request Apr 23, 2022

[fusion_actx] High compile times illinois-ceesd/mirgecom#644

Open

kaushikcfd force-pushed the better_loop_around_nest_map branch from 15cea20 to 5b55bcd Compare May 6, 2022 17:30

kaushikcfd force-pushed the better_loop_around_nest_map branch from 5b55bcd to f8dc959 Compare May 16, 2022 13:58

kaushikcfd force-pushed the better_loop_around_nest_map branch from f8dc959 to 80425bd Compare June 9, 2022 20:09

inducer reviewed Jun 9, 2022

View reviewed changes

inducer reviewed Jun 10, 2022

View reviewed changes

kaushikcfd mentioned this pull request Jun 17, 2022

Implements loop distribution transformation #629

Draft

4 tasks

kaushikcfd force-pushed the better_loop_around_nest_map branch 2 times, most recently from 65a6b62 to c21c092 Compare July 6, 2022 17:49

This was referenced Oct 11, 2022

[Update] generate_loop_schedule_v2 #690

Draft

add Tree implementation #694

Open

kaushikcfd added 4 commits June 1, 2023 20:23

adds an immutable tree implementation

43dd810

add helpers to figure out loop nestings from a kernel

2c78a9c

adds loopy scheduler v2

0f478dd

changes in docs to account for equivalent generated codes from the sa…

fadb652

…me LoopKernel

inducer force-pushed the better_loop_around_nest_map branch from c21c092 to fadb652 Compare June 2, 2023 01:23

inducer mentioned this pull request Feb 16, 2024

Add test_long_kernel #826

Draft


		iname_to_insns = kernel.iname_to_insns()
		loop_nest_around_map = defaultdict(frozenset)

		Updates tree to make inames_to_pull_out a loop nesting level in
		loop_nests

		@@ -930,6 +933,225 @@ def _get_persistent_hashable_arg(arg):

		return wrapper


		# {{{ tree data structure

		_parent_to_children: "PMap[T, FrozenSet[T]]"
		_child_to_parent: "PMap[T, Optional[T]]"


		loop_nests = sorted(loop_nests, key=lambda nest: tree.depth(nest))

		for outer, inner in zip(loop_nests[:-1], loop_nests[1:]):

		Returns a copy of tree that realizes inames_to_pull_out as loop
		nesting.

generate_loop_schedule_v2 #350

Are you sure you want to change the base?

generate_loop_schedule_v2 #350

Conversation

kaushikcfd commented Apr 28, 2021 • edited Loading

inducer left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kaushikcfd Jun 10, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

inducer commented May 26, 2021

kaushikcfd commented May 26, 2021

inducer commented Jun 9, 2022

inducer left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

inducer Jun 9, 2022 • edited Loading

Choose a reason for hiding this comment

inducer Jun 9, 2022 • edited Loading

Choose a reason for hiding this comment

inducer left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

inducer commented Jun 10, 2022

kaushikcfd commented Apr 28, 2021 •

edited

Loading

kaushikcfd Jun 10, 2021 •

edited

Loading

inducer Jun 9, 2022 •

edited

Loading

inducer Jun 9, 2022 •

edited

Loading