chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,435 @@
|
||||
// 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/kernels/array_kernel.h"
|
||||
|
||||
#include "paddle/common/layout.h"
|
||||
#include "paddle/phi/backends/cpu/cpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/concat_kernel.h"
|
||||
#include "paddle/phi/kernels/full_kernel.h"
|
||||
#include "paddle/phi/kernels/stack_kernel.h"
|
||||
|
||||
namespace phi {
|
||||
template <typename T, typename Context>
|
||||
void CreateArrayKernel(const Context& dev_ctx,
|
||||
DataType dtype,
|
||||
TensorArray* out) {}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CreateArrayLikeKernel(const Context& dev_ctx,
|
||||
const TensorArray& input,
|
||||
float val,
|
||||
TensorArray* out) {
|
||||
out->resize(input.size());
|
||||
for (size_t i = 0; i < input.size(); i++) {
|
||||
DenseTensor input_i = input[i];
|
||||
out->at(i).Resize(input_i.dims());
|
||||
FullLikeKernel<T, Context>(
|
||||
dev_ctx, input_i, val, input_i.dtype(), &out->at(i));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ArrayLengthKernel(const Context& dev_ctx,
|
||||
const TensorArray& x,
|
||||
DenseTensor* out) {
|
||||
out->Resize({1});
|
||||
dev_ctx.template Alloc<int64_t>(out);
|
||||
*out->data<int64_t>() = static_cast<int64_t>(x.size());
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ArrayReadKernel(const Context& dev_ctx,
|
||||
const TensorArray& array,
|
||||
const Scalar& i,
|
||||
DenseTensor* out) {
|
||||
size_t offset = i.to<int64_t>();
|
||||
PADDLE_ENFORCE_EQ(
|
||||
offset < array.size(),
|
||||
true,
|
||||
errors::InvalidArgument(
|
||||
"index %d exceed array size %d.", offset, array.size()));
|
||||
phi::Copy(dev_ctx, array[offset], dev_ctx.GetPlace(), false, out);
|
||||
out->set_lod(array[offset].lod());
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ArrayWriteKernel(const Context& dev_ctx,
|
||||
const TensorArray& array,
|
||||
const DenseTensor& x,
|
||||
const Scalar& i,
|
||||
TensorArray* out) {
|
||||
size_t offset = i.to<int64_t>();
|
||||
if (offset >= out->size()) {
|
||||
out->resize(offset + 1);
|
||||
}
|
||||
auto* out_tensor = &out->at(offset);
|
||||
out_tensor->set_lod(x.lod());
|
||||
if (x.memory_size() > 0) {
|
||||
phi::Copy(dev_ctx, x, dev_ctx.GetPlace(), false, out_tensor);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ArrayToTensorKernel(const Context& dev_ctx,
|
||||
const TensorArray& x,
|
||||
int axis,
|
||||
bool use_stack,
|
||||
DenseTensor* out,
|
||||
DenseTensor* out_index) {
|
||||
const size_t n = x.size();
|
||||
PADDLE_ENFORCE_GT(
|
||||
n,
|
||||
0,
|
||||
common::errors::InvalidArgument("Input tensor array size should > 0,"
|
||||
"but the received is %d",
|
||||
n));
|
||||
|
||||
std::vector<DenseTensor> tmp_inputs(x.size());
|
||||
std::vector<const DenseTensor*> inputs;
|
||||
|
||||
std::vector<DenseTensor> tmp_indices(x.size());
|
||||
std::vector<const DenseTensor*> indices;
|
||||
|
||||
for (size_t i = 0; i < x.size(); i++) {
|
||||
tmp_inputs[i].ShareDataWith(x[i]);
|
||||
inputs.push_back(&tmp_inputs[i]);
|
||||
FullKernel<int, Context>(
|
||||
dev_ctx, {1}, x[i].dims()[axis], DataType::INT32, &tmp_indices[i]);
|
||||
indices.push_back(&tmp_indices[i]);
|
||||
}
|
||||
|
||||
if (use_stack) {
|
||||
auto vec = vectorize<int>(x[0].dims());
|
||||
vec.insert(vec.begin() + axis, x.size()); // NOLINT
|
||||
out->Resize(vec);
|
||||
StackKernel<T, Context>(dev_ctx, inputs, axis, out);
|
||||
} else {
|
||||
auto out_dims = x[0].dims();
|
||||
size_t in_zero_dims_size = out_dims.size();
|
||||
for (size_t i = 1; i < n; i++) {
|
||||
for (size_t j = 0; j < in_zero_dims_size; j++) {
|
||||
if (j == static_cast<size_t>(axis)) {
|
||||
out_dims[axis] += x[i].dims()[static_cast<int>(j)];
|
||||
}
|
||||
}
|
||||
}
|
||||
auto vec = vectorize<int>(out_dims);
|
||||
out->Resize(vec);
|
||||
ConcatKernel<T, Context>(dev_ctx, inputs, axis, out);
|
||||
}
|
||||
|
||||
out_index->Resize({static_cast<int64_t>(x.size())});
|
||||
StackKernel<int, Context>(dev_ctx, indices, 0, out_index);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ArrayPopKernel(const Context& dev_ctx,
|
||||
const TensorArray& array,
|
||||
int index,
|
||||
TensorArray* array_out,
|
||||
DenseTensor* out) {
|
||||
PADDLE_ENFORCE_GT(
|
||||
array.size(),
|
||||
0,
|
||||
common::errors::InvalidArgument("Input tensorarray size should > 0,"
|
||||
"but the received is %d",
|
||||
array.size()));
|
||||
if (index < 0) {
|
||||
index += array.size();
|
||||
}
|
||||
|
||||
out->ShareDataWith(array[index]);
|
||||
array_out->pop(static_cast<size_t>(index));
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
PD_REGISTER_KERNEL(create_array,
|
||||
CPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CreateArrayKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
PD_REGISTER_KERNEL(create_array,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CreateArrayKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
#endif
|
||||
|
||||
#if defined(PADDLE_WITH_XPU)
|
||||
PD_REGISTER_KERNEL(create_array,
|
||||
XPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CreateArrayKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
#endif
|
||||
|
||||
PD_REGISTER_KERNEL(create_array_like,
|
||||
CPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CreateArrayLikeKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
PD_REGISTER_KERNEL(create_array_like,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CreateArrayLikeKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
#endif
|
||||
|
||||
#if defined(PADDLE_WITH_XPU)
|
||||
PD_REGISTER_KERNEL(create_array_like,
|
||||
XPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CreateArrayLikeKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
#endif
|
||||
|
||||
PD_REGISTER_KERNEL(array_length,
|
||||
CPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArrayLengthKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
PD_REGISTER_KERNEL(array_read,
|
||||
CPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArrayReadKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
PD_REGISTER_KERNEL(array_read,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArrayReadKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
#endif
|
||||
|
||||
#if defined(PADDLE_WITH_XPU)
|
||||
PD_REGISTER_KERNEL(array_read,
|
||||
XPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArrayReadKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
#endif
|
||||
|
||||
PD_REGISTER_KERNEL(array_write,
|
||||
CPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArrayWriteKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
PD_REGISTER_KERNEL(array_write,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArrayWriteKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
#endif
|
||||
|
||||
#if defined(PADDLE_WITH_XPU)
|
||||
PD_REGISTER_KERNEL(array_write,
|
||||
XPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArrayWriteKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
#endif
|
||||
|
||||
PD_REGISTER_KERNEL(array_to_tensor,
|
||||
CPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArrayToTensorKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
PD_REGISTER_KERNEL(array_to_tensor,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArrayToTensorKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
#endif
|
||||
|
||||
#if defined(PADDLE_WITH_XPU)
|
||||
PD_REGISTER_KERNEL(array_to_tensor,
|
||||
XPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArrayToTensorKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
#endif
|
||||
|
||||
PD_REGISTER_KERNEL(array_pop,
|
||||
CPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArrayPopKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
PD_REGISTER_KERNEL(array_pop,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArrayPopKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
#endif
|
||||
|
||||
#if defined(PADDLE_WITH_XPU)
|
||||
PD_REGISTER_KERNEL(array_pop,
|
||||
XPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArrayPopKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
#endif
|
||||
Reference in New Issue
Block a user