[Indexing failure] Need to find permissive index binding on the inner side of CA axis.

[shmsong]

🐛 Describe the bug

We have an unsupported case currently with inlined broadcasting.

Minimum repro:

TEST_F(NVFuserTest, FusionInlineBroadcastIndexing_CUDA) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeContigTensor(1);
  auto tv1 = makeContigTensor(2);
  fusion.addInput(tv0);
  fusion.addInput(tv1);
  auto tv2 = set(tv0);
  auto tv3 = broadcast(tv2, {true, false});
  auto tv4 = add(tv3, tv1);
  fusion.addOutput(tv4);

  tv4->merge(0);
  tv4->split(0, 128);

  tv0->computeAt(tv4, 1);

  // [I, 128]
  // This line triggers the failrue.
  tv2->split(-1, 32);

  FusionExecutor fe;
  fe.compileFusion(&fusion);
}

Error message:

C++ exception with description "index_map.find(root_dom[i]) != index_map.end() INTERNAL ASSERT FAILED at "torch/csrc/jit/codegen/cuda/index_compute.cpp":2024, please report a bug to PyTorch. Couldn't find root mapping for T2_l[ iS14{( ceilDiv(T0.size[0], 128) )}, iS18{( ceilDiv(128, 32) )}, iS19{32} ] ca_pos( 1 ) dim: 0 id: iS23{T0.size[0]}
E

Fusion debug print:

%kernel {
T2_l[ iS14{( ceilDiv(i1, 128) )}, iS18{( ceilDiv(128, 32) )}, iS19{32} ] ca_pos( 1 )
   = T0_g[ iS16{( ceilDiv(i1, 128) )}, iS17{128} ];
T3_l[ iS12{( ceilDiv(( 1 * i1 ), 128) )}, iS13{128} ] ca_pos( 1 ) produce_pos( 1)
   = broadcast( T2_l[ iS14{( ceilDiv(i1, 128) )}, iS18{( ceilDiv(128, 32) )}, iS19{32} ] ca_pos( 1 ) )
T4_g[ iS9{( ceilDiv(( i2 * i1 ), 128) )}, iS10{128} ] produce_pos( 1)
   = T3_l[ iS12{( ceilDiv(( 1 * i1 ), 128) )}, iS13{128} ] ca_pos( 1 ) produce_pos( 1)
   + T1_g[ iS1{i2}, iS2{i3} ];

TransformPrinter : 
T0_g[ iS16{( ceilDiv(i1, 128) )}, iS17{128} ]
 root domain : (iS0{i1})
  Split: iS0{i1} by factor 128 -> iS16{( ceilDiv(i1, 128) )}, iS17{128}, start offset: 0, stop offset: 0
T2_l[ iS14{( ceilDiv(i1, 128) )}, iS18{( ceilDiv(128, 32) )}, iS19{32} ] ca_pos( 1 )
 root domain : (iS3{i1})
  Split: iS3{i1} by factor 128 -> iS14{( ceilDiv(i1, 128) )}, iS15{128}, start offset: 0, stop offset: 0
  Split: iS15{128} by factor 32 -> iS18{( ceilDiv(128, 32) )}, iS19{32}, start offset: 0, stop offset: 0
T3_l[ iS12{( ceilDiv(( 1 * i1 ), 128) )}, iS13{128} ] ca_pos( 1 ) produce_pos( 1)
 root domain : (bS4{1},iS5{i1})
  Merge: bS4{1} and iS5{i1} -> iS11{( 1 * i1 )}
  Split: iS11{( 1 * i1 )} by factor 128 -> iS12{( ceilDiv(( 1 * i1 ), 128) )}, iS13{128}, start offset: 0, stop offset: 0
T1_g[ iS1{i2}, iS2{i3} ]
 root domain : (iS1{i2},iS2{i3})
T4_g[ iS9{( ceilDiv(( i2 * i1 ), 128) )}, iS10{128} ] produce_pos( 1)
 root domain : (iS6{i2},iS7{i1})
  Merge: iS6{i2} and iS7{i1} -> iS8{( i2 * i1 )}
  Split: iS8{( i2 * i1 )} by factor 128 -> iS9{( ceilDiv(( i2 * i1 ), 128) )}, iS10{128}, start offset: 0, stop offset: 0
}

This leads to indexing failure with T0 and T2 since the left of their ca pos are loop mapped to broadcasted Iterdomains so we'd need to use permissive mapping instead of exact mapping to bind the index for, say, iS15{128}, and currently we are only doing this for leaf domains / loop ids:

https://github.com/csarofeen/pytorch/blob/devel/torch/csrc/jit/codegen/cuda/lower_index_compute.cpp#L714-L715

Thoroughly fixing this one would take a bit of formalization, potentially with additional pass on concrete indices or a different traversal order. Would appreciate any suggestions.

Meanwhile I guess we could detect this usage and assert on them.

Versions

N/A