void SpatialMaxPool(OpKernelContext* context, Tensor* output,
                      const Tensor& tensor_in, const PoolParameters& params,
                      const Padding& padding) {
    // On GPU, use Eigen's Spatial Max Pooling.  On CPU, use an
    // EigenMatrix version that is currently faster than Eigen's
    // Spatial MaxPooling implementation.
    //
    // TODO(vrv): Remove this once we no longer need it.
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
    if (std::is_same<Device, GPUDevice>::value) {
      Eigen::PaddingType pt = BrainPadding2EigenPadding(padding);
      if (std::is_same<T, qint8>::value) {
        LaunchMaxPoolingNoMask_NCHW_VECT_C<GPUDevice>::launch(
            context, params, tensor_in, output);
      } else {
        functor::SpatialMaxPooling<Device, T>()(
            context->eigen_device<Device>(), output->tensor<T, 4>(),
            tensor_in.tensor<T, 4>(), params.window_rows, params.window_cols,
            params.row_stride, params.col_stride, pt);
      }
    } else
#endif
    {
      typedef Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>
          ConstEigenMatrixMap;
      typedef Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>
          EigenMatrixMap;

      ConstEigenMatrixMap in_mat(tensor_in.flat<T>().data(), params.depth,
                                 params.tensor_in_cols * params.tensor_in_rows *
                                     params.tensor_in_batch);
      EigenMatrixMap out_mat(
          output->flat<T>().data(), params.depth,
          params.out_width * params.out_height * params.tensor_in_batch);

      const DeviceBase::CpuWorkerThreads& worker_threads =
          *(context->device()->tensorflow_cpu_worker_threads());

      // The following code basically does the following:
      // 1. Flattens the input and output tensors into two dimensional arrays.
      //    tensor_in_as_matrix:
      //      depth by (tensor_in_cols * tensor_in_rows * tensor_in_batch)
      //    output_as_matrix:
      //      depth by (out_width * out_height * tensor_in_batch)
      //
      // 2. Walks through the set of columns in the flattened
      // tensor_in_as_matrix,
      //    and updates the corresponding column(s) in output_as_matrix with the
      //    max value.
      auto shard = [&params, &in_mat, &out_mat](int64_t start, int64_t limit) {
        const int32_t in_rows = params.tensor_in_rows;
        const int32_t in_cols = params.tensor_in_cols;
        const int32_t pad_top = params.pad_top;
        const int32_t pad_left = params.pad_left;
        const int32_t window_rows = params.window_rows;
        const int32_t window_cols = params.window_cols;
        const int32_t row_stride = params.row_stride;
        const int32_t col_stride = params.col_stride;
        const int32_t out_height = params.out_height;
        const int32_t out_width = params.out_width;

        {
          // Initializes the output tensor with MIN<T>.
          const int32_t output_image_size =
              out_height * out_width * params.depth;
          EigenMatrixMap out_shard(out_mat.data() + start * output_image_size,
                                   1, (limit - start) * output_image_size);
          out_shard.setConstant(Eigen::NumTraits<T>::lowest());
        }

        for (int32_t b = start; b < limit; ++b) {
          const int32_t out_offset_batch = b * out_height;
          for (int32_t h = 0; h < in_rows; ++h) {
            for (int32_t w = 0; w < in_cols; ++w) {
              // (h_start, h_end) * (w_start, w_end) is the range that the input
              // vector projects to.
              const int32_t hpad = h + pad_top;
              const int32_t wpad = w + pad_left;
              const int32_t h_start =
                  (hpad < window_rows) ? 0
                                       : (hpad - window_rows) / row_stride + 1;
              const int32_t h_end = std::min(hpad / row_stride + 1, out_height);
              const int32_t w_start =
                  (wpad < window_cols) ? 0
                                       : (wpad - window_cols) / col_stride + 1;
              const int32_t w_end = std::min(wpad / col_stride + 1, out_width);
              // compute elementwise max
              const int32_t in_offset = (b * in_rows + h) * in_cols + w;
              for (int32_t ph = h_start; ph < h_end; ++ph) {
                const int32_t out_offset_base =
                    (out_offset_batch + ph) * out_width;
                for (int32_t pw = w_start; pw < w_end; ++pw) {
                  const int32_t out_offset = out_offset_base + pw;
                  out_mat.col(out_offset) =
                      out_mat.col(out_offset).cwiseMax(in_mat.col(in_offset));
                }
              }
            }
          }
        }
      };

      // TODO(andydavis) Consider sharding across batch x rows x cols.
      // TODO(andydavis) Consider a higher resolution shard cost model.
      const int64_t shard_cost =
          params.tensor_in_rows * params.tensor_in_cols * params.depth;
      Shard(worker_threads.num_threads, worker_threads.workers,
            params.tensor_in_batch, shard_cost, shard);
    }
  }