WIP: IterDomain Graphs #32

csarofeen · 2023-03-19T16:10:12Z

Build out of Iter domain graphs as infrastructure. I kept these concepts separate from the compute_at_map as that may need to be reimplemented later based on this new concept/infrastructure.

This new concept of IterDomainGraphs will eventually replace all our index and parallelization logic. This infrastructure is to make it easier to work with iter domain graphs for processes like accurate broadcast resolution/promotion. IdGraph or IterDomainGraphs could also directly replace BestEffortReplay and similar mappings across producer-consumers.

I added a couple interesting tests that I want to make work but still don't.

csarofeen · 2023-03-19T21:18:54Z

Helps if I include the right files.

jacobhinkle

Started reviewing. I haven't gone through all of buildLoopPromotionMap yet.

csrc/disjoint_set.h

jacobhinkle · 2023-03-21T19:39:28Z

csrc/id_graphs.h

+  //     (1) The disjoint set of the provided Iter Domain if it exists,
+  //     otherwise a null shared ptr
+  //     (2) If the disjoint set of the provided Iter Domain exists


How does this differ from returning a std::optional<IdGroup>?

Returning an optional probably makes a lot of sense, but actually I think most instances I just want to assert there's actually an id set. So probably I'll just put an alias.

I just stole the pattern for https://cplusplus.com/reference/unordered_set/unordered_set/emplace/ for development, no reason not to switch to optional.

Ah OK. Makes sense. I think std::optional was introduced in C++17, but std::unordered_set has been around since C++11 so they had to use the explicit pair<..., bool> . From side discussion it looks like we might be OK with using C++17 going forward on NVFuser. Either way not a big deal of course.

jacobhinkle · 2023-03-21T20:03:37Z

csrc/id_graphs.h

+  //! all IterDomains in the disjoint set to that PType.
+  void validateAndPropagatePType() const;
+
+  void buildLoopPromotionMap(const std::vector<Expr*>& exprs);


Basic question, but I see "promotion" mentioned many times here. In this context what does it mean to promote an ID?

Promotion is the concept that an iteration domain that a TensorView has in its root->domain might not really be what's required for the generated kernel to index into that TensorView. Promotion is for example:

Producer has a broadcast merged with an iteration domain

Consumer has a (mapped to the producer) iteration domain merged with an iteration domain

Based on other transformations the producer might have to "promote" it's broadcast domain to the iteration domain in it's consumer

If there's a producer of that producer, then we still might need the "broadcast promotion" but there isn't a broadcast that maps in that producer's producer

The "end goal" of promotion in this pass (still very WIP), is that each leaf iter domain of a tensor view might be "promoted" to a larger iteration domain representative of the for loops. That larger iter domain still needs connections that we can traverse to index into the tensor view's buffer.

Revisiting this after some thought and going through some of the machinery. Just to try and reiterate the example above:

t0[(b0*i1), i2] // producer t1[(i3 * i4)] = f(t0) // consumer, through expression f

i3 is defined by an expression in the variables i1, i2.

We may need to alter b0 to match i3 if i4 matches i1.

/This might cascade, since (b0*i1) might match another merged broadcast in its producer.

End goal of promotion: map each leaf IterDomain in each tensorview (importantly including bcast domains and transforms thereof) to an Iteration IterDomain that is written as a transform of that TVs root_domain, so that we can index into it.

naoyam

Took a quick look, but that's not enough to make any meaningful feedback. I'd need to spend a significant amount of time (i.e., a few days). Should I wait? Is it ready?

csrc/id_graphs.h

csarofeen · 2023-03-22T13:20:02Z

@naoyam I think it's worth trying to read through, I wouldn't worry about really nailing down interfaces, but the building and relationship of an Iter Domain Graph, and how we operate on a collection of Iter Domain Graphs. The infrastructure is here, but buildLoopPromotionMap is the real challenge. That's the function where I'm trying to resolve all the broadcasts in a way we can write index pass directly on the graph iter domain graphs, or really (hopefully) on an iter domain graph.

csarofeen · 2023-03-22T13:20:23Z

i.e. buildLoopPromotionMap is something you could initially ignore.

naoyam · 2023-03-22T16:48:59Z

