pool_kernel.cc 5.2 KB
Newer Older
Z
zhangkaihuo 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

15
#include "paddle/phi/kernels/sparse/pool_kernel.h"
16

Z
zhangkaihuo 已提交
17 18
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/core/tensor_meta.h"
19
#include "paddle/phi/core/visit_type.h"
Z
zhangkaihuo 已提交
20 21
#include "paddle/phi/kernels/funcs/pooling.h"
#include "paddle/phi/kernels/funcs/sparse/convolution.h"
22
#include "paddle/phi/kernels/sparse/cpu/conv.h"
Z
zhangkaihuo 已提交
23 24 25 26 27 28 29 30

namespace phi {
namespace sparse {

/**
 * x: (N, D, H, W, C)
 * kernel: (D, H, W, C, OC)
 * out: (N, D, H, W, OC)
31
 **/
32
template <typename T, typename IntT = int>
33 34 35 36 37 38 39
void MaxPoolCooCPUKernel(const CPUContext& dev_ctx,
                         const SparseCooTensor& x,
                         const std::vector<int>& kernel_sizes,
                         const std::vector<int>& paddings,
                         const std::vector<int>& dilations,
                         const std::vector<int>& strides,
                         SparseCooTensor* out,
40 41
                         DenseTensor* rulebook,
                         DenseTensor* counter) {
Z
zhangkaihuo 已提交
42 43 44 45 46 47 48 49 50
  const auto& x_dims = x.dims();
  int kernel_size = kernel_sizes[0] * kernel_sizes[1] * kernel_sizes[2];
  const std::vector<int>& real_kernel_sizes =
      phi::funcs::sparse::PoolResetKernel(kernel_sizes, x_dims[4], x_dims[4]);
  DDim out_dims = {1, 1, 1, 1, 1};
  phi::funcs::sparse::GetOutShape(
      x_dims, real_kernel_sizes, paddings, dilations, strides, &out_dims);
  const int in_channels = real_kernel_sizes[3];

51
  std::vector<int> counter_per_kernel(kernel_size, 0);
Z
zhangkaihuo 已提交
52 53 54

  const T* in_features_ptr = x.non_zero_elements().data<T>();
  // 1. product rule book
55 56 57 58 59 60 61 62 63
  ProductRuleBook<T, CPUContext, IntT>(dev_ctx,
                                       x,
                                       real_kernel_sizes,
                                       paddings,
                                       dilations,
                                       strides,
                                       out_dims,
                                       false,
                                       rulebook,
64
                                       counter_per_kernel.data());
65 66

  UpdateRulebookAndOutIndex<T, CPUContext, IntT>(
Z
zhangkaihuo 已提交
67 68 69
      dev_ctx, x, kernel_size, in_channels, out_dims, rulebook, out);

  int rulebook_len = rulebook->dims()[1];
70
  const IntT* rulebook_ptr = rulebook->data<IntT>();
71 72 73 74

  counter->Resize({kernel_size});
  int* counter_ptr = dev_ctx.template HostAlloc<int>(counter);
  memcpy(counter_ptr, counter_per_kernel.data(), kernel_size * sizeof(int));
Z
zhangkaihuo 已提交
75 76 77 78 79 80 81 82 83 84

  std::vector<int> offsets(kernel_size + 1);
  phi::funcs::sparse::PrefixSum(counter_ptr, &offsets[0], kernel_size);
  std::vector<bool> out_flags(out->nnz(), false);

  // 2. max pool
  T* out_features_ptr = out->mutable_non_zero_elements()->data<T>();
  phi::funcs::MaxPool<T> max_pool_functor;
  for (int i = 0; i < kernel_size; i++) {
    for (int j = 0; j < counter_ptr[i]; j++) {
85 86
      IntT in_i = rulebook_ptr[rulebook_len + offsets[i] + j];
      IntT out_i = rulebook_ptr[rulebook_len * 2 + offsets[i] + j];
Z
zhangkaihuo 已提交
87 88 89 90 91 92 93 94 95 96 97 98 99 100 101
      if (!out_flags[out_i]) {
        out_flags[out_i] = true;
        memcpy(&out_features_ptr[out_i * in_channels],
               &in_features_ptr[in_i * in_channels],
               in_channels * sizeof(T));
      } else {
        for (int c = 0; c < in_channels; c++) {
          max_pool_functor.compute(in_features_ptr[in_i * in_channels + c],
                                   &out_features_ptr[out_i * in_channels + c]);
        }
      }
    }
  }
}

102
template <typename T, typename Context>
103 104 105 106 107 108 109
void MaxPoolCooKernel(const Context& dev_ctx,
                      const SparseCooTensor& x,
                      const std::vector<int>& kernel_sizes,
                      const std::vector<int>& paddings,
                      const std::vector<int>& dilations,
                      const std::vector<int>& strides,
                      SparseCooTensor* out,
110 111
                      DenseTensor* rulebook,
                      DenseTensor* counter) {
112
  PD_VISIT_INTEGRAL_TYPES(
113 114 115 116 117 118 119 120
      x.non_zero_indices().dtype(), "MaxPoolCooCPUKernel", ([&] {
        MaxPoolCooCPUKernel<T, data_t>(dev_ctx,
                                       x,
                                       kernel_sizes,
                                       paddings,
                                       dilations,
                                       strides,
                                       out,
121 122
                                       rulebook,
                                       counter);
123 124 125
      }));
}

Z
zhangkaihuo 已提交
126 127 128
}  // namespace sparse
}  // namespace phi

129
PD_REGISTER_KERNEL(maxpool_coo,
Z
zhangkaihuo 已提交
130 131
                   CPU,
                   ALL_LAYOUT,
132
                   phi::sparse::MaxPoolCooKernel,
Z
zhangkaihuo 已提交
133 134 135 136
                   float,
                   double) {
  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
}