chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,428 @@
|
||||
// Copyright (c) 2023 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 <climits>
|
||||
#include <numeric>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/phi/kernels/unique_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/xpu/enforce_xpu.h"
|
||||
#include "paddle/phi/common/memory_utils.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/utils/data_type.h"
|
||||
#include "paddle/phi/core/visit_type.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename Context, typename T, typename IndexT>
|
||||
void XPUFlattenUniqueKernelImpl(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
bool return_index,
|
||||
bool return_inverse,
|
||||
bool return_counts,
|
||||
DenseTensor* out,
|
||||
DenseTensor* indices,
|
||||
DenseTensor* index,
|
||||
DenseTensor* counts) {
|
||||
using XPUType = typename XPUTypeTrait<T>::Type;
|
||||
const auto* x_data = x.data<T>();
|
||||
int64_t x_len = x.numel();
|
||||
int r = 0;
|
||||
xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
|
||||
int64_t unique_len_cpu = 0;
|
||||
int64_t* unique_len_xpu = RAII_GUARD.alloc_l3_or_gm<int64_t>(1);
|
||||
if (x_len != 0) {
|
||||
r = xpu::unique_count<XPUType, IndexT>(
|
||||
dev_ctx.x_context(),
|
||||
reinterpret_cast<const XPUType*>(x_data),
|
||||
unique_len_xpu,
|
||||
x_len,
|
||||
nullptr,
|
||||
nullptr,
|
||||
nullptr,
|
||||
nullptr,
|
||||
nullptr,
|
||||
false);
|
||||
PADDLE_ENFORCE_XDNN_SUCCESS(r, "unique_count");
|
||||
memory_utils::Copy(CPUPlace(),
|
||||
&unique_len_cpu,
|
||||
dev_ctx.GetPlace(),
|
||||
unique_len_xpu,
|
||||
sizeof(int64_t));
|
||||
}
|
||||
out->Resize({unique_len_cpu});
|
||||
auto* out_data = dev_ctx.template Alloc<T>(out);
|
||||
IndexT* indices_data = nullptr;
|
||||
if (return_index) {
|
||||
indices->Resize({unique_len_cpu});
|
||||
indices_data = dev_ctx.template Alloc<IndexT>(indices);
|
||||
}
|
||||
|
||||
IndexT* inverse_data = nullptr;
|
||||
if (return_inverse) {
|
||||
index->Resize({x_len});
|
||||
inverse_data = dev_ctx.template Alloc<IndexT>(index);
|
||||
}
|
||||
|
||||
IndexT* counts_data = nullptr;
|
||||
if (return_counts) {
|
||||
counts->Resize({unique_len_cpu});
|
||||
counts_data = dev_ctx.template Alloc<IndexT>(counts);
|
||||
}
|
||||
if (x_len == 0) {
|
||||
return;
|
||||
}
|
||||
r = xpu::unique_compute<XPUType, IndexT>(
|
||||
dev_ctx.x_context(),
|
||||
reinterpret_cast<const XPUType*>(x_data),
|
||||
reinterpret_cast<XPUType*>(out_data),
|
||||
x_len,
|
||||
unique_len_cpu,
|
||||
indices_data,
|
||||
counts_data,
|
||||
inverse_data,
|
||||
nullptr,
|
||||
nullptr,
|
||||
nullptr,
|
||||
nullptr,
|
||||
nullptr,
|
||||
false);
|
||||
PADDLE_ENFORCE_XDNN_SUCCESS(r, "unique_compute");
|
||||
}
|
||||
|
||||
template <typename Context, typename T, typename IndexT>
|
||||
void XPUDimUniqueKernelImpl(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
bool return_index,
|
||||
bool return_inverse,
|
||||
bool return_counts,
|
||||
int axis,
|
||||
DenseTensor* out,
|
||||
DenseTensor* indices,
|
||||
DenseTensor* index,
|
||||
DenseTensor* counts) {
|
||||
using XPUType = typename XPUTypeTrait<T>::Type;
|
||||
xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
|
||||
int r = 0;
|
||||
const auto* x_data = x.data<T>();
|
||||
auto* x_trans_data = RAII_GUARD.alloc_l3_or_gm<XPUType>(x.numel());
|
||||
std::vector<int64_t> permute(x.dims().size());
|
||||
std::iota(permute.begin(), permute.end(), 0);
|
||||
permute[axis] = 0;
|
||||
permute[0] = axis;
|
||||
if (axis != 0) {
|
||||
auto x_shape = vectorize<int64_t>(x.dims());
|
||||
r = xpu::transpose<XPUType>(dev_ctx.x_context(),
|
||||
reinterpret_cast<const XPUType*>(x_data),
|
||||
x_trans_data,
|
||||
x_shape,
|
||||
permute);
|
||||
PADDLE_ENFORCE_XDNN_SUCCESS(r, "transpose");
|
||||
} else {
|
||||
r = xpu::copy<XPUType>(dev_ctx.x_context(),
|
||||
reinterpret_cast<const XPUType*>(x_data),
|
||||
x_trans_data,
|
||||
x.numel());
|
||||
PADDLE_ENFORCE_XDNN_SUCCESS(r, "copy");
|
||||
}
|
||||
|
||||
DDim x_trans_dims = x.dims();
|
||||
x_trans_dims[0] = x.dims()[axis];
|
||||
x_trans_dims[axis] = x.dims()[0];
|
||||
DDim x_trans_flat_dims = common::flatten_to_2d(x_trans_dims, 1);
|
||||
int64_t axis_len = x_trans_flat_dims[0];
|
||||
int64_t slice_size = x_trans_flat_dims[1];
|
||||
auto x_trans_flat_dims_vec = vectorize<int64_t>(x_trans_flat_dims);
|
||||
|
||||
auto* sorted_axis_idx = RAII_GUARD.alloc_l3_or_gm<IndexT>(axis_len);
|
||||
auto* sort_in_tmp = RAII_GUARD.alloc_l3_or_gm<XPUType>(axis_len);
|
||||
auto* sort_out_tmp = RAII_GUARD.alloc_l3_or_gm<XPUType>(axis_len);
|
||||
auto* x_trans_tmp = RAII_GUARD.alloc_l3_or_gm<XPUType>(x.numel());
|
||||
auto* ori_idx_xpu = RAII_GUARD.alloc_l3_or_gm<IndexT>(axis_len);
|
||||
auto* ori_idx_xpu_tmp = RAII_GUARD.alloc_l3_or_gm<IndexT>(axis_len);
|
||||
auto* sort_offset = RAII_GUARD.alloc_l3_or_gm<IndexT>(axis_len);
|
||||
r = xpu::range<IndexT>(
|
||||
dev_ctx.x_context(), sort_offset, 0, slice_size, axis_len);
|
||||
PADDLE_ENFORCE_XDNN_SUCCESS(r, "range");
|
||||
|
||||
r = xpu::range<IndexT>(dev_ctx.x_context(), ori_idx_xpu, 0, 1, axis_len);
|
||||
PADDLE_ENFORCE_XDNN_SUCCESS(r, "range");
|
||||
|
||||
// radix sort
|
||||
for (int64_t i = slice_size - 1; i >= 0; --i) {
|
||||
r = xpu::paddle_gather<XPUType, IndexT>(dev_ctx.x_context(),
|
||||
x_trans_data + i,
|
||||
sort_offset,
|
||||
sort_in_tmp,
|
||||
{x.numel() - i},
|
||||
axis_len,
|
||||
0);
|
||||
PADDLE_ENFORCE_XDNN_SUCCESS(r, "paddle_gather");
|
||||
r = xpu::stable_sort<XPUType, IndexT>(dev_ctx.x_context(),
|
||||
sort_in_tmp,
|
||||
sort_out_tmp,
|
||||
sorted_axis_idx,
|
||||
1,
|
||||
axis_len,
|
||||
false);
|
||||
PADDLE_ENFORCE_XDNN_SUCCESS(r, "stable_sort");
|
||||
r = xpu::paddle_gather<XPUType, IndexT>(dev_ctx.x_context(),
|
||||
x_trans_data,
|
||||
sorted_axis_idx,
|
||||
x_trans_tmp,
|
||||
x_trans_flat_dims_vec,
|
||||
axis_len,
|
||||
0);
|
||||
PADDLE_ENFORCE_XDNN_SUCCESS(r, "paddle_gather");
|
||||
std::swap(x_trans_data, x_trans_tmp);
|
||||
|
||||
r = xpu::paddle_gather<IndexT, IndexT>(dev_ctx.x_context(),
|
||||
ori_idx_xpu,
|
||||
sorted_axis_idx,
|
||||
ori_idx_xpu_tmp,
|
||||
{axis_len},
|
||||
axis_len,
|
||||
0);
|
||||
PADDLE_ENFORCE_XDNN_SUCCESS(r, "paddle_gather");
|
||||
std::swap(ori_idx_xpu, ori_idx_xpu_tmp);
|
||||
}
|
||||
|
||||
// adjacent difference
|
||||
int64_t compare_num = (axis_len - 1) * slice_size;
|
||||
auto* compare_results = RAII_GUARD.alloc_l3_or_gm<bool>(compare_num);
|
||||
if (compare_num > 0) {
|
||||
r = xpu::broadcast_equal<XPUType>(dev_ctx.x_context(),
|
||||
x_trans_data + slice_size,
|
||||
x_trans_data,
|
||||
compare_results,
|
||||
{compare_num},
|
||||
{compare_num});
|
||||
PADDLE_ENFORCE_XDNN_SUCCESS(r, "broadcast_equal");
|
||||
}
|
||||
|
||||
std::vector<IndexT> unique_axis;
|
||||
std::vector<IndexT> indices_cpu;
|
||||
std::vector<IndexT> inverse_cpu(axis_len);
|
||||
std::vector<IndexT> counts_cpu;
|
||||
std::vector<IndexT> ori_idx_cpu(axis_len);
|
||||
memory_utils::Copy(CPUPlace(),
|
||||
ori_idx_cpu.data(),
|
||||
dev_ctx.GetPlace(),
|
||||
ori_idx_xpu,
|
||||
sizeof(IndexT) * axis_len);
|
||||
unique_axis.push_back(0);
|
||||
indices_cpu.push_back(ori_idx_cpu[0]);
|
||||
inverse_cpu[ori_idx_cpu[0]] = 0;
|
||||
IndexT unique_len = 1;
|
||||
IndexT repeat_cnt = 1;
|
||||
if (axis_len > 1) {
|
||||
DenseTensor adj_identical_cpu;
|
||||
adj_identical_cpu.Resize({axis_len - 1});
|
||||
bool* adj_identical_cpu_data =
|
||||
dev_ctx.template HostAlloc<bool>(&adj_identical_cpu);
|
||||
auto* adj_identical_xpu = RAII_GUARD.alloc_l3_or_gm<bool>(axis_len - 1);
|
||||
r = xpu::reduce_all<bool>(dev_ctx.x_context(),
|
||||
compare_results,
|
||||
adj_identical_xpu,
|
||||
{axis_len - 1, slice_size},
|
||||
{1});
|
||||
PADDLE_ENFORCE_XDNN_SUCCESS(r, "reduce_all");
|
||||
|
||||
memory_utils::Copy(CPUPlace(),
|
||||
adj_identical_cpu_data,
|
||||
dev_ctx.GetPlace(),
|
||||
adj_identical_xpu,
|
||||
(axis_len - 1) * sizeof(bool));
|
||||
|
||||
for (IndexT i = 1; i < axis_len; ++i) {
|
||||
if (!adj_identical_cpu_data[i - 1]) {
|
||||
unique_axis.push_back(i);
|
||||
indices_cpu.push_back(ori_idx_cpu[i]);
|
||||
counts_cpu.push_back(repeat_cnt);
|
||||
++unique_len;
|
||||
repeat_cnt = 1;
|
||||
} else {
|
||||
++repeat_cnt;
|
||||
}
|
||||
inverse_cpu[ori_idx_cpu[i]] = unique_len - 1;
|
||||
}
|
||||
}
|
||||
counts_cpu.push_back(repeat_cnt);
|
||||
DDim out_dims = x.dims();
|
||||
out_dims[axis] = unique_len;
|
||||
out->Resize(out_dims);
|
||||
auto* out_data = dev_ctx.template Alloc<T>(out);
|
||||
|
||||
auto* unique_axis_idx_xpu = RAII_GUARD.alloc_l3_or_gm<IndexT>(unique_len);
|
||||
auto* out_trans_data =
|
||||
RAII_GUARD.alloc_l3_or_gm<XPUType>(unique_len * slice_size);
|
||||
memory_utils::Copy(dev_ctx.GetPlace(),
|
||||
unique_axis_idx_xpu,
|
||||
CPUPlace(),
|
||||
unique_axis.data(),
|
||||
unique_len * sizeof(IndexT));
|
||||
r = xpu::paddle_gather<XPUType, IndexT>(dev_ctx.x_context(),
|
||||
x_trans_data,
|
||||
unique_axis_idx_xpu,
|
||||
out_trans_data,
|
||||
x_trans_flat_dims_vec,
|
||||
unique_len,
|
||||
0);
|
||||
PADDLE_ENFORCE_XDNN_SUCCESS(r, "paddle_gather");
|
||||
DDim out_trans_dims = x_trans_dims;
|
||||
out_trans_dims[0] = unique_len;
|
||||
auto out_trans_dims_vec = vectorize<int64_t>(out_trans_dims);
|
||||
if (axis != 0) {
|
||||
r = xpu::transpose<XPUType>(dev_ctx.x_context(),
|
||||
out_trans_data,
|
||||
reinterpret_cast<XPUType*>(out_data),
|
||||
out_trans_dims_vec,
|
||||
permute);
|
||||
PADDLE_ENFORCE_XDNN_SUCCESS(r, "transpose");
|
||||
} else {
|
||||
r = xpu::copy<XPUType>(dev_ctx.x_context(),
|
||||
out_trans_data,
|
||||
reinterpret_cast<XPUType*>(out_data),
|
||||
unique_len * slice_size);
|
||||
PADDLE_ENFORCE_XDNN_SUCCESS(r, "copy");
|
||||
}
|
||||
if (return_index) {
|
||||
indices->Resize({unique_len});
|
||||
auto* indices_data = dev_ctx.template Alloc<IndexT>(indices);
|
||||
memory_utils::Copy(dev_ctx.GetPlace(),
|
||||
indices_data,
|
||||
CPUPlace(),
|
||||
indices_cpu.data(),
|
||||
sizeof(IndexT) * unique_len);
|
||||
}
|
||||
|
||||
if (return_inverse) {
|
||||
index->Resize({axis_len});
|
||||
auto* reverse_data = dev_ctx.template Alloc<IndexT>(index);
|
||||
memory_utils::Copy(dev_ctx.GetPlace(),
|
||||
reverse_data,
|
||||
CPUPlace(),
|
||||
inverse_cpu.data(),
|
||||
sizeof(IndexT) * axis_len);
|
||||
}
|
||||
|
||||
if (return_counts) {
|
||||
counts->Resize({unique_len});
|
||||
auto* counts_data = dev_ctx.template Alloc<IndexT>(counts);
|
||||
memory_utils::Copy(dev_ctx.GetPlace(),
|
||||
counts_data,
|
||||
CPUPlace(),
|
||||
counts_cpu.data(),
|
||||
sizeof(IndexT) * unique_len);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void UniqueKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
bool return_index,
|
||||
bool return_inverse,
|
||||
bool return_counts,
|
||||
const std::vector<int>& axis,
|
||||
DataType dtype,
|
||||
DenseTensor* out,
|
||||
DenseTensor* indices,
|
||||
DenseTensor* index,
|
||||
DenseTensor* counts) {
|
||||
bool is_sorted = true;
|
||||
UniqueRawKernel<T, Context>(dev_ctx,
|
||||
x,
|
||||
return_index,
|
||||
return_inverse,
|
||||
return_counts,
|
||||
axis,
|
||||
dtype,
|
||||
is_sorted,
|
||||
out,
|
||||
indices,
|
||||
index,
|
||||
counts);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void UniqueRawKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
bool return_index,
|
||||
bool return_inverse,
|
||||
bool return_counts,
|
||||
const std::vector<int>& axis,
|
||||
DataType dtype,
|
||||
bool is_sorted,
|
||||
DenseTensor* out,
|
||||
DenseTensor* indices,
|
||||
DenseTensor* index,
|
||||
DenseTensor* counts) {
|
||||
if (dtype == DataType::INT32) {
|
||||
PADDLE_ENFORCE_LE(
|
||||
x.numel(),
|
||||
INT_MAX,
|
||||
common::errors::InvalidArgument(
|
||||
"The number of elements in Input(X) should be less than or "
|
||||
"equal to INT_MAX, but received num is %d. Please set `dtype` to "
|
||||
"int64.",
|
||||
x.numel()));
|
||||
}
|
||||
|
||||
if (axis.empty()) {
|
||||
PD_VISIT_BASE_INTEGRAL_TYPES(dtype, "XPUFlattenUniqueKernelImpl", [&] {
|
||||
XPUFlattenUniqueKernelImpl<Context, T, data_t>(dev_ctx,
|
||||
x,
|
||||
return_index,
|
||||
return_inverse,
|
||||
return_counts,
|
||||
out,
|
||||
indices,
|
||||
index,
|
||||
counts);
|
||||
});
|
||||
} else {
|
||||
int axis_value = axis[0];
|
||||
axis_value = (axis_value == -1) ? (x.dims().size() - 1) : axis_value;
|
||||
PD_VISIT_BASE_INTEGRAL_TYPES(dtype, "XPUDimUniqueKernelImpl", [&] {
|
||||
XPUDimUniqueKernelImpl<Context, T, data_t>(dev_ctx,
|
||||
x,
|
||||
return_index,
|
||||
return_inverse,
|
||||
return_counts,
|
||||
axis_value,
|
||||
out,
|
||||
indices,
|
||||
index,
|
||||
counts);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(
|
||||
unique, XPU, ALL_LAYOUT, phi::UniqueKernel, float, int, int64_t) {
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::UNDEFINED);
|
||||
kernel->OutputAt(2).SetDataType(phi::DataType::UNDEFINED);
|
||||
kernel->OutputAt(3).SetDataType(phi::DataType::UNDEFINED);
|
||||
}
|
||||
|
||||
PD_REGISTER_KERNEL(
|
||||
unique_raw, XPU, ALL_LAYOUT, phi::UniqueRawKernel, float, int, int64_t) {
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::UNDEFINED);
|
||||
kernel->OutputAt(2).SetDataType(phi::DataType::UNDEFINED);
|
||||
kernel->OutputAt(3).SetDataType(phi::DataType::UNDEFINED);
|
||||
}
|
||||
Reference in New Issue
Block a user