chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,267 @@
|
||||
/* 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. */
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <set>
|
||||
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/sparse_coo_tensor.h"
|
||||
#include "paddle/phi/core/tensor_meta.h"
|
||||
#include "paddle/phi/kernels/funcs/blas/blas.h"
|
||||
#include "paddle/phi/kernels/sparse/conv_kernel.h"
|
||||
|
||||
namespace phi {
|
||||
namespace sparse {
|
||||
|
||||
using Dims4D = funcs::sparse::Dims4D;
|
||||
|
||||
// such as: kernel(3, 3, 3), kernel_size = 27
|
||||
// counter_per_weight: (kernel_size)
|
||||
// TODO(zhangkaihuo): optimize performance with multithreading
|
||||
template <typename T, typename Context, typename IntT = int>
|
||||
void ProductRuleBook(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,
|
||||
const DDim& out_dims,
|
||||
const bool subm,
|
||||
DenseTensor* rulebook,
|
||||
int* counter_per_kernel) {
|
||||
const bool is2D = out_dims.size() == 4 ? true : false;
|
||||
const int64_t non_zero_num = x.nnz();
|
||||
const auto& indices = x.indices();
|
||||
const IntT* indices_ptr = indices.data<IntT>();
|
||||
int kernel_size = is2D ? kernel_sizes[0] * kernel_sizes[1]
|
||||
: kernel_sizes[0] * kernel_sizes[1] * kernel_sizes[2];
|
||||
memset(counter_per_kernel, 0, kernel_size * sizeof(int));
|
||||
|
||||
int rulebook_len = 0;
|
||||
// calc the rulebook_len
|
||||
const auto& x_dims = x.dims();
|
||||
|
||||
int xdim0, xdim1, xdim2, xdim3;
|
||||
int kdim0, kdim1, kdim2, kdim3;
|
||||
int odim0, odim1, odim2, odim3;
|
||||
int pdim0, pdim1, pdim2, pdim3;
|
||||
int sdim0, sdim1, sdim2, sdim3;
|
||||
int ddim0, ddim1, ddim2, ddim3;
|
||||
|
||||
xdim0 = x_dims[0];
|
||||
xdim1 = is2D ? x_dims[2] : x_dims[3];
|
||||
xdim2 = is2D ? x_dims[1] : x_dims[2];
|
||||
xdim3 = is2D ? 1 : x_dims[1];
|
||||
|
||||
kdim0 = 1;
|
||||
kdim1 = is2D ? kernel_sizes[1] : kernel_sizes[2];
|
||||
kdim2 = is2D ? kernel_sizes[0] : kernel_sizes[1];
|
||||
kdim3 = is2D ? 1 : kernel_sizes[0];
|
||||
|
||||
odim0 = out_dims[0];
|
||||
odim1 = is2D ? out_dims[2] : out_dims[3];
|
||||
odim2 = is2D ? out_dims[1] : out_dims[2];
|
||||
odim3 = is2D ? 1 : out_dims[1];
|
||||
|
||||
pdim0 = 1;
|
||||
pdim1 = is2D ? paddings[1] : paddings[2];
|
||||
pdim2 = is2D ? paddings[0] : paddings[1];
|
||||
pdim3 = is2D ? 1 : paddings[0];
|
||||
|
||||
sdim0 = 1;
|
||||
sdim1 = is2D ? strides[1] : strides[2];
|
||||
sdim2 = is2D ? strides[0] : strides[1];
|
||||
sdim3 = is2D ? 1 : strides[0];
|
||||
|
||||
ddim0 = 1;
|
||||
ddim1 = is2D ? dilations[1] : dilations[2];
|
||||
ddim2 = is2D ? dilations[0] : dilations[1];
|
||||
ddim3 = is2D ? 1 : dilations[0];
|
||||
|
||||
const Dims4D c_x_dims(xdim0, xdim1, xdim2, xdim3);
|
||||
const Dims4D c_kernel_dims(kdim0, kdim1, kdim2, kdim3);
|
||||
const Dims4D c_out_dims(odim0, odim1, odim2, odim3);
|
||||
const Dims4D c_paddings(pdim0, pdim1, pdim2, pdim3);
|
||||
const Dims4D c_strides(sdim0, sdim1, sdim2, sdim3);
|
||||
const Dims4D c_dilations(ddim0, ddim1, ddim2, ddim3);
|
||||
|
||||
std::set<IntT> hash_in;
|
||||
if (subm) {
|
||||
for (int i = 0; i < non_zero_num; i++) {
|
||||
IntT batch = indices_ptr[i];
|
||||
IntT in_z = is2D ? 0 : indices_ptr[i + non_zero_num];
|
||||
IntT in_y = is2D ? indices_ptr[i + non_zero_num]
|
||||
: indices_ptr[i + 2 * non_zero_num];
|
||||
IntT in_x = is2D ? indices_ptr[i + 2 * non_zero_num]
|
||||
: indices_ptr[i + 3 * non_zero_num];
|
||||
IntT index = funcs::sparse::PointToIndex<Dims4D>(
|
||||
batch, in_x, in_y, in_z, c_x_dims);
|
||||
hash_in.insert(index);
|
||||
}
|
||||
}
|
||||
|
||||
auto f_calc_rulebook = [&](IntT* rulebook_ptr) {
|
||||
int kernel_index = 0, rulebook_index = 0;
|
||||
int zceil = is2D ? 1 : kernel_sizes[0];
|
||||
int yceil = is2D ? kernel_sizes[0] : kernel_sizes[1];
|
||||
int xceil = is2D ? kernel_sizes[1] : kernel_sizes[2];
|
||||
for (int kz = 0; kz < zceil; kz++) {
|
||||
for (int ky = 0; ky < yceil; ky++) {
|
||||
for (int kx = 0; kx < xceil; kx++) {
|
||||
++kernel_index;
|
||||
for (int64_t i = 0; i < non_zero_num; i++) {
|
||||
IntT batch = indices_ptr[i];
|
||||
IntT in_z = is2D ? 0 : indices_ptr[i + non_zero_num];
|
||||
IntT in_y = is2D ? indices_ptr[i + non_zero_num]
|
||||
: indices_ptr[i + 2 * non_zero_num];
|
||||
IntT in_x = is2D ? indices_ptr[i + 2 * non_zero_num]
|
||||
: indices_ptr[i + 3 * non_zero_num];
|
||||
|
||||
IntT out_z =
|
||||
is2D ? 0
|
||||
: (in_z + paddings[0] - kz * dilations[0]) / strides[0];
|
||||
IntT out_y =
|
||||
(in_y + c_paddings[2] - ky * c_dilations[2]) / c_strides[2];
|
||||
IntT out_x =
|
||||
(in_x + c_paddings[3] - kx * c_dilations[3]) / c_strides[3];
|
||||
if (funcs::sparse::Check(c_x_dims,
|
||||
c_kernel_dims,
|
||||
c_paddings,
|
||||
c_dilations,
|
||||
c_strides,
|
||||
in_x,
|
||||
in_y,
|
||||
in_z,
|
||||
kx,
|
||||
ky,
|
||||
kz)) {
|
||||
if (subm) {
|
||||
IntT out_index = funcs::sparse::PointToIndex<Dims4D>(
|
||||
batch, out_x, out_y, out_z, c_out_dims);
|
||||
if (hash_in.find(out_index) == hash_in.end()) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
if (rulebook_ptr == nullptr) {
|
||||
counter_per_kernel[kernel_index - 1] += 1;
|
||||
++rulebook_len;
|
||||
} else {
|
||||
rulebook_ptr[rulebook_index] = kernel_index - 1;
|
||||
rulebook_ptr[rulebook_index + rulebook_len] = i; // in_i
|
||||
rulebook_ptr[rulebook_index + rulebook_len * 2] =
|
||||
funcs::sparse::PointToIndex<Dims4D>(
|
||||
batch, out_x, out_y, out_z, c_out_dims); // out_index
|
||||
++rulebook_index;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
f_calc_rulebook(nullptr);
|
||||
// alloc the rulebook
|
||||
*rulebook = Empty(dev_ctx,
|
||||
DenseTensorMeta(phi::CppTypeToDataType<IntT>::Type(),
|
||||
{3, rulebook_len},
|
||||
DataLayout::NCHW));
|
||||
IntT* rulebook_ptr = rulebook->data<IntT>();
|
||||
f_calc_rulebook(rulebook_ptr);
|
||||
}
|
||||
|
||||
template <typename T, typename Context, typename IntT = int>
|
||||
void UpdateRulebookAndOutIndex(const Context& dev_ctx,
|
||||
const SparseCooTensor& x,
|
||||
const int kernel_size UNUSED,
|
||||
const int out_channels,
|
||||
const DDim& out_dims,
|
||||
DenseTensor* rulebook,
|
||||
SparseCooTensor* out) {
|
||||
const bool is2D = out_dims.size() == 4 ? true : false;
|
||||
|
||||
std::set<IntT> tmp_indices;
|
||||
int64_t n = rulebook->dims()[1];
|
||||
IntT* rulebook_ptr = rulebook->data<IntT>();
|
||||
for (int64_t i = 0; i < n; i++) {
|
||||
tmp_indices.insert(rulebook_ptr[i + n * 2]);
|
||||
}
|
||||
|
||||
int out_non_zero_num = tmp_indices.size();
|
||||
const int64_t sparse_dim = is2D ? 3 : 4;
|
||||
DenseTensorMeta indices_meta(phi::CppTypeToDataType<IntT>::Type(),
|
||||
{sparse_dim, out_non_zero_num},
|
||||
DataLayout::NCHW);
|
||||
DenseTensorMeta values_meta(
|
||||
x.dtype(), {out_non_zero_num, out_channels}, x.values().layout());
|
||||
DenseTensor out_indices = Empty(dev_ctx, std::move(indices_meta));
|
||||
DenseTensor out_values = Empty(dev_ctx, std::move(values_meta));
|
||||
IntT* out_indices_ptr = out_indices.data<IntT>();
|
||||
int64_t idx = 0;
|
||||
|
||||
int odim0, odim1, odim2, odim3;
|
||||
odim0 = out_dims[0];
|
||||
odim1 = is2D ? out_dims[2] : out_dims[3];
|
||||
odim2 = is2D ? out_dims[1] : out_dims[2];
|
||||
odim3 = is2D ? 1 : out_dims[1];
|
||||
const Dims4D c_out_dims(odim0, odim1, odim2, odim3);
|
||||
|
||||
for (auto it = tmp_indices.begin(); it != tmp_indices.end(); it++, idx++) {
|
||||
const IntT index = *it;
|
||||
IntT batch, x, y, z;
|
||||
funcs::sparse::IndexToPoint<Dims4D>(index, c_out_dims, &batch, &x, &y, &z);
|
||||
out_indices_ptr[idx] = batch;
|
||||
if (is2D) {
|
||||
out_indices_ptr[idx + out_non_zero_num] = y;
|
||||
out_indices_ptr[idx + out_non_zero_num * 2] = x;
|
||||
} else {
|
||||
out_indices_ptr[idx + out_non_zero_num] = z;
|
||||
out_indices_ptr[idx + out_non_zero_num * 2] = y;
|
||||
out_indices_ptr[idx + out_non_zero_num * 3] = x;
|
||||
}
|
||||
}
|
||||
for (int64_t i = 0; i < n; i++) {
|
||||
IntT out_index = rulebook_ptr[i + n * 2];
|
||||
rulebook_ptr[i + n * 2] =
|
||||
std::distance(tmp_indices.begin(), tmp_indices.find(out_index));
|
||||
}
|
||||
|
||||
out->SetMember(out_indices, out_values, out_dims, true);
|
||||
}
|
||||
|
||||
template <typename T, typename IntT = int>
|
||||
void Gather(
|
||||
const T* x, const IntT* indices, const int n, const int channels, T* out) {
|
||||
for (int i = 0; i < n; i++) {
|
||||
IntT real_i = indices[i];
|
||||
memcpy(out + i * channels, x + real_i * channels, channels * sizeof(T));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IntT = int>
|
||||
void Scatter(
|
||||
const T* x, const IntT* indices, const int n, const int channels, T* out) {
|
||||
for (int i = 0; i < n; i++) {
|
||||
IntT real_i = indices[i];
|
||||
for (int j = 0; j < channels; j++) {
|
||||
out[real_i * channels + j] += x[i * channels + j];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace sparse
|
||||
} // namespace phi
|
||||
Reference in New Issue
Block a user