wehub-resource-sync
133 lines
4.0 KiB
C++
133 lines
4.0 KiB
C++
// 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
|
|
|
|
#if defined(__NVCC__) || defined(__HIPCC__)
|
|
#include <thrust/device_vector.h>
|
|
#include <thrust/host_vector.h>
|
|
#endif
|
|
|
|
#include "paddle/phi/kernels/diag_embed_kernel.h"
|
|
|
|
#include <algorithm>
|
|
|
|
#include "paddle/phi/kernels/funcs/for_range.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T>
|
|
struct DiagEmbedFunctor {
|
|
DiagEmbedFunctor(const T* input,
|
|
int64_t numel,
|
|
const int64_t* dim,
|
|
int64_t offset,
|
|
int64_t dims_size,
|
|
T* output,
|
|
const int64_t* strides)
|
|
: input_(input),
|
|
numel_(numel),
|
|
dim_(dim),
|
|
offset_(offset),
|
|
dims_size_(dims_size),
|
|
output_(output),
|
|
strides_(strides) {}
|
|
|
|
HOSTDEVICE void operator()(size_t idx) const {
|
|
int64_t position = 0;
|
|
auto numel = numel_;
|
|
int64_t num = idx;
|
|
for (int64_t i = 0; i < dims_size_; i++) {
|
|
numel = numel / dim_[i];
|
|
position += num / numel * strides_[i];
|
|
num = num % numel;
|
|
}
|
|
output_[position + offset_] = input_[idx];
|
|
}
|
|
|
|
const T* input_;
|
|
int64_t numel_;
|
|
const int64_t* dim_;
|
|
int64_t offset_;
|
|
int64_t dims_size_;
|
|
T* output_;
|
|
const int64_t* strides_;
|
|
};
|
|
|
|
template <typename T, typename Context>
|
|
void DiagEmbedKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
int offset,
|
|
int dim1,
|
|
int dim2,
|
|
DenseTensor* out) {
|
|
auto* input_data = x.data<T>();
|
|
T* out_data = dev_ctx.template Alloc<T>(out);
|
|
if (out && out->numel() == 0) {
|
|
return;
|
|
}
|
|
funcs::SetConstant<Context, T> set_zero;
|
|
|
|
set_zero(dev_ctx, out, static_cast<T>(0.0));
|
|
|
|
auto out_dims = out->dims();
|
|
int dim1_ = dim1 < 0 ? out_dims.size() + dim1 : dim1;
|
|
int dim2_ = dim2 < 0 ? out_dims.size() + dim2 : dim2;
|
|
auto stride = common::stride(out_dims);
|
|
int64_t diag_size;
|
|
int64_t storage_offset = 0;
|
|
if (offset >= 0) {
|
|
int64_t dim = out_dims[dim2_] - offset;
|
|
diag_size = std::max<int64_t>(std::min(out_dims[dim1_], dim), 0);
|
|
} else {
|
|
int64_t dim = out_dims[dim1_] + offset;
|
|
diag_size = std::max<int64_t>(std::min(dim, out_dims[dim2_]), 0);
|
|
}
|
|
if (diag_size == 0) {
|
|
// skip
|
|
} else if (offset >= 0) {
|
|
storage_offset += offset * stride[dim2_];
|
|
} else {
|
|
storage_offset -= offset * stride[dim1_];
|
|
}
|
|
auto strides = vectorize(stride);
|
|
strides.erase(strides.begin() + std::max(dim1_, dim2_));
|
|
strides.erase(strides.begin() + std::min(dim1_, dim2_));
|
|
strides.push_back(stride[dim1_] + stride[dim2_]);
|
|
const auto dims = vectorize(x.dims());
|
|
|
|
#if defined(__NVCC__) || defined(__HIPCC__)
|
|
thrust::device_vector<int64_t> dims_vec(dims);
|
|
const int64_t* dims_arr = thrust::raw_pointer_cast(dims_vec.data());
|
|
thrust::device_vector<int64_t> strides_vec(strides);
|
|
const int64_t* strides_arr = thrust::raw_pointer_cast(strides_vec.data());
|
|
#else
|
|
const int64_t* dims_arr = dims.data();
|
|
const int64_t* strides_arr = strides.data();
|
|
#endif
|
|
|
|
funcs::ForRange<Context> for_range(dev_ctx, x.numel());
|
|
DiagEmbedFunctor<T> functor(input_data,
|
|
x.numel(),
|
|
dims_arr,
|
|
storage_offset,
|
|
dims.size(),
|
|
out_data,
|
|
strides_arr);
|
|
for_range(functor);
|
|
}
|
|
|
|
} // namespace phi
|