conv_kernel.cc 8.2 KB
Newer Older
Z
zhangkaihuo 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
/* 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. */

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

namespace phi {
namespace sparse {

/**
 * x: (N, D, H, W, C)
 * kernel: (D, H, W, C, OC)
 * out: (N, D, H, W, OC)
29
 **/
30
template <typename T, typename IntT = int>
Z
zhangkaihuo 已提交
31 32 33 34 35 36
void Conv3dCooCPUKernel(const CPUContext& dev_ctx,
                        const SparseCooTensor& x,
                        const DenseTensor& kernel,
                        const std::vector<int>& paddings,
                        const std::vector<int>& dilations,
                        const std::vector<int>& strides,
37
                        const int groups UNUSED,
Z
zhangkaihuo 已提交
38
                        const bool subm,
39
                        const std::string& key,
Z
zhangkaihuo 已提交
40
                        SparseCooTensor* out,
41 42
                        DenseTensor* rulebook,
                        DenseTensor* counter) {
Z
zhangkaihuo 已提交
43 44 45 46 47
  // update padding and dilation
  // Currently, only support x.layout is NDHWC, groups = 1
  // if x.layout != NDHWC then transpose(x), transpose(weight)

  const auto& x_dims = x.dims();
48
  const bool is2D = x_dims.size() == 4 ? true : false;
Z
zhangkaihuo 已提交
49
  const auto& kernel_dims = kernel.dims();
50 51 52 53 54 55 56
  int kernel_size = is2D ? kernel_dims[0] * kernel_dims[1]
                         : kernel_dims[0] * kernel_dims[1] * kernel_dims[2];

  int count_tmp = is2D ? 4 : 5;
  std::vector<int> out_dims_vec(count_tmp, 1);
  DDim out_dims = make_ddim(out_dims_vec);

Z
zhangkaihuo 已提交
57 58 59 60 61
  std::vector<int> kernel_sizes(kernel_dims.size());
  for (int i = 0; i < kernel_dims.size(); i++) {
    kernel_sizes[i] = kernel_dims[i];
  }

62 63
  std::vector<int> subm_paddings(paddings), subm_strides(strides);
  if (subm) {
64 65
    // the out shape of subm_conv is same as input shape
    // reset the padding=kernel_size/2 and strides=1
66 67 68
    phi::funcs::sparse::ResetSubmKernelSizeAndStrides(
        kernel.dims(), &subm_paddings, &subm_strides);
  }
69 70 71

  phi::funcs::sparse::GetOutShape(
      x_dims, kernel_sizes, subm_paddings, dilations, subm_strides, &out_dims);
72 73
  const int in_channels = is2D ? kernel_dims[2] : kernel_dims[3];
  const int out_channels = is2D ? kernel_dims[3] : kernel_dims[4];
74

Z
zhangkaihuo 已提交
75 76 77
  // Second algorithm:
  // https://pdfs.semanticscholar.org/5125/a16039cabc6320c908a4764f32596e018ad3.pdf
  // 1. product rulebook
78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120
  DenseTensor h_counter, h_offsets;
  h_counter.Resize({kernel_size});
  h_offsets.Resize({kernel_size + 1});
  int* h_counter_ptr = dev_ctx.template HostAlloc<int>(&h_counter);
  int* h_offsets_ptr = dev_ctx.template HostAlloc<int>(&h_offsets);

  // DenseTensor* rulebook = nullptr;
  const IntT* rulebook_ptr = nullptr;
  int n = 0;
  bool need_product_rulebook = true;
  if (subm && !key.empty()) {
    rulebook_ptr = phi::funcs::sparse::PrepareSubm<T, IntT, CPUContext>(
        dev_ctx,
        x,
        key,
        out_dims,
        out,
        h_counter_ptr,
        h_offsets_ptr,
        &n,
        &need_product_rulebook);
  }
  if (need_product_rulebook) {
    DenseTensor tmp_rulebook;
    ProductRuleBook<T, CPUContext, IntT>(dev_ctx,
                                         x,
                                         kernel_sizes,
                                         subm_paddings,
                                         dilations,
                                         subm_strides,
                                         out_dims,
                                         subm,
                                         &tmp_rulebook,
                                         h_counter_ptr);

    UpdateRulebookAndOutIndex<T, CPUContext, IntT>(
        dev_ctx, x, kernel_size, out_channels, out_dims, &tmp_rulebook, out);
    n = tmp_rulebook.dims()[1];
    rulebook_ptr = tmp_rulebook.data<IntT>();

    phi::funcs::sparse::SaveToTable(
        dev_ctx, x, key, tmp_rulebook, h_counter, out, rulebook, counter);
  }
Z
zhangkaihuo 已提交
121 122 123 124 125 126 127 128 129 130 131 132 133

  // 2. gather
  DenseTensorMeta in_features_meta(
      x.dtype(), {n, in_channels}, DataLayout::NHWC);
  DenseTensorMeta out_features_meta(
      x.dtype(), {n, out_channels}, DataLayout::NHWC);
  phi::DenseTensor in_features =
      phi::Empty(dev_ctx, std::move(in_features_meta));
  phi::DenseTensor out_features =
      phi::Empty(dev_ctx, std::move(out_features_meta));
  T* in_features_ptr = in_features.data<T>();
  T* out_features_ptr = out_features.data<T>();

134 135
  Gather<T, IntT>(
      x.values().data<T>(), rulebook_ptr + n, n, in_channels, in_features_ptr);
Z
zhangkaihuo 已提交
136 137

  // 3. call gemm for every werght
138
  auto blas = phi::funcs::GetBlas<CPUContext, T>(dev_ctx);
Z
zhangkaihuo 已提交
139 140
  int offset = 0;
  for (int i = 0; i < kernel_size; i++) {
141 142
    h_offsets_ptr[i] = offset;
    offset += h_counter_ptr[i];
Z
zhangkaihuo 已提交
143
  }
144
  h_offsets_ptr[kernel_size] = offset;
Z
zhangkaihuo 已提交
145 146 147

  const T* kernel_ptr = kernel.data<T>();
  for (int i = 0; i < kernel_size; i++) {
148
    if (h_counter_ptr[i] <= 0) {
Z
zhangkaihuo 已提交
149 150 151 152
      continue;
    }

    // call gemm: (n, in_channels) * (in_channels, out_channels)
153
    const int M = h_counter_ptr[i];
Z
zhangkaihuo 已提交
154 155
    const int K = in_channels;   // in_channels
    const int N = out_channels;  // out_channels
156
    T* tmp_in_ptr = in_features_ptr + h_offsets_ptr[i] * in_channels;
Z
zhangkaihuo 已提交
157
    const T* tmp_kernel_ptr = kernel_ptr + i * K * N;
158
    T* tmp_out_ptr = out_features_ptr + h_offsets_ptr[i] * out_channels;
Z
zhangkaihuo 已提交
159 160 161 162 163 164 165 166 167 168 169 170 171
    blas.GEMM(CblasNoTrans,
              CblasNoTrans,
              M,
              N,
              K,
              static_cast<T>(1),
              tmp_in_ptr,
              tmp_kernel_ptr,
              static_cast<T>(0),
              tmp_out_ptr);
  }

  // 4. scatter
172
  T* out_values_ptr = out->mutable_values()->data<T>();
Z
zhangkaihuo 已提交
173
  memset(out_values_ptr, 0, sizeof(T) * out->nnz() * out_channels);
174 175
  Scatter<T, IntT>(
      out_features_ptr, rulebook_ptr + n * 2, n, out_channels, out_values_ptr);
176 177 178
}

template <typename T, typename Context>
Z
zhangkaihuo 已提交
179 180 181 182 183 184 185 186
void Conv3dCooKernel(const Context& dev_ctx,
                     const SparseCooTensor& x,
                     const DenseTensor& kernel,
                     const std::vector<int>& paddings,
                     const std::vector<int>& dilations,
                     const std::vector<int>& strides,
                     const int groups,
                     const bool subm,
187
                     const std::string& key,
Z
zhangkaihuo 已提交
188
                     SparseCooTensor* out,
189 190
                     DenseTensor* rulebook,
                     DenseTensor* counter) {
191 192 193 194 195 196 197 198 199 200 201 202 203 204
  PD_VISIT_BASE_INTEGRAL_TYPES(x.indices().dtype(), "Conv3dCooCPUKernel", ([&] {
                                 Conv3dCooCPUKernel<T, data_t>(dev_ctx,
                                                               x,
                                                               kernel,
                                                               paddings,
                                                               dilations,
                                                               strides,
                                                               groups,
                                                               subm,
                                                               key,
                                                               out,
                                                               rulebook,
                                                               counter);
                               }));
Z
zhangkaihuo 已提交
205 206 207 208 209
}
}  // namespace sparse
}  // namespace phi

PD_REGISTER_KERNEL(
Z
zhangkaihuo 已提交
210
    conv3d_coo, CPU, ALL_LAYOUT, phi::sparse::Conv3dCooKernel, float, double) {
Z
zhangkaihuo 已提交
211
  kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
212 213 214
  kernel->OutputAt(0).SetDataType(paddle::DataType::UNDEFINED);
  kernel->OutputAt(1).SetDataType(paddle::DataType::INT32);
  kernel->OutputAt(2).SetDataType(paddle::DataType::INT32);
Z
zhangkaihuo 已提交
215
}