eraseInputDistinctRootDomains supports general logical-to-allocation transforms

[wujingyue]

This is a spin-off from #3444.

The current code assumes that logical-to-allocation has to be a permutation. This assumption won't hold any more with #2563. So this PR tries to extend eraseInputDistinctRootDomains to support more general transforms.

This can happen to single-GPU, although not as common. The tests added in this PR are for single-GPU because #3444 hasn't landed. #3444 will add some multi-GPU tests.

[wujingyue]

!test

[wujingyue]

@jacobhinkle FYI

[naoyam]

I'm not sure why this could cause that problem. It may be just because of the concern I mentioned below.

[wujingyue]

This issue persisted after I moved the split after addInput/addOutput. FYI, below is my failed attempt to fix dynamic transform for this test.

diff --git a/csrc/dynamic_transform.cpp b/csrc/dynamic_transform.cpp
index 24404db8..d287149d 100644
--- a/csrc/dynamic_transform.cpp
+++ b/csrc/dynamic_transform.cpp
@@ -1056,7 +1056,7 @@ void DynamicTransformConcretizer::mutate(TensorView* tv) {
     // beyond the logical domain as asserted above
     auto all_id_exprs = StmtSort::getExprsBetween(
         {tv->getRootDomain().begin(), tv->getRootDomain().end()},
-        {tv->getLogicalDomain().begin(), tv->getLogicalDomain().end()});
+        {tv->getMaybeAllocationDomain().begin(), tv->getMaybeAllocationDomain().end()});
     for (auto expr : all_id_exprs) {
       // Assume outputs of IterDomain exprs are always IterDomains. If
       // the assumption is invalidated, the logic here would need to

[jacobhinkle]

Previously we had assumed there were no loop or allocation transforms so that each of those would just be permutations of logical. I think what we should do is actually use IRBFS here to propagate from root to all of the other domains (logical, loop, and allocation). IRBFS is not needed actually since we can safely assume at concretization that we don't have an uncommon situation like a loop domain that's a producer for transforms that lead to the root, instead of the other direction. If we assume that root is a producer for all of the other domains we can just use StmtSort::getExprsBetween like above, but we need to add not just the logical or allocation, but all three other domains as the "to" argument. i.e.

std::unordered_set<Val*> to{tv->getLogicalDomain().begin(), tv->getLogicalDomain().end()};
to.insert(tv->getMaybeAllocationDomain().begin(), tv->getMaybeAllocationDomain().end());
to.insert(tv->getLoopDomain().begin(), tv->getLoopDomain().end());
auto all_id_exprs =
    StmtSort::getExprsBetween({tv->getRootDomain().begin(), tv->getRootDomain().end()}, to);

[wujingyue]

Did you try it? While I agree with what you said, I doubt that helps this particular test case where root->maybeallocation includes all expressions. Btw, I believe

Fuser/csrc/ir/nodes.cpp

Line 3762 in b5e5182

std::vector<Expr*> TensorDomain::allExprs() const {

can be used to capture all Exprs in a TensorView.

[jacobhinkle]

Yes, the issue here is that allocation and logical are disconnected because there's been a replacement performed on logical, but not on allocation. I think the issue is that StmtSort::getStmts is only giving us the logical domains of the input TVs whereas we should expect it to provide all IDs for processing before the input TensorDomain and the TV itself.

[naoyam]

Ha, looks like it's actually just loop domains.

https://github.com/NVIDIA/Fuser/blob/main/csrc/iter_visitor.cpp#L73

We should probably change this to visit all of root, logical and loop domains.

[jacobhinkle]

In this particular case though, loop is equal to allocation. I originally thought this was the issue here. I agree that that line is an issue if we have unconventional tensor domains like loops that are producers of logical, and we should ensure all the domains are available there, but I think in this case the bad behavior is specific to input TVs.

[naoyam]

@wujingyue Is this a blocker?

[wujingyue]

Not for this PR. I can stick with concrete sizes for the tests for now

[wujingyue]

!test

[wujingyue]

!test

[wujingyue]

!test

[naoyam]

LGTM

[wujingyue]

!test