  void Compute(OpKernelContext* context) override {
    const Tensor& input = context->input(0);
    const TensorShape& input_shape = input.shape();

    const Tensor& out_backprop = context->input(2);
    const TensorShape& out_backprop_shape = out_backprop.shape();

    TensorShape filter_shape;
    if (takes_shape_) {
      const Tensor& filter_sizes = context->input(1);
      OP_REQUIRES(context, TensorShapeUtils::IsVector(filter_sizes.shape()),
                  errors::InvalidArgument(
                      "filter_sizes shape must be rank 1 but is rank ",
                      filter_sizes.shape().dims()));
      OP_REQUIRES_OK(context, TensorShapeUtils::MakeShape(
                                  filter_sizes.vec<int32>(), &filter_shape));
    } else {
      filter_shape = context->input(1).shape();
    }

    OP_REQUIRES(context, input_shape.dims() == 5,
                errors::InvalidArgument("input tensor must have 5 dimensions"));
    OP_REQUIRES(
        context, filter_shape.dims() == 5,
        errors::InvalidArgument("filter_sizes tensor must have 5 dimensions"));
    OP_REQUIRES(
        context, out_backprop_shape.dims() == 5,
        errors::InvalidArgument("out_backprop tensor must have 5 dimensions"));
    OP_REQUIRES(
        context, input_shape.dim_size(4) == filter_shape.dim_size(3),
        errors::InvalidArgument("input and filter_sizes must have the same "
                                "number of channels. Got ",
                                input_shape.dim_size(4), " for input and ",
                                filter_shape.dim_size(3), " for filter_sizes"));
    OP_REQUIRES(
        context, out_backprop_shape.dim_size(4) == filter_shape.dim_size(4),
        errors::InvalidArgument("out_backprop and filter_sizes must have the "
                                "same number of channels. Got ",
                                out_backprop_shape.dim_size(4),
                                " for out_backprop and ",
                                filter_shape.dim_size(4), " for filter_sizes"));

    ConvBackpropDimensions dims;
    OP_REQUIRES_OK(context,
                   ConvBackpropComputeDimensions(
                       "Conv3DBackpropFilterOp", /*num_spatial_dims=*/3,
                       input_shape, filter_shape, out_backprop_shape, stride_,
                       padding_, data_format_, &dims));

    Tensor* filter_backprop;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, filter_shape, &filter_backprop));

    if (input_shape.num_elements() == 0) {
      filter_backprop->template flat<T>().setZero();
      return;
    }

    int64_t top_pad_planes, bottom_pad_planes;
    int64_t top_pad_rows, bottom_pad_rows;
    int64_t left_pad_cols, right_pad_cols;

    OP_REQUIRES_OK(context, GetWindowedOutputSizeVerbose(
                                dims.spatial_dims[0].input_size,
                                dims.spatial_dims[0].filter_size,
                                dims.spatial_dims[0].stride, padding_,
                                &dims.spatial_dims[0].output_size,
                                &top_pad_planes, &bottom_pad_planes));
    OP_REQUIRES_OK(context, GetWindowedOutputSizeVerbose(
                                dims.spatial_dims[1].input_size,
                                dims.spatial_dims[1].filter_size,
                                dims.spatial_dims[1].stride, padding_,
                                &dims.spatial_dims[1].output_size,
                                &top_pad_rows, &bottom_pad_rows));
    OP_REQUIRES_OK(context, GetWindowedOutputSizeVerbose(
                                dims.spatial_dims[2].input_size,
                                dims.spatial_dims[2].filter_size,
                                dims.spatial_dims[2].stride, padding_,
                                &dims.spatial_dims[2].output_size,
                                &left_pad_cols, &right_pad_cols));

    // TODO(ezhulenev): Extract work size and shard estimation to shared
    // functions in conv_grad_ops, and update 2d convolution backprop.

    // The total dimension size of each kernel.
    const int64_t filter_total_size =
        dims.spatial_dims[0].filter_size * dims.spatial_dims[1].filter_size *
        dims.spatial_dims[2].filter_size * dims.in_depth;
    // The output image size is the spatial size of the output.
    const int64_t output_image_size = dims.spatial_dims[0].output_size *
                                      dims.spatial_dims[1].output_size *
                                      dims.spatial_dims[2].output_size;

    // Shard 'batch' images (volumes) into 'shard_size' groups of images
    // (volumes) to be fed into the parallel matmul. Calculate 'shard_size' by
    // dividing the L3 cache size ('target_working_set_size') by the matmul size
    // of an individual image ('work_unit_size').

    const auto cache_sizes = Eigen::internal::CacheSizes();
    const ptrdiff_t l3_cache_size = cache_sizes.m_l3;

    // TODO(andydavis)
    // *) Consider reducing 'target_working_set_size' if L3 is shared by
    //    other concurrently running tensorflow ops.
    const size_t target_working_set_size = l3_cache_size / sizeof(T);

    const int64_t size_A = output_image_size * filter_total_size;

    const int64_t size_B = output_image_size * dims.out_depth;

    const int64_t size_C = filter_total_size * dims.out_depth;

    const int64_t work_unit_size = size_A + size_B + size_C;

    OP_REQUIRES(
        context, work_unit_size > 0,
        errors::InvalidArgument("input, filter_sizes and out_backprop tensors "
                                "must all have at least 1 element"));

    const size_t shard_size =
        (target_working_set_size + work_unit_size - 1) / work_unit_size;

    // Total number of elements in all the tensors used by this kernel.
    int64_t total_tensor_elements = input_shape.num_elements() +
                                    filter_shape.num_elements() +
                                    out_backprop_shape.num_elements();

    // Shape of the temporary workspace buffer.
    TensorShape col_buffer_shape = {static_cast<int64_t>(shard_size),
                                    static_cast<int64_t>(output_image_size),
                                    static_cast<int64_t>(filter_total_size)};
    int64_t col_buffer_elements = col_buffer_shape.num_elements();

    // If the temporary allocation overhead is too large, fallback on Eigen
    // implementation which requires much less memory.
    int64_t col_buffer_overhead = col_buffer_elements / total_tensor_elements;
    if (col_buffer_overhead > kMaxTempAllocationOverhead) {
      VLOG(2) << "Fallback on Eigen implementation of Conv3DBackpropFilterOp: "
                 "col_buffer_overhead="
              << col_buffer_overhead;

      functor::CuboidConvolutionBackwardFilter<Device, T>()(
          context->eigen_device<Device>(),
          filter_backprop->tensor<T, 5>(),                 // filter_backward
          input.tensor<T, 5>(),                            // input
          out_backprop.tensor<T, 5>(),                     // output_backward
          static_cast<int>(dims.spatial_dims[0].stride),   // stride_planes
          static_cast<int>(dims.spatial_dims[1].stride),   // stride_rows
          static_cast<int>(dims.spatial_dims[2].stride));  // stride_cols

      return;
    }

    Tensor col_buffer;
    OP_REQUIRES_OK(context,
                   context->allocate_temp(DataTypeToEnum<T>::value,
                                          col_buffer_shape, &col_buffer));

    // The input offset corresponding to a single input image.
    const int64_t input_offset =
        dims.spatial_dims[0].input_size * dims.spatial_dims[1].input_size *
        dims.spatial_dims[2].input_size * dims.in_depth;
    // The output offset corresponding to a single output image.
    const int64_t output_offset =
        dims.spatial_dims[0].output_size * dims.spatial_dims[1].output_size *
        dims.spatial_dims[2].output_size * dims.out_depth;

    const T* input_data = input.template flat<T>().data();
    T* col_buffer_data = col_buffer.template flat<T>().data();
    const T* out_backprop_data = out_backprop.template flat<T>().data();
    T* filter_backprop_data = filter_backprop->template flat<T>().data();

    typedef Eigen::TensorMap<Eigen::Tensor<T, 2, Eigen::RowMajor>,
                             Eigen::Unaligned>
        TensorMap;
    typedef Eigen::TensorMap<Eigen::Tensor<const T, 2, Eigen::RowMajor>,
                             Eigen::Unaligned>
        ConstTensorMap;

    TensorMap C(filter_backprop_data, filter_total_size, dims.out_depth);
    C.setZero();

    // Initialize contraction dims (we need to transpose 'A' below).
    Eigen::array<Eigen::IndexPair<Eigen::DenseIndex>, 1> contract_dims;
    contract_dims[0].first = 0;
    contract_dims[0].second = 0;

    auto worker_threads = *(context->device()->tensorflow_cpu_worker_threads());

    for (int image_id = 0; image_id < dims.batch_size; image_id += shard_size) {
      const int shard_limit =
          std::min(static_cast<int>(shard_size),
                   static_cast<int>(dims.batch_size) - image_id);

      auto shard = [&input_data, &col_buffer_data, &dims, &top_pad_planes,
                    &top_pad_rows, &left_pad_cols, &bottom_pad_planes,
                    &bottom_pad_rows, &right_pad_cols, &input_offset,
                    &size_A](int64_t start, int64_t limit) {
        for (int shard_id = start; shard_id < limit; ++shard_id) {
          const T* input_data_shard = input_data + shard_id * input_offset;
          T* col_data_shard = col_buffer_data + shard_id * size_A;

          // When we compute the gradient with respect to the filters, we need
          // to do im2col to allow gemm-type computation.
          Im2col<T>(input_data_shard, dims.in_depth,
                    // Input spatial dimensions.
                    dims.spatial_dims[0].input_size,  // input planes
                    dims.spatial_dims[1].input_size,  // input rows
                    dims.spatial_dims[2].input_size,  // input cols
                    // Filter spatial dimensions.
                    dims.spatial_dims[0].filter_size,  // filter planes
                    dims.spatial_dims[1].filter_size,  // filter rows
                    dims.spatial_dims[2].filter_size,  // filter cols
                    // Spatial padding.
                    top_pad_planes, top_pad_rows, left_pad_cols,
                    bottom_pad_planes, bottom_pad_rows, right_pad_cols,
                    // Spatial striding.
                    dims.spatial_dims[0].stride,  // stride planes
                    dims.spatial_dims[1].stride,  // stride rows
                    dims.spatial_dims[2].stride,  // stride cols
                    col_data_shard);
        }
      };
      Shard(worker_threads.num_threads, worker_threads.workers, shard_limit,
            size_A, shard);

      ConstTensorMap A(col_buffer_data, output_image_size * shard_limit,
                       filter_total_size);
      ConstTensorMap B(out_backprop_data, output_image_size * shard_limit,
                       dims.out_depth);

      // Gradient with respect to filter.
      C.device(context->eigen_cpu_device()) += A.contract(B, contract_dims);

      input_data += input_offset * shard_limit;
      out_backprop_data += output_offset * shard_limit;
    }
  }
  void Compute(OpKernelContext* context) override {
    const Tensor& input = context->input(0);
    const TensorShape& input_shape = input.shape();

    const Tensor& out_backprop = context->input(2);
    const TensorShape& out_backprop_shape = out_backprop.shape();

    TensorShape filter_shape;
    if (takes_shape_) {
      const Tensor& filter_sizes = context->input(1);
      OP_REQUIRES(context, TensorShapeUtils::IsVector(filter_sizes.shape()),
                  errors::InvalidArgument(
                      "filter_sizes shape must be rank 1 but is rank ",
                      filter_sizes.shape().dims()));
      OP_REQUIRES_OK(context, TensorShapeUtils::MakeShape(
                                  filter_sizes.vec<int32>(), &filter_shape));
    } else {
      filter_shape = context->input(1).shape();
    }

    OP_REQUIRES(context, input_shape.dims() == 5,
                errors::InvalidArgument("input tensor must have 5 dimensions"));
    OP_REQUIRES(
        context, filter_shape.dims() == 5,
        errors::InvalidArgument("filter_sizes tensor must have 5 dimensions"));
    OP_REQUIRES(
        context, out_backprop_shape.dims() == 5,
        errors::InvalidArgument("out_backprop tensor must have 5 dimensions"));
    OP_REQUIRES(
        context, input_shape.dim_size(4) == filter_shape.dim_size(3),
        errors::InvalidArgument("input and filter_sizes must have the same "
                                "number of channels. Got ",
                                input_shape.dim_size(4), " for input and ",
                                filter_shape.dim_size(3), " for filter_sizes"));
    OP_REQUIRES(
        context, out_backprop_shape.dim_size(4) == filter_shape.dim_size(4),
        errors::InvalidArgument("out_backprop and filter_sizes must have the "
                                "same number of channels. Got ",
                                out_backprop_shape.dim_size(4),
                                " for out_backprop and ",
                                filter_shape.dim_size(4), " for filter_sizes"));

    ConvBackpropDimensions dims;
    OP_REQUIRES_OK(context,
                   ConvBackpropComputeDimensions(
                       "Conv3DBackpropFilterOp",  3,
                       input_shape, filter_shape, out_backprop_shape, stride_,
                       padding_, data_format_, &dims));

    Tensor* filter_backprop;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, filter_shape, &filter_backprop));

    if (input_shape.num_elements() == 0) {
      filter_backprop->template flat<T>().setZero();
      return;
    }

    int64_t top_pad_planes, bottom_pad_planes;
    int64_t top_pad_rows, bottom_pad_rows;
    int64_t left_pad_cols, right_pad_cols;

    OP_REQUIRES_OK(context, GetWindowedOutputSizeVerbose(
                                dims.spatial_dims[0].input_size,
                                dims.spatial_dims[0].filter_size,
                                dims.spatial_dims[0].stride, padding_,
                                &dims.spatial_dims[0].output_size,
                                &top_pad_planes, &bottom_pad_planes));
    OP_REQUIRES_OK(context, GetWindowedOutputSizeVerbose(
                                dims.spatial_dims[1].input_size,
                                dims.spatial_dims[1].filter_size,
                                dims.spatial_dims[1].stride, padding_,
                                &dims.spatial_dims[1].output_size,
                                &top_pad_rows, &bottom_pad_rows));
    OP_REQUIRES_OK(context, GetWindowedOutputSizeVerbose(
                                dims.spatial_dims[2].input_size,
                                dims.spatial_dims[2].filter_size,
                                dims.spatial_dims[2].stride, padding_,
                                &dims.spatial_dims[2].output_size,
                                &left_pad_cols, &right_pad_cols));

     
     

     
    const int64_t filter_total_size =
        dims.spatial_dims[0].filter_size * dims.spatial_dims[1].filter_size *
        dims.spatial_dims[2].filter_size * dims.in_depth;
     
    const int64_t output_image_size = dims.spatial_dims[0].output_size *
                                      dims.spatial_dims[1].output_size *
                                      dims.spatial_dims[2].output_size;

     
     
     
     

    const auto cache_sizes = Eigen::internal::CacheSizes();
    const ptrdiff_t l3_cache_size = cache_sizes.m_l3;

     
     
     
    const size_t target_working_set_size = l3_cache_size / sizeof(T);

    const int64_t size_A = output_image_size * filter_total_size;

    const int64_t size_B = output_image_size * dims.out_depth;

    const int64_t size_C = filter_total_size * dims.out_depth;

    const int64_t work_unit_size = size_A + size_B + size_C;

    OP_REQUIRES(
        context, work_unit_size > 0,
        errors::InvalidArgument("input, filter_sizes and out_backprop tensors "
                                "must all have at least 1 element"));

    const size_t shard_size =
        (target_working_set_size + work_unit_size - 1) / work_unit_size;

     
    int64_t total_tensor_elements = input_shape.num_elements() +
                                    filter_shape.num_elements() +
                                    out_backprop_shape.num_elements();

     
    TensorShape col_buffer_shape = {static_cast<int64_t>(shard_size),
                                    static_cast<int64_t>(output_image_size),
                                    static_cast<int64_t>(filter_total_size)};
    int64_t col_buffer_elements = col_buffer_shape.num_elements();

     
     
    int64_t col_buffer_overhead = col_buffer_elements / total_tensor_elements;
    if (col_buffer_overhead > kMaxTempAllocationOverhead) {
      VLOG(2) << "Fallback on Eigen implementation of Conv3DBackpropFilterOp: "
                 "col_buffer_overhead="
              << col_buffer_overhead;

      functor::CuboidConvolutionBackwardFilter<Device, T>()(
          context->eigen_device<Device>(),
          filter_backprop->tensor<T, 5>(),                  
          input.tensor<T, 5>(),                             
          out_backprop.tensor<T, 5>(),                      
          static_cast<int>(dims.spatial_dims[0].stride),    
          static_cast<int>(dims.spatial_dims[1].stride),    
          static_cast<int>(dims.spatial_dims[2].stride));   

      return;
    }

    Tensor col_buffer;
    OP_REQUIRES_OK(context,
                   context->allocate_temp(DataTypeToEnum<T>::value,
                                          col_buffer_shape, &col_buffer));

     
    const int64_t input_offset =
        dims.spatial_dims[0].input_size * dims.spatial_dims[1].input_size *
        dims.spatial_dims[2].input_size * dims.in_depth;
     
    const int64_t output_offset =
        dims.spatial_dims[0].output_size * dims.spatial_dims[1].output_size *
        dims.spatial_dims[2].output_size * dims.out_depth;

    const T* input_data = input.template flat<T>().data();
    T* col_buffer_data = col_buffer.template flat<T>().data();
    const T* out_backprop_data = out_backprop.template flat<T>().data();
    T* filter_backprop_data = filter_backprop->template flat<T>().data();

    typedef Eigen::TensorMap<Eigen::Tensor<T, 2, Eigen::RowMajor>,
                             Eigen::Unaligned>
        TensorMap;
    typedef Eigen::TensorMap<Eigen::Tensor<const T, 2, Eigen::RowMajor>,
                             Eigen::Unaligned>
        ConstTensorMap;

    TensorMap C(filter_backprop_data, filter_total_size, dims.out_depth);
    C.setZero();

     
    Eigen::array<Eigen::IndexPair<Eigen::DenseIndex>, 1> contract_dims;
    contract_dims[0].first = 0;
    contract_dims[0].second = 0;

    auto worker_threads = *(context->device()->tensorflow_cpu_worker_threads());

    for (int image_id = 0; image_id < dims.batch_size; image_id += shard_size) {
      const int shard_limit =
          std::min(static_cast<int>(shard_size),
                   static_cast<int>(dims.batch_size) - image_id);

      auto shard = [&input_data, &col_buffer_data, &dims, &top_pad_planes,
                    &top_pad_rows, &left_pad_cols, &bottom_pad_planes,
                    &bottom_pad_rows, &right_pad_cols, &input_offset,
                    &size_A](int64_t start, int64_t limit) {
        for (int shard_id = start; shard_id < limit; ++shard_id) {
          const T* input_data_shard = input_data + shard_id * input_offset;
          T* col_data_shard = col_buffer_data + shard_id * size_A;

           
           
          Im2col<T>(input_data_shard, dims.in_depth,
                     
                    dims.spatial_dims[0].input_size,   
                    dims.spatial_dims[1].input_size,   
                    dims.spatial_dims[2].input_size,   
                     
                    dims.spatial_dims[0].filter_size,   
                    dims.spatial_dims[1].filter_size,   
                    dims.spatial_dims[2].filter_size,   
                     
                    top_pad_planes, top_pad_rows, left_pad_cols,
                    bottom_pad_planes, bottom_pad_rows, right_pad_cols,
                     
                    dims.spatial_dims[0].stride,   
                    dims.spatial_dims[1].stride,   
                    dims.spatial_dims[2].stride,   
                    col_data_shard);
        }
      };
      Shard(worker_threads.num_threads, worker_threads.workers, shard_limit,
            size_A, shard);

      ConstTensorMap A(col_buffer_data, output_image_size * shard_limit,
                       filter_total_size);
      ConstTensorMap B(out_backprop_data, output_image_size * shard_limit,
                       dims.out_depth);

       
      C.device(context->eigen_cpu_device()) += A.contract(B, contract_dims);

      input_data += input_offset * shard_limit;
      out_backprop_data += output_offset * shard_limit;
    }
  }