Allow inlining past loop broadcasts

[jacobhinkle]

Stacked on #3414

This PR enables us to inline an MmaOp properly when its inputs are missing broadcast dimensions. We do this by always allowing inlining past loop broadcasts or their transforms. For example

tv0:
  logical [ iS1{i0} ]
  loop [ iS1{i0} bS5{1} ]
tv1:
  logical [ iS2{i1} ]
  loop [ bS6{1} iS2{i1} ]
tv2 = foo(tv0, tv1)
  logical [ iS3{i0} iS4{i1} ]

As long as the operation foo properly maps its arguments despite the missing logical dimensions (as MmaOp does as of #3391), then we should be able to fully inline this case because the loop broadcasts bS5 and bS6 are imaginary in the sense that they don't impact indexing.

[jacobhinkle]

After this, we can actually generate a proper kernel and run it. I will rebase #3406 onto this and modify the test to compile and run in that PR so we can inspect the generated kernel there. We can keep this PR for discussing the inlining changes only.

[naoyam]

Does this only apply to broadcast IDs added by TensorView::broadcast()?

[jacobhinkle]

Does this only apply to broadcast IDs added by TensorView::broadcast()?

Yes, that's the intention. I am using tv->domain()->additionalIDs(), which I think is only those broadcasts?

[naoyam]

Does this only apply to broadcast IDs added by TensorView::broadcast()?

Yes, that's the intention. I am using tv->domain()->additionalIDs(), which I think is only those broadcasts?

Yes. @zasdfgbnm, when you added this, were you thinking about having non-broadcast IDs in additional_ids_?

[jacobhinkle]

Does this only apply to broadcast IDs added by TensorView::broadcast()?

Yes, that's the intention. I am using tv->domain()->additionalIDs(), which I think is only those broadcasts?

Yes. @zasdfgbnm, when you added this, were you thinking about having non-broadcast IDs in additional_ids_?

To be safe I'll check the IterType when skipping.

[jacobhinkle]

This includes all IDs that are between additionalIDs() and loop domain. However, we could have something like this:

tv->broadcast(0, 16);
tv->merge(0);

In this case, we'll be merging the new broadcast ID with a pre-existing loop ID, so we should not ignore that. I think instead maybe what we should do is traverse from the root domain to the loop domain instead and the complement will then be the "pure" loop broadcasts which we can ignore.

[jacobhinkle]

Actually I suppose that would also automatically allow us to inline past regular broadcasts that are created using BroadcastOp since those new Broadcast IDs are not reachable from the root domain either, but we already inline past those IDs anyway I believe.