Invalid code generation

[naoyam]

🐛 Bug

The testGPU_FusionSoftmax test in this branch generates invalid code:

https://github.com/naoyam/pytorch/tree/codegen-bug

Here's the generated kernel:

__global__ void CUDAGeneratedKernel(Tensor<float, 3> T0, Tensor<float, 3> T9){
  __shared__ float shared_mem[1024];
  float T6[1];
  float T5[1];
  T5[ 0 ]
     = float(0);
  float T11[1];
  if ( ( ( ( 0 * 32 ) + threadIdx.x ) < T9.size[2] ) ) {
    T11[ 0 ]
       = float(0);
  }
  float T2[1];
  float T1[1];
  T1[ 0 ]
     = float(0);
  float T10[1];
  if ( ( ( ( 0 * 32 ) + threadIdx.x ) < T9.size[2] ) ) {
    T10[ 0 ]
       = float(0);
  }
  for(size_t i84 = 0; i84 < ( ceilDiv(T9.size[2], 32) ); ++i84 ) {
    if ( ( ( ( i84 * 32 ) + threadIdx.x ) < T9.size[2] ) ) {
      T10[ 0 ]
         = fmaxf(T10[ 0 ]
        , T0[ ( blockIdx.x * T0.stride[0] ) + ( blockIdx.y * T0.stride[1] ) + ( ( ( i84 * 32 ) + threadIdx.x ) * T0.stride[2] ) ]);
    }
  }
  blockReduce< true, false, false > ( T1[ 0 ], T10[ 0 ], reduction_fmaxf_float, threadIdx, blockDim, reinterpret_cast<float*>(shared_mem));
  if ( ( threadIdx.x == 0 ) ) {
    T2[ 0 ]
       = T1[ 0 ];
  }
  for(size_t i87 = 0; i87 < ( ceilDiv(T9.size[2], 32) ); ++i87 ) {
    float T7[1];
    if ( ( ( ( i87 * 32 ) + threadIdx.x ) < T9.size[2] ) ) {
      T7[ 0 ]
         = T0[ ( blockIdx.x * T0.stride[0] ) + ( blockIdx.y * T0.stride[1] ) + ( ( ( i87 * 32 ) + threadIdx.x ) * T0.stride[2] ) ]
         - T2[ 0 ];
    }
    float T8[1];
    if ( ( ( ( i87 * 32 ) + threadIdx.x ) < T9.size[2] ) ) {
      T8[ 0 ]
         = expf(T7[ 0 ]);
    }
    float T3[1];
    if ( ( ( ( i87 * 32 ) + threadIdx.x ) < T9.size[2] ) ) {
      T3[ 0 ]
         = T0[ ( blockIdx.x * T0.stride[0] ) + ( blockIdx.y * T0.stride[1] ) + ( ( ( i87 * 32 ) + threadIdx.x ) * T0.stride[2] ) ]
         - T2[ 0 ];
    }
    float T4[1];
    if ( ( ( ( i87 * 32 ) + threadIdx.x ) < T9.size[2] ) ) {
      T4[ 0 ]
         = expf(T3[ 0 ]);
    }
  }
  for(size_t i93 = 0; i93 < ( ceilDiv(T9.size[2], 32) ); ++i93 ) {
    if ( ( ( ( i93 * 32 ) + threadIdx.x ) < T9.size[2] ) ) {
      T11[ 0 ]
         = T11[ 0 ]
         + T4[ 0 ];
    }
  }
  blockReduce< true, false, false > ( T5[ 0 ], T11[ 0 ], reduction_add_float, threadIdx, blockDim, reinterpret_cast<float*>(shared_mem));
  if ( ( threadIdx.x == 0 ) ) {
    T6[ 0 ]
       = T5[ 0 ];
  }
  for(size_t i96 = 0; i96 < ( ceilDiv(T9.size[2], 32) ); ++i96 ) {
    if ( ( ( ( i96 * 32 ) + threadIdx.x ) < T9.size[2] ) ) {
      T9[ ( blockIdx.x * T9.stride[0] ) + ( blockIdx.y * T9.stride[1] ) + ( ( ( i96 * 32 ) + threadIdx.x ) * T9.stride[2] ) ]
         = T8[ 0 ]
         / T6[ 0 ];
    }
  }
}

T8 is used in the final loop nest but is defined and computed in a different loop nest.

The final loop nest corresponds to this expression:

TensorView* output_tv6 = div(exp_tv4_2, bcast_sum_tv6);

Interestingly, swapping the two operands, i.e., div(bcast_sum_tv6, exp_tv4_2), seems to be fine. See https://github.com/naoyam/pytorch/blob/d2913c9fb2c7e94c786ce3b73a6b5ab87b4eee3b/test/cpp/jit/test_gpu.cpp#L2798. The fuser generates code as expected, although they are not the same computation anymore.

Note that this issue was encountered in the block broadcast PR (#100) but shows up with the dev branch as well.

[naoyam]

Here's a very simple reproducer: naoyam@979d421#diff-538e381bc787011cc53c15caf54e126dR3344-R3369

The generated code looks like this:

__global__ void kernel(Tensor<float, 1> T0, Tensor<float, 1> T3, Tensor<float, 1> T2){
  for(size_t i26 = 0; i26 < T2.size[0]; ++i26 ) {
    float T1[1];
    T1[ 0 ]
       = T0[ ( i26 * T0.stride[0] ) ]
       + float(1);
  }
  for(size_t i28 = 0; i28 < T2.size[0]; ++i28 ) {
    T2[ ( i28 * T2.stride[0] ) ]
       = T0[ ( i28 * T0.stride[0] ) ]
       + float(1);
  }
  for(size_t i29 = 0; i29 < T2.size[0]; ++i29 ) {
    T3[ ( i29 * T3.stride[0] ) ]
       = T1[ 0 ]
       + T2[ ( i29 * T2.stride[0] ) ];
  }
}

Here, T1 should be computed at T3 in the final loop, but it is effectively blocked by T2, with which it is grouped into the same loop nest.

[csarofeen]

Copying a conversation between Naoya and myself:

Naoya: I believe this is because in the first case, the topological order is T1 -> T2 -> T3 , while in the second case, it's T2 -> T1 -> T3. Both are a valid topological order, but in the former case, T1 is not going to be in the same loop as T3 since there is a middle node, T2, which is not computed at T3.

Christian: You are 100% correct. The way the Exprs are generated don't take into consideration the computeAt structure and they need to. Our graph traversals are output to input. So we likely need to know which TVs are computeAt-ing into a given output of an Expr. We then need to sort those TVs giving priority to those with a higher computeAt position. i.e. if we have the structure:

tv4->computeAt(tv6, 2)
tv5->computeAt(tv6, 1)

with Exprs:

tv4 = ...
tv5 = ...
tv6 = tv4 + tv5

then we need the ordering:

tv5 = ...
tv4 = ...
tv6 = tv4 + tv5

Because when we generate loops, the structuring would need to be:

for(tv6->axis(0))
  for (tv5->axis(1:...) )
    tv5 = ...
  for(tv6->axis(1))
    for(tv4->axis(2:...) )
      tv4 = ...
    for(tv6->axis(2:...)
      tv6 = tv4 + tv5

We should insert a special ordering pass in the lowering before the other passes to get the loop structure in the correct order. I believe we could also avoid the reordering that's done right now of the compute at structure.