Add initial_loop_domain_ to TensorDomain

[naoyam]

It turned out it's necessary to remember the initial loop domain, which is either the logical domain set by the constructor or the new loop domain set by setLoopDomain. When a loop domain has an extra ID, TensorDomain::allIDs may miss IDs that solely depend on the extra ID.

[naoyam]

!build

[csarofeen]

When a loop domain has an extra ID

When does that happen?

[naoyam]

I'm still prototyping, but, for example:

t0 = [i0, i1] // input
t1 = [i0] // input

t2 = t0[:, n:] // slice t0 to [i0, i2] 
t3 = broadcast(t1, {false, true}) // [i0, b3]
t4 = t2 + t3 // [i0, i2]

The producer-based scheduling would use the loop domain of t0, [i0, i1], as the reference loop domain. For t1, I think it'd be the most natural to create a new broadcast ID and set its loop domain as [i0, b4]. This b4 could be also exactly mapped with b3.

[zasdfgbnm]

Should we just append everything in loop_domain_ here to additional_ids_? In allIDs, we might want to do:

for (auto i : c10::irange(all_domains.size() - 1)) {
  for (auto j : c10::irange(all_domains.size() - 1)) {

instead of

for (auto i : c10::irange(all_domains.size() - 1)) {
  for (auto j : c10::irange(i + 1, all_domains.size())) {

[naoyam]

I did think about it, but since additional_ids_ are just "additional", it seems to me it could be useful in some cases. For example, we could cache all IDs between logical and initial_loop to speed up TensorDomain::allIDs.

[zasdfgbnm]

additional_ids_ was the mechanism designed to allow loop domain to have a new broadcast, and the reason why it exists was to solve almost exactly the same problem as this PR solves. Moving forward, looks like we will not continue with the "allow loop domain to have a new broadcast" direction. But I don't want to have two mechanisms for the same thing. We should pick either of them, but only one should be picked. It makes sense to speed up TensorDomain::allIDs, and it might also make sense to call the data structure we use for the speedup additional_ids_, but it will be a new thing that is totally unrelated to today's additional_ids_ except they happen to have the same name.

[zasdfgbnm]

Which mechanism to pick I believe is mostly a question of interface design. Today, we have:

void TensorDomain::broadcast(int64_t axis, Val* extent) {
  axis = nvfuser::wrapDim(axis, nDims() + 1);
  IterDomain* id = IterDomainBuilder(fusion()->zeroVal(), extent)
                       .iter_type(IterType::Broadcast)
                       .build();
  loop_domain_.insert(loop_domain_.begin() + axis, id);
  additional_ids_.push_back(id);
}

That creates a new broadcast on the loop domain. But if we just do setLoopDomain(something with manually created new broadcasts), we will fail.

And in this PR, we are not using an interface like TensorDomain::broadcast, so it will not have the problem of the above mechanism. But if we just do a setLoopDomain(getLoopDomain()), we will fail.

In summary, I believe neither the additional_ids_ mechanism nor the one in this PR is safe. The additional_ids_ is more convenient for the TensorDomain::broadcast-like interface, that is, we have a method function that says "add this new ID to this axis". And the one in this PR is more convenient for setLoopDomain way of inserting new IDs.

I think my question for now is: do we want to pause a bit to think about interface design, or do we want to just quickly unblock for now and do some cleanup later?

[naoyam]

Per our offline discussion, we decided that it's more important to move forward and get some progress before too much worrying about interface designs. See also #3000