Invalid code generation with computeAt

[naoyam]

🐛 Bug

This test generates invalid code:
(https://github.com/naoyam/pytorch/blob/compute-at-bug/test/cpp/jit/test_gpu.cpp)

void testGPU_FusionComputeAtBug() {
  torch::jit::fuser::cuda::CudaKernel prog;
  prog.setFusionPtr(std::make_unique<Fusion>());
  Fusion* fusion = prog.fusion();
  FusionGuard fg(fusion);

  // Set up your input tensor views
  TensorView* tv0 = makeDummyTensor(2);
  fusion->addInput(tv0);

  TensorView* tv1 = mul(tv0, new Float(1));
  TensorView* tv2 = add(tv0, new Float(2));
  TensorView* tv3 = add(tv1, new Float(3));
  TensorView* tv4 = add(tv1, new Float(4));

  fusion->addOutput(tv2);
  fusion->addOutput(tv3);
  fusion->addOutput(tv4);

  std::cout << "tv1->computeAt(tv3, 1)\n";
  tv1->computeAt(tv3, -1);
  fusion->printMath();
  fusion->printKernel();
  fusion->printMath();

  prog.setDevice(0);
  prog.grid(1);
  prog.block(1);
  torch::jit::fuser::cuda::compileKernel(&prog);
}

Here's the generated kernel:

__global__ void kernel(Tensor<float, 2> T0, Tensor<float, 2> T2, Tensor<float, 2> T3, Tensor<float, 2> T4){
  for(size_t i36 = 0; i36 < T4.size[0]; ++i36 ) {
    for(size_t i37 = 0; i37 < T4.size[1]; ++i37 ) {
      float T1[1];
      T1[ 0 ]
         = T0[ ( i36 * T0.stride[0] ) + ( i37 * T0.stride[1] ) ]
         * float(1);
      T4[ ( i36 * T4.stride[0] ) + ( i37 * T4.stride[1] ) ]
         = T1[ 0 ]
         + float(4);
    }
  }
  for(size_t i39 = 0; i39 < T4.size[0]; ++i39 ) {
    for(size_t i40 = 0; i40 < T4.size[1]; ++i40 ) {
      T2[ ( i39 * T2.stride[0] ) + ( i40 * T2.stride[1] ) ]
         = T0[ ( i39 * T0.stride[0] ) + ( i40 * T0.stride[1] ) ]
         + float(2);
    }
  }
  for(size_t i41 = 0; i41 < T4.size[0]; ++i41 ) {
    for(size_t i42 = 0; i42 < T4.size[1]; ++i42 ) {
      T3[ ( i41 * T3.stride[0] ) + ( i42 * T3.stride[1] ) ]
         = T1[ 0 ]
         + float(3);
    }
  }
}

Notice that the loop for T3 references T1, but that's only defined in the first loop nest. The T3 loop should be actually placed in the same loop nest as T4, but it's "blocked" by the T2 loop. We had a similar issue recently, which we fixed by sorting inputs of expressions when traversing fusion expressions (see #112). This issue is similar, but happens with output expressions.

[naoyam]

This also fails:

void testGPU_FusionComputeAtBug2() {
  torch::jit::fuser::cuda::CudaKernel prog;
  prog.setFusionPtr(std::make_unique<Fusion>());
  Fusion* fusion = prog.fusion();
  FusionGuard fg(fusion);

  // Set up your input tensor views
  TensorView* tv0 = makeDummyTensor(2);
  fusion->addInput(tv0);

  TensorView* tv1 = mul(tv0, new Float(1));
  TensorView* tv2 = mul(tv1, new Float(2));

  TensorView* tv3 = add(tv0, new Float(3));
  TensorView* tv4 = add(tv3, new Float(4));

  TensorView* tv5 = add(tv1, tv3);

  fusion->addOutput(tv2);
  fusion->addOutput(tv4);
  fusion->addOutput(tv5);

  fusion->printMath();
  fusion->printKernel();

  std::cout << "apply computeAt\n";
  tv1->computeAt(tv5, -1);
  tv3->computeAt(tv5, -1);

  fusion->printMath();
  fusion->printKernel();
  fusion->printMath();

  prog.setDevice(0);
  prog.grid(1);
  prog.block(1);
  torch::jit::fuser::cuda::compileKernel(&prog);
}

__global__ void kernel(Tensor<float, 2> T0, Tensor<float, 2> T2, Tensor<float, 2> T4, Tensor<float, 2> T5){
  for(size_t i78 = 0; i78 < T5.size[0]; ++i78 ) {
    for(size_t i79 = 0; i79 < T5.size[1]; ++i79 ) {
      float T1[1];
      T1[ 0 ]
         = T0[ ( i78 * T0.stride[0] ) + ( i79 * T0.stride[1] ) ]
         * float(1);
      float T3[1];
      T3[ 0 ]
         = T0[ ( i78 * T0.stride[0] ) + ( i79 * T0.stride[1] ) ]
         + float(3);
      T5[ ( i78 * T5.stride[0] ) + ( i79 * T5.stride[1] ) ]
         = T1[ 0 ]
         + T3[ 0 ];
      T2[ ( i78 * T2.stride[0] ) + ( i79 * T2.stride[1] ) ]
         = T1[ 0 ]
         * float(2);
    }
  }
  for(size_t i82 = 0; i82 < T5.size[0]; ++i82 ) {
    for(size_t i83 = 0; i83 < T5.size[1]; ++i83 ) {
      T4[ ( i82 * T4.stride[0] ) + ( i83 * T4.stride[1] ) ]
         = T3[ 0 ]
         + float(4);
    }
  }
}

[naoyam]

This also fails:
(See testGPU_FusionComputeAtOutputOrder4)

// First tree
  TensorView* tv0 = makeDummyTensor(1);
  fusion->addInput(tv0);
  TensorView* tv1 = add(tv0, new Float(1));
  TensorView* tv2 = add(tv1, new Float(2));
  TensorView* tv3 = add(tv1, new Float(3));
  fusion->addOutput(tv2);
  fusion->addOutput(tv3);

  // Second tree
  TensorView* tv4 = makeDummyTensor(1);
  fusion->addInput(tv4);
  TensorView* tv5 = add(tv4, new Float(5));
  TensorView* tv6 = add(tv5, new Float(6));
  TensorView* tv7 = add(tv5, new Float(7));
  fusion->addOutput(tv6);
  fusion->addOutput(tv7);

  tv1->computeAt(tv2, -1);
  tv5->computeAt(tv6, -1);

All the four output TVs must be properly ordered, but there is no such ordering enforcement right now.

[naoyam]

This also fails:

 TensorView* tv0 = makeDummyTensor(2);
  fusion->addInput(tv0);
  TensorView* tv1 = add(tv0, new Float(1));
  TensorView* tv2 = add(tv1, new Float(2));
  TensorView* tv3 = add(tv0, new Float(3));
  TensorView* tv4 = add(tv3, new Float(4));
  TensorView* tv5 = add(tv2, tv4);

  fusion->addOutput(tv1);
  fusion->addOutput(tv3);
  fusion->addOutput(tv5);

  tv2->computeAt(tv5, -1);
  tv4->computeAt(tv5, -1);

__global__ void kernel(Tensor<float, 2> T0, Tensor<float, 2> T1, Tensor<float, 2> T3, Tensor<float, 2> T5){
  for(size_t i53 = 0; i53 < T5.size[0]; ++i53 ) {
    for(size_t i54 = 0; i54 < T5.size[1]; ++i54 ) {
      T1[ ( i53 * T1.stride[0] ) + ( i54 * T1.stride[1] ) ]
         = T0[ ( i53 * T0.stride[0] ) + ( i54 * T0.stride[1] ) ]
         + float(1);
    }
  }
  for(size_t i55 = 0; i55 < T5.size[0]; ++i55 ) {
    for(size_t i56 = 0; i56 < T5.size[1]; ++i56 ) {
      float T2[1];
      T2[ 0 ]
         = T1[ ( i55 * T1.stride[0] ) + ( i56 * T1.stride[1] ) ]
         + float(2);
    }
  }
  for(size_t i58 = 0; i58 < T5.size[0]; ++i58 ) {
    for(size_t i59 = 0; i59 < T5.size[1]; ++i59 ) {
      T3[ ( i58 * T3.stride[0] ) + ( i59 * T3.stride[1] ) ]
         = T0[ ( i58 * T0.stride[0] ) + ( i59 * T0.stride[1] ) ]
         + float(3);
    }
  }
  for(size_t i60 = 0; i60 < T5.size[0]; ++i60 ) {
    for(size_t i61 = 0; i61 < T5.size[1]; ++i61 ) {
      float T4[1];
      T4[ 0 ]
         = T3[ ( i60 * T3.stride[0] ) + ( i61 * T3.stride[1] ) ]
         + float(4);
      T5[ ( i60 * T5.stride[0] ) + ( i61 * T5.stride[1] ) ]
         = T2[ 0 ]
         + T4[ 0 ];
    }
  }
}

[naoyam]

Closed by #174