(I'm sorry I accidentally clicked the close button)

naoyam · 2023-03-22T16:50:44Z

i.e. buildLoopPromotionMap is something you could initially ignore.

Good to know. That's where I intended to focus on as I think that's one of the main tasks in this PR. I'll wait further updates for the function. Will look though the rest.

csarofeen · 2023-03-23T13:49:40Z

@naoyam you could start looking at buildLoopPromotionMap with all the debug spew I have. Interesting cases in my opinion are:
FusionInlineBroadcastIndexing0_CUDA
FusionIndexing19_CUDA
FusionIndexSplitMerge_CUDA

I still need to make a backward replay before we try to perform indexing, as I don't want to index all the tensor views one by one, so instead I intend to build a graph that we can naively traverse all at once to get consumer indices.

…nGraphs

naoyam · 2023-03-28T22:36:02Z

(Sorry again accidentally clicked the close button)

FusionInlineBroadcastIndexing0_CUDA
FusionIndexing19_CUDA
FusionIndexSplitMerge_CUDA

All these tests resulted in a failure. Is it expected?

naoyam · 2023-03-28T23:05:05Z

(Sorry again accidentally clicked the close button)

FusionInlineBroadcastIndexing0_CUDA
FusionIndexing19_CUDA
FusionIndexSplitMerge_CUDA

All these tests resulted in a failure. Is it expected?

It seems expected as there's TORCH_INERNAL_ASSERT(false) at the end of buldLoopPromotionMap

naoyam

Comments on IdGraph

csrc/id_graphs.cpp

csarofeen · 2023-03-29T12:57:51Z

(Sorry again accidentally clicked the close button)

FusionInlineBroadcastIndexing0_CUDA
FusionIndexing19_CUDA
FusionIndexSplitMerge_CUDA

All these tests resulted in a failure. Is it expected?

Yes, indexing is not hooked up, so they're not yet supported. I'm throwing an error where I leave off in the analysis.

csrc/id_graphs.cpp

naoyam

Comments so far on buildLoopPromotionMap

csrc/id_graphs.cpp

naoyam · 2023-03-30T16:13:08Z

csrc/id_graphs.cpp

+      VectorOfUniqueEntries<IterDomain*> all_producer_ca_deps;
+      {
+        auto ca_dep_vals = DependencyCheck::getAllValsBetween(
+            {producer_root.begin(), producer_root.end()},
+            {producer_domain.begin(),
+             producer_domain.begin() + producer->getComputeAtPosition()});
+        auto ca_deps_filter = ir_utils::filterByType<IterDomain>(ca_dep_vals);
+
+        all_producer_ca_deps.insert(
+            ca_deps_filter.begin(), ca_deps_filter.end());
+      }


nit: This pattern appears quite often. We should create a utility function.

Probably. Something that goes directly from...

DependencyCheck::getAllValsBetween( {producer_root.begin(), producer_root.end()}, {producer_domain.begin(), producer_domain.begin() + producer->getComputeAtPosition()})

To:

VectorOfUniqueEntries<IterDomain*> all_producer_ca_deps;

?

naoyam · 2023-03-30T16:30:50Z

csrc/id_graphs.cpp

+  }
+}
+
+const IdGraph& IterDomainGraphs::idGraph(IdMappingMode mode) const {


Can this assert if the returned graph is already constructed? For example, the LOOP map is not available before lowering, so it would be nice if we could assert it's accidentally queried.

Sure, do you think the check in replay as is enough to do this effectively or do you think that we should have some flag in IterDomainGraphs marking when each mode is initialized?

The latter seems to make more sense to me.

csrc/id_graphs.cpp

naoyam · 2023-03-30T16:42:29Z

csrc/id_graphs.cpp

+
+  buildPermissiveMap(tv_exprs);
+
+  // Only build loop map during lowering


We may also sometimes just want to build only an exact and/or permissive map. The scheduler is one example where we sometimes use an exact map only. Maybe a permissive map is also used anyway?

I hope it's not too expensive, but I'd be happy to have finer granularity on the building of the graphs. Being able to specify which are needed, or generating the maps lazily could also be cool. Leaving to future work for now.

csrc/id_graphs.cpp

naoyam · 2023-03-30T18:00:18Z

csrc/id_graphs.cpp

+       idGraph(IdMappingMode::LOOP).disjointIdSets().disjointSets()) {
+    if (group->size() == 1) {
+      p2c_ca_terminal_loop_ids.pushBack(group->front());
+      id_consumer_terminal_loop_ids.pushBack(group->front());


nit: insert continue

Yeah, I need to clean this up.

naoyam · 2023-03-30T22:45:46Z

csrc/id_graphs.cpp

+  // T4 = T1[i0, b1] + T3[i0, i1]
+  // T6 = T2[i0, b1] + T5[i0, i2]
+  //
+  // The almost exact map will map T1's and T2's b1 together, but they're being


Can you elaborate a little more why they are mapped? Are these domains merged?

Yes, thank you, this assumes merge(i0, b1), merge(i0, i1), and merge(i0, i2)

Then merge(i0, b1) of T1 and merge(i0, b1) of T2 are almost exact mapped together from id graph propagation.

Let me rewrite the expressions as below. I think this is more accurate.

T1[i0, b1] = T0[i0] T2[i0, b2] = T0[i0] // Not T2[i0, b1] T4 = T1[i0, b1] + T3[i0, i1] T6 = T2[i0, b2] + T5[i0, i2]

Then merge(0, 1) with all tensors except for T0.

The almost-exact map would map i0 and i0*b1 and i0*b2. Does it also map b1 and b2?

I'm not sure, if it maps b1 and b2 together. Should be benign if so.

csrc/id_graphs.cpp

naoyam · 2024-05-21T04:11:34Z

!build

naoyam · 2024-06-04T18:56:54Z

Closing this PR. Most of the code was already merged into main, and the remaining code is currently not used and will be revisited it if necessary.

The next work is indexing, which is tracked in #2238.

Renamed test_gpu_indexing.cpp to test_indexing_advanced.cpp. Changed the tests to exercise both the legacy and new indexers. Added several tests originally developed for IdModel (#32). Some of them are disabled as they are not yet supported.

Builds and can start debugging again.

16f064e

csarofeen requested review from jacobhinkle and naoyam March 19, 2023 16:10

Helps if I include the right files.

45959ef

jacobhinkle reviewed Mar 21, 2023

View reviewed changes

naoyam reviewed Mar 21, 2023

View reviewed changes

csrc/id_graphs.h Outdated Show resolved Hide resolved

naoyam closed this Mar 22, 2023

naoyam reopened this Mar 22, 2023

Forward replay now works, need to implement backward index replay.

eee9bf7

csarofeen added 3 commits March 25, 2023 10:10

Merge branch 'main' of https://github.com/nvidia/fuser into IterDomai…

8ce0963

…nGraphs

Finish computing loop group promotions.

f745410

Prepare for backward graph replays.

6edbba3

naoyam closed this Mar 28, 2023

naoyam reopened this Mar 28, 2023

naoyam reviewed Mar 29, 2023

View reviewed changes

csrc/id_graphs.cpp Outdated Show resolved Hide resolved

csrc/id_graphs.cpp Outdated Show resolved Hide resolved

csrc/id_graphs.cpp Outdated Show resolved Hide resolved

naoyam reviewed Mar 30, 2023

View reviewed changes

csrc/id_graphs.cpp Outdated Show resolved Hide resolved

naoyam reviewed Mar 30, 2023

View reviewed changes

naoyam reviewed Mar 31, 2023

View reviewed changes

csrc/id_graphs.cpp Outdated Show resolved Hide resolved

csarofeen commented Apr 1, 2023

View reviewed changes

csrc/id_graphs.cpp Outdated Show resolved Hide resolved

Build promoted tensor domains.

1a2b261

naoyam added 15 commits May 10, 2024 09:24

Merge branch 'idmodel_step5' into IterDomainGraphs

0027ee9

Merge branch 'main' into IterDomainGraphs

2bcbcb4

LoopPromotionBuilder class

fa8455a

Copied all loop promotion code

c05fd85

enable idmodel

31d28e5

Switch to the new builder

d0bdc8e

const

aa9b414

cleanup

7aef29d

Remove loop promotion code from IdModel

24a0c8c

replace tester with callback

c7c04b1

clang-format

2b605cd

comment

0f5ab07

Merge branch 'idmodel_loop_promotion_cleanup' into IterDomainGraphs

423fde1

Merge branch 'main' into IterDomainGraphs

c78da86

cleanup

f0aeab7

naoyam added 7 commits May 21, 2024 15:00

Merge branch 'main' into IterDomainGraphs

ce5f8ee

Merge branch 'main' into IterDomainGraphs

3214bc7

Merge branch 'main' into IterDomainGraphs

d8aacea

enable

18de523

Merge branch 'main' into IterDomainGraphs

83edc44

Merge branch 'main' into IterDomainGraphs

e71cb5a

Merge branch 'main' into IterDomainGraphs

009a11a

naoyam closed this Jun 4, 2024

This was referenced Jun 6, 2024

Squeezed IterDomain ?S536{1} must concretize to IterType::Broadcast but found ?S536{1}. #2359

Closed

Merging IterDomains requires that their iteration types match. #2317

Closed

naoyam mentioned this pull request Sep 26, 2024

Indexing test cleanup #3028

Merged

wujingyue mentioned this pull request Oct 18, 2024

OpInfo has problems testing define_tensor. #3225

Closed


		buildPermissiveMap(tv_exprs);

		// Only build loop map during lowering

WIP: IterDomain Graphs #32

WIP: IterDomain Graphs #32

Conversation

csarofeen commented Mar 19, 2023 • edited Loading

csarofeen commented Mar 19, 2023

jacobhinkle 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

csarofeen Mar 21, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

naoyam left a comment

Choose a reason for hiding this comment

csarofeen commented Mar 22, 2023

csarofeen commented Mar 22, 2023

naoyam commented Mar 22, 2023

naoyam commented Mar 22, 2023

csarofeen commented Mar 23, 2023

naoyam commented Mar 28, 2023

naoyam commented Mar 28, 2023

naoyam left a comment

Choose a reason for hiding this comment

csarofeen commented Mar 29, 2023

naoyam 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

csarofeen Apr 1, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

csarofeen May 20, 2023 • edited Loading

Choose a reason for hiding this comment

naoyam commented May 21, 2024

naoyam commented Jun 4, 2024

csarofeen commented Mar 19, 2023 •

edited

Loading

csarofeen Mar 21, 2023 •

edited

Loading

csarofeen Apr 1, 2023 •

edited

Loading

csarofeen May 20, 2023 •

edited

Loading