Function ShuffleInTensor3SimpleVector

tensorflow/core/kernels/conv_2d_gpu.h:214–260

template <typename T, int sp0, int sp1, int sp2, bool conjugate = false>
__global__ void ShuffleInTensor3SimpleVector(int nthreads,
                                             const T* __restrict__ input,
                                             Dimension<3> input_dims,
                                             T* __restrict__ output) {
  Dimension<3> output_dims;
  output_dims[sp0] = input_dims[0];
  output_dims[sp1] = input_dims[1];
  output_dims[sp2] = input_dims[2];

  const int stride = blockDim.x * gridDim.x * kUnroll;
  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
  T buf[kUnroll];

  int output_index;
  for (output_index = tid * kUnroll; output_index + kUnroll - 1 < nthreads;
       output_index += stride) {
#pragma unroll
    for (int i = 0; i < kUnroll; i++) {
      int output_index_i = output_index + i;
      Index<3> output_tensor_index = FlatToTensorIndex(output_index_i,
                                                       output_dims);
      Index<3> input_tensor_index;
      input_tensor_index[0] = output_tensor_index[sp0];
      input_tensor_index[1] = output_tensor_index[sp1];
      input_tensor_index[2] = output_tensor_index[sp2];

      int input_index_i = TensorIndexToFlat(input_tensor_index, input_dims);
      buf[i] = maybe_conj<T, conjugate>::run(ldg(input + input_index_i));
    }
    float2 *out = reinterpret_cast<float2*>(output + output_index);
    *out = *reinterpret_cast<float2*>(buf);
  }

  for(; output_index < nthreads; output_index++) {
    Index<3> output_tensor_index = FlatToTensorIndex(output_index, output_dims);

    Index<3> input_tensor_index;
    input_tensor_index[0] = output_tensor_index[sp0];
    input_tensor_index[1] = output_tensor_index[sp1];
    input_tensor_index[2] = output_tensor_index[sp2];

    int input_index = TensorIndexToFlat(input_tensor_index, input_dims);

    output[output_index] =
        maybe_conj<T, conjugate>::run(ldg(input + input_index));
  }
}

// Use shared memory tiles to swap dimension-1 and dimension-2 of a 3D tensor,
// where dimensions are zero-based: output[i][j][k] = input[i][k][j].