Inlining error in Hopper matmul with AxisMapping and grid swizzling

[jacobhinkle]

The inlining logic for MmaOp with AxisMapping checks that unmapped dimensions are Broadcast. We expect to have something like this

t0 {
  logical: [ iS0{M} iS1{K} ]
  loop: [ iS2{ ceilDiv(M, 256) }  bS5{1}  iS3{256} bS6{256} ... ]
    Split iS0 by 256 -> iS2, iS3
    Split bS4 by 256 -> bS5, bS6
      ...
  additional ids: bS4{1}
}
t1 {
  logical: [ iS10{N} iS11{K} ]
  loop: [ bS13{1} iS15{ ceilDiv(N, 256) } bS14{256} iS16{256} ... ]
    Split bS12 by 256 -> bS13, bS14
    Split iS10 by 256 -> iS15, iS16
      ...
  additional ids: bS12{1}
}

In this case, we are able to inline the mma operation that consumes these two tensors, but we check that the unmapped IDs 5, 6, 13, and 14 are Broadcast and that the operation is an MmaOp.

In the case of grid swizzling by a factor 4, we will do some further scheduling here. For example we will have

t0 {
  logical: [ iS0{M} iS1{K} ]
  loop: [ iS21{ ceilDiv( ceilDiv(M, 256), 4) } iS22{4} iS3{256} bS6{256} ... ]
    Split iS0 by 256 -> iS2, iS3
    Split bS4 by 256 -> bS5, bS6
    Split iS2 by 4 -> iS20, iS21{4}
    Merge bS5 with iS21{4} -> iS22{4}
      ...
  additional ids: bS4{1}
}

Now we have mixed the first two outer dimensions with this swizzle and what used to be a simple split of a loop broadcast (bS5) is now an iteration ID iS22{4} resulting from the merge.

I am not sure yet how to address this. I don't think we can just inline here without some other changes since when I disable this check I get errors in expression sorting.