Transform propagator skip replay when possible

[zasdfgbnm]

This comment in the code describes what this PR is doing:

  // Note: [Using multiple TransformPropagators]
  // There are cases that we use multiple TransformPropagators along different
  // spanning trees with different references in the same fusion. Some of these
  // spanning trees could overlap. In cases when there are overlapping nodes,
  // TransformPropagator needs to respect the replay of others, because the
  // current TransformPropagator might not contain the most amount of
  // information on how to do the correct transformation. The logic below tells
  // TransformPropagator to skip the replay when not necessary.

[zasdfgbnm]

This PR is rebased

[csarofeen]

I feel like I'm missing something here, why would new_pos be less than 0 here?

[zasdfgbnm]

This is by protocol. <0 means impossible to find matched position without replay, that is, replay is required. See comment:

  // Returns the leaf position in consumer that matches with `producer_pos` in
  // producer. Returns -1 if matching is impossible. This function can be used
  // to test if replay is needed for getting matching outer dims.

[csarofeen]

Yeah, this was what was confusing me. We're not returning the actual mismatched location, we're just returning the provided position if the mismatched position is further to the right. So this function can only return the provided position, or -1 meaning there's a mismatch before we hit the provided position.

[csarofeen]

We should either return this function to just being a bool, or we should actually return the position of the mismatch and compare that position outside of this function.

[zasdfgbnm]

we're just returning the provided position if the mismatched position is further to the right.

I think we are not just returning the provided position, but instead, if you provide a consumer position, we are returning the corresponding producer position of the given consumer position. In this case, we can not return the actual mismatch position, because it does not provide the information about the corresponding producer position of the given consumer position?

I copy-pasted this function from computeAt, and actually, there is one thing that I don't understand about this function. Why are we skipping unmappable dims in the consumer but not in the producer? What breaks the symmetry?

[csarofeen]

I copy-pasted this function from computeAt

Yeah it's my fault, this function is really just returning yes or no, it shouldn't be returning an int. It could return an int, but should be rewritten to do so.

Why are we skipping unmappable dims in the consumer but not in the producer? What breaks the symmetry?
Unmappable dims (might be a bad name) come from patterns associated with reductions, and really associated with normalization patterns.

If I have:
T1 = set(T0)
T2 = sum(T1, {1})
T3 = broadcast(T2, {false, true})
T4 = add(T1, T3)

T1 cannot be fully inlined to T4. So T1's second dimension (I believe) is marked as unmappable to T4. It's not okay to fully inline T2 into T4, but it's fine to fully inline T1 with T2.

[csarofeen]

@naoyam can tell me if I'm wrong on the details of the example, but I believe the principles stand.

[zasdfgbnm]

it shouldn't be returning an int. It could return an int, but should be rewritten to do so.

I still don't understand. Are you saying that the algorithm we used to compute the corresponding position is wrong, and we should rewrite to compute it differently? Or are you saying that the algorithm is OK, but we should split up the testing of "needs play" and "find the corresponding position" into two things? Or something else?

If I have:
T1 = set(T0)
T2 = sum(T1, {1})
T3 = broadcast(T2, {false, true})
T4 = add(T1, T3)

Hmm, I am not sure if this example is a related. I think getMatchedLeafPosWithoutReplay only looks at PairwiseRootDomainMap, which will map all the dims in T4 to T1?

TEST_F(NVFuserTest, TMP) {
  auto fusion = std::make_unique<Fusion>();
  FusionGuard fg(fusion.get());

  auto tv0 = makeSymbolicTensor(2);
  fusion->addInput(tv0);
  auto tv1 = set(tv0);
  auto tv2 = sum(tv1, {1});
  auto tv3 = broadcast(tv2, {false, true});
  auto tv4 = add(tv1, tv3);
  fusion->addOutput(tv4);

  auto producer = tv1;
  auto consumer = tv4;

  const auto c2p_root_map =
      PairwiseRootDomainMap(producer, consumer)
          .mapConsumerToProducer(consumer->domain(), producer->domain());

  // IterDomains in consumer root also in producer root
  std::unordered_set<Val*> mapped_consumer_roots;
  for (auto entry : c2p_root_map) {
    mapped_consumer_roots.emplace(entry.first);
  }

  fusion->print();
  std::cout << ir_utils::toString(mapped_consumer_roots) << std::endl;
}

This gives iS8{i1}, iS9{i2}.

I think "unmapped IDs in consumer but not in producer" here refers to new broadcasting dims, and we are saying that, if we see a leaf ID that completely comes from new broadcasting dims, then we can ignore it. But why don't we want to symmetrically do the same for reductions in the producer? For non-trivial reduction, I think this is because it could not be inlined, but for trivial reduction, it should be safe to just ignore it as well?

[csarofeen]

I believe the algorithm is correct, and it should not be symmetric. I was considering trivial reduction, but trivial reduction could likely be ignored as well.

I was trying to say the asymmetry is due to the unmappable dims, which are related to particular graph patterns with reduction (the pattern mentioned above).

[csarofeen]

