csarofeen / pytorch

Tensors and Dynamic neural networks in Python with strong GPU acceleration
http://pytorch.org
Other
26 stars 7 forks source link

[Question] Replaying post-view tensors cancels view transformations #1555

Open naoyam opened 2 years ago

naoyam commented 2 years ago

There's some inconsistency with this fusion:

  Fusion fusion;
  FusionGuard fg(&fusion);

  const std::vector<int64_t> shape1{10, 12};
  const std::vector<int64_t> shape2{10, 3, 4};

  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);

  auto tv1 = add(tv0, IrBuilder::create<Double>(1));

  auto tv2 = view(tv1, shape1, shape2);
  fusion.addOutput(tv2);

  auto tv3 = add(tv1, IrBuilder::create<Double>(1));
  fusion.addOutput(tv3);

  tv3->split(-1, 5);
  tv1->computeAt(tv3, -1);

tv1 is an input to view and add. Before the computeAt, the fusion looks like:

%kernel_math {
T1_l[ iS2{i1}, iS3{i2} ]
   = T0_g[ iS0{i1}, iS1{i2} ]
   + double(1);
T2_g[ iS4{i1}, iS6{3}rf, iS7{( ceilDiv(i2, 3) )}rf ] = view( T1_l[ iS2{i1}, iS3{i2} ] )
T3_g[ iS8{i1}, iS10{( ceilDiv(i2, 5) )}, iS11{5} ]
   = T1_l[ iS2{i1}, iS3{i2} ]
   + double(1);
}

With tv1->computeAt(tv3, -1), tv2 is transformed like tv3, effectively cancelling the the view transformation.

 T2_g[ iS4{i1}, iS14{( ceilDiv(i2, 5) )}, iS15{5} ] produce_pos( 3)

I'm not really sure if this makes sense as the effect of view disappears. Maybe fine as long as the shapes of input and output tensors are not incorrectly changed.

However, T2 is in an inconsistent state as it keeps its "old" rfactor domain.

RFactor domain: iS4{i1} iS6{3}rf iS7{( ceilDiv(i2, 3) )}rf
Root domain: iS4{i1} iS5{i2}

Note that the fusion fails when lowered due to the inconsistency.

So, my questions are:

  1. Currently, replaying post-view tensors can result in cancelling view transformations. Does it make sense? Should it be disallowed?
  2. Assuming it's allowed, we would need to make sure rfactor domains get removed when post-view tensors are replayed
rdspring1 commented 2 years ago

Just some thoughts:

  1. We need to propagate the view transformations to the dependent tensors before scheduling. Alternative: We can schedule this fusion if we select the view tensor as the reference tensor in the pointwise scheduler.
  // This schedule passes
  TransformPropagator::from(tv2);
  tv3->split(-1, 5);
  tv1->computeAt(tv3, -1);
  1. view also changes the dimensions of the tensor, so it has an effect even if its domain transformations are ignored.
naoyam commented 2 years ago
  1. We need to propagate the view transformations to the dependent tensors before scheduling. Alternative: We can schedule this fusion if we select the view tensor as the reference tensor in the pointwise scheduler.
  // This schedule passes
  TransformPropagator::from(tv2);
  tv3->split(-1, 5);
  tv1->computeAt(tv3, -1);

Yes, that would work, but my question is what should be done if the view tensor is not selected as the reference. Or, what about if there are multiple view tensors?