Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Rework HaloInfo interface #1987

Merged
merged 1 commit into from
Sep 16, 2022
Merged

Rework HaloInfo interface #1987

merged 1 commit into from
Sep 16, 2022

Conversation

csarofeen
Copy link
Owner

@csarofeen csarofeen commented Sep 16, 2022

Trying to rework the HaloInfo interface a bit. When working on contiguity.h I realized it was challenging to separate HaloInfo from GpuLower. I don't know if there's value in being able to support:
https://github.com/csarofeen/pytorch/blob/devel/torch/csrc/jit/codegen/cuda/contiguity.cpp#L52-L55
when we aren't in the lowering process. However, the only thing HaloInfo really takes to construct is the compute at map and the fusion. So I separated out the reliance on GpuLower in HaloInfo. I also realized there aren't any accesses to HaloInfo that are non-constant, so I removed the non-const interfaces from the public API.

The other thing I wanted to do was to try using shared_ptr to pass the information around just to be a bit safer on lifetime and ownership. I'm generally concerned with the way we use GpuLower we could be holding on to a reference that's not really safe to hold onto since our classes don't have ownership, GpuLower does. I don't think we have much chance to suffer from this today, but I think the pattern is safer.

The other thing I like about using a shared_ptr is it's easier to use a pattern I'm often looking for which is "Take this reference, or construct a new object." Specifically we sometimes share analysis tools in lowering when we're doing something like scheduling. At scheduling we either might be creating a new ComputeAtMap, or we might be using one that exists. So if we have a class that requires a ComputeAtMap we could either pass in a shared pointer, or create a new one.

Copy link
Collaborator

@naoyam naoyam left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Thanks for the cleanup.