This is my mistake, I misread the return statements, I thought we were checking the return value was equal to the provided value to return that value. We're swapping produce/consume so they're different.

[zasdfgbnm]

The algorithm is actually not correct😉: #1791

[naoyam]

it shouldn't be returning an int. It could return an int, but should be rewritten to do so.

I still don't understand. Are you saying that the algorithm we used to compute the corresponding position is wrong, and we should rewrite to compute it differently? Or are you saying that the algorithm is OK, but we should split up the testing of "needs play" and "find the corresponding position" into two things? Or something else?

If I have:
T1 = set(T0)
T2 = sum(T1, {1})
T3 = broadcast(T2, {false, true})
T4 = add(T1, T3)

Hmm, I am not sure if this example is a related. I think getMatchedLeafPosWithoutReplay only looks at PairwiseRootDomainMap, which will map all the dims in T4 to T1?

TEST_F(NVFuserTest, TMP) {
  auto fusion = std::make_unique<Fusion>();
  FusionGuard fg(fusion.get());

  auto tv0 = makeSymbolicTensor(2);
  fusion->addInput(tv0);
  auto tv1 = set(tv0);
  auto tv2 = sum(tv1, {1});
  auto tv3 = broadcast(tv2, {false, true});
  auto tv4 = add(tv1, tv3);
  fusion->addOutput(tv4);

  auto producer = tv1;
  auto consumer = tv4;

  const auto c2p_root_map =
      PairwiseRootDomainMap(producer, consumer)
          .mapConsumerToProducer(consumer->domain(), producer->domain());

  // IterDomains in consumer root also in producer root
  std::unordered_set<Val*> mapped_consumer_roots;
  for (auto entry : c2p_root_map) {
    mapped_consumer_roots.emplace(entry.first);
  }

  fusion->print();
  std::cout << ir_utils::toString(mapped_consumer_roots) << std::endl;
}

This gives iS8{i1}, iS9{i2}.

I think "unmapped IDs in consumer but not in producer" here refers to new broadcasting dims, and we are saying that, if we see a leaf ID that completely comes from new broadcasting dims, then we can ignore it. But why don't we want to symmetrically do the same for reductions in the producer? For non-trivial reduction, I think this is because it could not be inlined, but for trivial reduction, it should be safe to just ignore it as well?

The constraint about reductions not being able to get inlined is reflected in ComputeAtRootDomainMap. PairwiseRootDomainMap does not consider constraints due to inlining but just looks at a pair of a producer and consumer.

[csarofeen]

LGTM

[naoyam]

Please add comments on the parameters and the return value.

[zasdfgbnm]

I agree that these functions need a better doc. I will work on it in a new PR.

[naoyam]

We often use names like consumer_pos to indicate a position in a consumer domain, so this could be confusing. Maybe something like is_producer_as_consumer?

[zasdfgbnm]

This part of the code has been replace by #1791, no more such issue 😉

[naoyam]

it shouldn't be returning an int. It could return an int, but should be rewritten to do so.

I still don't understand. Are you saying that the algorithm we used to compute the corresponding position is wrong, and we should rewrite to compute it differently? Or are you saying that the algorithm is OK, but we should split up the testing of "needs play" and "find the corresponding position" into two things? Or something else?

If I have:
T1 = set(T0)
T2 = sum(T1, {1})
T3 = broadcast(T2, {false, true})
T4 = add(T1, T3)

Hmm, I am not sure if this example is a related. I think getMatchedLeafPosWithoutReplay only looks at PairwiseRootDomainMap, which will map all the dims in T4 to T1?

TEST_F(NVFuserTest, TMP) {
  auto fusion = std::make_unique<Fusion>();
  FusionGuard fg(fusion.get());

  auto tv0 = makeSymbolicTensor(2);
  fusion->addInput(tv0);
  auto tv1 = set(tv0);
  auto tv2 = sum(tv1, {1});
  auto tv3 = broadcast(tv2, {false, true});
  auto tv4 = add(tv1, tv3);
  fusion->addOutput(tv4);

  auto producer = tv1;
  auto consumer = tv4;

  const auto c2p_root_map =
      PairwiseRootDomainMap(producer, consumer)
          .mapConsumerToProducer(consumer->domain(), producer->domain());

  // IterDomains in consumer root also in producer root
  std::unordered_set<Val*> mapped_consumer_roots;
  for (auto entry : c2p_root_map) {
    mapped_consumer_roots.emplace(entry.first);
  }

  fusion->print();
  std::cout << ir_utils::toString(mapped_consumer_roots) << std::endl;
}

This gives iS8{i1}, iS9{i2}.

I think "unmapped IDs in consumer but not in producer" here refers to new broadcasting dims, and we are saying that, if we see a leaf ID that completely comes from new broadcasting dims, then we can ignore it. But why don't we want to symmetrically do the same for reductions in the producer? For non-trivial reduction, I think this is because it could not be inlined, but for trivial reduction, it should be safe to just ignore it as well?

The constraint about reductions not being able to get inlined is reflected in ComputeAtRootDomainMap. PairwiseRootDomainMap does not consider constraints due to inlining but just looks at a pair of a producer and consumer.