Merging bcast and non-bcast dim + computeAt

[csarofeen]

The following case throws a computeAt error

    Fusion fusion;
    FusionGuard fg(&fusion);

    // Set up your input tensor views
    TensorView* tv0 = makeDummyTensor(2);
    fusion.addInput(tv0);
    TensorView* tv1 = add(tv0, new Float(2.0));
    TensorView* tv2 = broadcast(tv1, {true, false, false});
    TensorView* tv3 = makeDummyTensor(3);
    fusion.addInput(tv3);
    TensorView* tv4 = mul(tv2, tv3);
    fusion.addOutput(tv4);

    tv4->merge(0);
    tv4->split(-1, 8);
    tv4->split(-1, 4);

    tv2->computeAt(tv4, 2);

IR:

T1[ iS{i1}, iS{i3} ]
   = T0[ iS{i1}, iS{i3} ]
   + float(2);
T2[ bS{1}, iS{i1}, iS{i3} ] = broadcast( T1[ iS{i1}, iS{i3} ] )
T4[ iS{( i8 * i1 )}, iS{( ceilDiv(i3, 8) )}, iS{( ceilDiv(8, 4) )}, iS{4} ]
   = T2[ bS{1}, iS{i1}, iS{i3} ]
   * T3[ iS{i8}, iS{i10}, iS{i12} ];

error:

C++ exception with description "BestEffortReplay::findFirstMismatchedID( consumer_->domain(), tv_data.at(consumer_).getOriginalDomain()) == consumer_->domain()->nDims() INTERNAL ASSERT FAILED at "../torch/csrc/jit/codegen/cuda/compute_at.cpp":382, please report a bug to PyTorch. ComputeAt logic changed the consumer domain which should not happen. Domain was [ iS{( i8 * i1 )}, iS{( ceilDiv(i3, 8) )}, iS{( ceilDiv(8, 4) )}, iS{4} ] but is now: [ iS{( i8 * i1 )}, iS{( ceilDiv(i3, 8) )}, iS{8}, iS{( ceilDiv(8, 4) )}, iS{4} ]

Need to figure out why this error is being produced. This is something we should support.

[csarofeen]

Thought this was from bcast merge but is also happening with:

    Fusion fusion;
    FusionGuard fg(&fusion);

    // Set up your input tensor views
    TensorView* tv0 = makeDummyTensor(2);
    fusion.addInput(tv0);
    TensorView* tv1 = makeDummyTensor(2);
    fusion.addInput(tv1);
    TensorView* tv2 = add(tv0, new Float(2.0));
    TensorView* tv3 = mul(tv1, tv2);
    fusion.addOutput(tv3);

    tv3->merge(0);
    tv3->split(-1, 8);
    tv3->split(-1, 4);
  
    tv2->computeAt(tv3, 1);