void HaloInfo::build(Fusion* fusion) {
HaloInfo::HaloInfo(Fusion* fusion, std::shared_ptr<const ComputeAtMap> ca_map)
// Make a copy of the permissive map for extent comparators
: permissive_map_(ca_map->idGraph().permissiveNodes()) {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are you trying to decouple the states held by GpuLower from each other? I like it as that seems like a better and more robust design.

Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I'm trying to pull the dependencies apart a bit as I was thinking of building a HaloInfo structure in contiguity which is also used outside GpuLower.

@@ -801,7 +796,8 @@ std::unordered_map<IterDomain*, Val*> HaloInfo::buildConcreteHaloExtentMap(
merged_shifted_ids.insert(ir_utils::caMapExactConcreteId(merge->out()));
// Note that halo_width_map_ is not updated
} else {
setHaloWidth(ir_utils::caMapExactConcreteId(merge->out()), 0);
local_halo_info.setHaloWidth(
ir_utils::caMapExactConcreteId(merge->out()), 0);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Regarding the decoupling, here we still use the CA map held by GpuLower. And, also just realized it's a little counter-intuitive that ir_utils implicitly uses GpuLower. It doesn't seem trivial to remove this interface as it's used everywhere.

Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is making this local_halo_info right?

Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will look again at the inter-class dependencies.

Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, this class is really hard set to the lowering process as it's taking indexing in. That's why I didn't bother looking more closely.

Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ir_utils::caMapExactConcreteId seems like a weak alias, will get rid of that either before this is merged or in another PR.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Creating local_halo_info is right. It was part of Shiming's indexing refactoring. Previously, reference domains are added to the singleton instance of HaloInfo, but as we don't create reference IterDomains anymore, the properties generated for reference IterDomains are instead generated for concrete IterDomains. A complexity about halo is that even exact mapped IterDomains may have different halo sizes, so it's not possible to associate a unique halo attribute with a group of mapped IterDomains. As the halo size information depends on the loop nest where an expression appears, what Shiming did was to clone the pre-built HaloInfo and modify the halo information for the used concrete IterDomains at that loop nest.

That's my understanding from my reading of his refactoring PR. Tagging @shmsong as he may want to chime in.

Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I was assuming it was just a typo.

@csarofeen csarofeen merged commit eb1dad1 into devel Sep 16, 2022
jjsjann123 added a commit that referenced this pull request Nov 9, 2022
Syncing nvfuser devel branch to upstream master. https://github.com/csarofeen/pytorch/

Codegen changes include:

* codegen improvement:
    i. allow non-root trivial reductions, allow empty/no-op fusion
    ii. fixes vectorization checks and size calculation
    iii. bank conflict handle improvement
    iv. enables transpose scheduler

* misc:
    i. CI tests failure fixes
    ii. cpp tests file clean up
    iii. trivial forwarding supports added in codegen runtime
    iv. added factory methods support in codegen

Commits that's in this PR from the devel branch:

```
7117a7e patching nvfuser conv cudnn test numerics mismatch (#2048)
65af1a4 Inserting sync for redundant parallel types is already done at the (#2023)
6ac74d1 Fix sync map (#2047)
f5bca33 Bank conflict checker improvements (#2032)
d2ca7e3 Minor update on cp.async code generation. (#1901)
d36cf61 Test file cleanup (#2040)
0b8e83f Allow non-root trivial reductions (#2037)
a2dfe40 Fix vectorize size calculation (#2035)
e040676 Use withPredicate to replace setPredicate to maintain Exprs immutable (#2025)
197221b removing ci workflow (#2034)
40e2703 Reduction rand like patch (#2031)
bc77266 Add utility for checking bank conflict of shared memory (#2029)
ddd1cf7 Add back FusionReductionWithTrivialReduction_CUDA (#2030)
fbd97e5 Revert "Cleanup trivial reduction workarounds (#2006)" (#2024)
bca20c1 Cleanup trivial reduction workarounds (#2006)
e4b6585 Trivial forwarding (#1995)
1a0e355 Fix contiguity analysis of predicates to match updated contiguity. (#1991)
a4effa6 Enable output allocation cache (#2010)
35440b7 Patching bn inference (#2016)
0f9f0b4 Add matmul benchmark (#2007)
45045cd Enable tests previously disabled due to an aliasing bug (#2005)
967aa77 Contiguous indexing for View operations (#1990)
a43cb20 Make inlining even more modular (#2004)
dc45835 Test util cleanup (#2003)
3ca21eb More strict validation (#2000)
a7a7d57 Fix build problem (#1999)
fc235b0 Just fixes comments (#1998)
482386c cleanup (#1997)
4cbe0db Improve divisible split detection (#1970)
42ccc52 Minor build fix. (#1996)
fcf8c09 Cleanup of lower_utils.cpp: Isolate out GpuLower usage (#1989)
15f2f6d Move ConcretizedBroadcastDomains to shared_ptr in GpuLower. (#1988)
8f1c7f5 Minor cleanup lower_unroll.cpp (#1994)
1d9858c Minor cleanup (#1992)
f262d9c Add support for uniform RNG (#1986)
eb1dad1 Remove non-const functions, remove GpuLower instance on build, pass in ca_map. (#1987)
634820c Add support for some empty fusion (#1981)
eabe8d8 Segment self mapping fusions (#1954)
e96aacf Enable Transpose operation (#1882)
425dce2 Add a null scheduler that helps segmenting away no-op schedules (#1835)
306d4a6 Fix canScheduleCompileTime check of transpose scheduler (#1969)
b1bd32c Minor fix (#1967)
bd93578 Enable transpose scheduler (#1927)
b7a206e Move scheduler vectorize utilities into their own file (#1959)
d9420e4 View scheduling (#1928)
c668e13 Upstream push ci fixes (#1965)
c40202b Fix dump effective bandwidth (#1962)
93505bc WAR on index mapping when exact and permissive maps differ (#1960)
45e95fd Allow splitting inner-most ID to create virtual innermost ID in transpose scheduler (#1930)
a3ecb33 Improve the comments at the beginning of index_compute.h (#1946)
f7bc341 Remove unused variables (#1955)
df3393a Some cleanup (#1957)
7d1d7c8 TVDomainGuard factory (#1953)
357ba22 Fill allocation with nan on tests (#1956)
8eafc54 Fix detection of unmappable root domains (#1952)
90a51f2 Some indexing cleanups, Add eye support (#1940)
ddc01e4 Exclude unsupported data types (#1951)
992e17c test the groups the same order as they are merged (#1949)
208262b Move detection of self mapping IDs to IterDomainGraph from (#1941)
ac4de38 Merge pull request #1945 from csarofeen/master_merge_0828
6310948 Add full, full_like, zeros, zeros_like, ones, ones_like (#1943)
aab10bc Merge remote-tracking branch 'upstream/viable/strict' into HEAD
4c254c0 Fix arange when step is negative (#1942)
89330aa Tensor factories must set the output shape as its input (#1939)
```

RUN_TORCHBENCH: nvfuser

Differential Revision: [D40869846](https://our.internmc.facebook.com/intern/diff/D40869846)
Pull Request resolved: pytorch#87779
Approved by: https://github.com/davidberard98
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants