wehub-resource-sync
830 lines
27 KiB
C++
830 lines
27 KiB
C++
// Copyright (c) 2025 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 <vector>
|
|
#include "paddle/common/array.h"
|
|
#include "paddle/phi/backends/context_pool.h"
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/common/int_array.h"
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/common/place.h"
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/phi/kernels/cast_kernel.h"
|
|
#include "paddle/phi/kernels/elementwise_add_kernel.h"
|
|
#include "paddle/phi/kernels/elementwise_kernel.h"
|
|
#include "paddle/phi/kernels/elementwise_multiply_kernel.h"
|
|
#include "paddle/phi/kernels/expand_kernel.h"
|
|
#include "paddle/phi/kernels/full_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/indexing.h"
|
|
#include "paddle/phi/kernels/nonzero_kernel.h"
|
|
#include "paddle/phi/kernels/slice_kernel.h"
|
|
#include "paddle/phi/kernels/transpose_kernel.h"
|
|
|
|
#if defined(__NVCC__) || defined(__HIPCC__)
|
|
#ifdef __NVCC__
|
|
#include <cuda.h>
|
|
#include <cuda_runtime.h>
|
|
#elif defined(__HIPCC__)
|
|
#include <hip/hip_runtime.h>
|
|
#endif
|
|
#endif
|
|
|
|
#ifdef PADDLE_WITH_CUDA
|
|
#include "paddle/phi/backends/gpu/cuda/cuda_graph_with_memory_pool.h"
|
|
#endif
|
|
|
|
namespace phi {
|
|
|
|
namespace funcs {
|
|
|
|
static inline std::vector<int64_t> infer_size_dimvector(
|
|
const std::vector<int64_t>& a, const std::vector<int64_t>& b) {
|
|
// Use ptrdiff_t to ensure signed comparison.
|
|
auto dimsA = a.size();
|
|
auto dimsB = b.size();
|
|
auto ndim = dimsA > dimsB ? dimsA : dimsB;
|
|
std::vector<int64_t> expandedSizes = std::vector<int64_t>(ndim, 0);
|
|
|
|
for (int64_t i = ndim - 1; i >= 0; --i) {
|
|
int64_t offset = ndim - 1 - i;
|
|
int64_t dimA = dimsA - 1 - offset;
|
|
int64_t dimB = dimsB - 1 - offset;
|
|
auto sizeA = (dimA >= 0) ? a[dimA] : 1;
|
|
auto sizeB = (dimB >= 0) ? b[dimB] : 1;
|
|
|
|
expandedSizes[i] = sizeA == 1 ? sizeB : sizeA;
|
|
}
|
|
|
|
return expandedSizes;
|
|
}
|
|
|
|
static inline std::vector<int64_t> compute_strides(
|
|
const std::vector<int64_t>& input_dims, // value_tensor
|
|
const std::vector<int64_t>& input_strides,
|
|
const int64_t& input_elesize,
|
|
const int64_t& ndim,
|
|
const std::vector<int64_t>* shape_,
|
|
std::vector<int64_t>* stride_size) {
|
|
std::vector<int64_t> stride_bytes(ndim, 0);
|
|
const auto& original_shape = input_dims;
|
|
const auto& original_stride = input_strides;
|
|
int64_t element_size_in_bytes = input_elesize;
|
|
int offset = ndim - original_shape.size();
|
|
if (offset > 0)
|
|
stride_bytes.resize(ndim, 0);
|
|
else
|
|
stride_bytes.resize(ndim);
|
|
for (size_t i = 0; i < original_shape.size(); i++) {
|
|
if (original_shape[i] == 1 && (*shape_)[offset + i] != 1) {
|
|
stride_bytes[offset + i] = 0;
|
|
} else {
|
|
stride_bytes[offset + i] = original_stride[i] * element_size_in_bytes;
|
|
}
|
|
}
|
|
stride_size->push_back(stride_bytes.size());
|
|
return stride_bytes;
|
|
}
|
|
|
|
static inline std::vector<int64_t> compute_shapes(
|
|
const std::vector<std::vector<int64_t>>& input_dims) {
|
|
std::vector<int64_t> shape_;
|
|
for (size_t i = 0; i < input_dims.size(); i++) {
|
|
auto shape = input_dims[i];
|
|
if (shape_.empty()) {
|
|
shape_ = shape;
|
|
} else if (!(shape == shape_)) {
|
|
shape_ = infer_size_dimvector(shape_, shape);
|
|
}
|
|
}
|
|
return shape_;
|
|
}
|
|
|
|
template <int N>
|
|
static inline void permute_dimensions(const std::vector<int64_t>& stride_size,
|
|
const std::vector<int64_t>& perm,
|
|
std::array<int64_t*, N>* strides_array,
|
|
std::vector<int64_t>* shape_) {
|
|
auto reorder = [perm](std::vector<int64_t> data) {
|
|
auto res = std::vector<int64_t>(data.size(), 0);
|
|
for (size_t i = 0; i < perm.size(); i++) {
|
|
res[i] = data[perm[i]];
|
|
}
|
|
return res;
|
|
};
|
|
|
|
// Update shape and strides
|
|
*shape_ = reorder(*shape_);
|
|
std::array<std::vector<int64_t>, N> temp_strides;
|
|
for (int64_t i = 0; i < N; i++) {
|
|
if ((*strides_array)[i] != nullptr) {
|
|
std::vector<int64_t> original_data((*strides_array)[i],
|
|
(*strides_array)[i] + stride_size[i]);
|
|
temp_strides[i] = reorder(original_data);
|
|
for (int64_t j = 0; j < stride_size[i]; j++) {
|
|
(*strides_array)[i][j] = temp_strides[i][j];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <int N>
|
|
static inline void reorder_dimensions(const std::vector<int64_t>& stride_size,
|
|
std::vector<int64_t>* shape_,
|
|
std::array<int64_t*, N>* strides_array) {
|
|
// Sort the dimensions based on strides in ascending order with reduced dims
|
|
// at the front. NOTE: that this inverts the order of C-contiguous tensors.
|
|
// strides[0] is the fastest moving dimension instead of strides[ndim - 1].
|
|
// See NOTE: [Computing output strides] and inline comments for more detailed
|
|
// description
|
|
auto ndim = shape_->size();
|
|
std::vector<int64_t> perm_;
|
|
|
|
perm_.resize(ndim);
|
|
if (ndim == 1) {
|
|
perm_[0] = 0;
|
|
return;
|
|
}
|
|
|
|
// initialize perm with n-1, n-2, ..., 1, 0
|
|
std::iota(perm_.rbegin(), perm_.rend(), 0);
|
|
// returns 1 if the dim0 should come after dim1, -1 if dim0 should come
|
|
// before dim1, and 0 if the comparison is ambiguous.
|
|
auto should_swap = [&](size_t dim0, size_t dim1) {
|
|
for (int64_t arg = 0; arg < N; arg++) {
|
|
// ignore undefined or incorrectly sized tensors
|
|
if ((*strides_array)[arg] == nullptr) {
|
|
continue;
|
|
}
|
|
int64_t stride0 = (*strides_array)[arg][dim0];
|
|
int64_t stride1 = (*strides_array)[arg][dim1];
|
|
// move on to the next input if one of the dimensions is broadcasted
|
|
if (stride0 == 0 || stride1 == 0) {
|
|
continue;
|
|
// it is important to return here only with strict comparisons, for
|
|
// equal strides we try to break the tie later by comparing
|
|
// corresponding dimensions or if that does not work, moving on to the
|
|
// next tensor
|
|
} else if (stride0 < stride1) {
|
|
return -1;
|
|
} else if (stride0 > stride1) {
|
|
return 1;
|
|
} else { // equal strides, use dimensions themselves as the tie-breaker.
|
|
// at this point, with zero strides out of the way, we are guaranteed
|
|
// that operand dimensions are equal to shape_
|
|
auto t_dim0 = (*shape_)[dim0];
|
|
auto t_dim1 = (*shape_)[dim1];
|
|
// return only if dimensions should be swapped, otherwise move on to the
|
|
// next tensor
|
|
if (t_dim0 > t_dim1) {
|
|
return 1;
|
|
}
|
|
}
|
|
}
|
|
return 0;
|
|
};
|
|
// insertion sort with support for ambiguous comparisons
|
|
for (size_t i = 1; i < ndim; i++) {
|
|
int dim1 = i;
|
|
for (int dim0 = i - 1; dim0 >= 0; dim0--) {
|
|
int comparison = should_swap(perm_[dim0], perm_[dim1]);
|
|
if (comparison > 0) {
|
|
std::swap(perm_[dim0], perm_[dim1]);
|
|
dim1 = dim0;
|
|
} else if (comparison < 0) {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
// perform re-ordering of shape and strides
|
|
permute_dimensions<N>(stride_size, perm_, strides_array, shape_);
|
|
}
|
|
|
|
static inline std::vector<int64_t> compatible_stride(
|
|
const std::vector<int64_t>* shape_,
|
|
const int64_t& ndim,
|
|
const int64_t& element_size) {
|
|
std::vector<int64_t> stride;
|
|
int64_t next_stride = element_size;
|
|
|
|
for (int64_t dim = 0; dim < ndim; ++dim) {
|
|
stride.push_back(next_stride);
|
|
next_stride *= (*shape_)[dim];
|
|
}
|
|
return stride;
|
|
}
|
|
|
|
template <int N>
|
|
static inline void allocate_or_resize_outputs(
|
|
const std::vector<int64_t>* shape_,
|
|
const int64_t element_size,
|
|
const int64_t ndim,
|
|
std::array<int64_t*, N>* strides_array) {
|
|
std::vector<int64_t> stride_bytes =
|
|
compatible_stride(shape_, ndim, static_cast<int64_t>(element_size));
|
|
|
|
if (strides_array && (*strides_array)[0]) {
|
|
std::copy(stride_bytes.begin(), stride_bytes.end(), (*strides_array)[0]);
|
|
}
|
|
}
|
|
|
|
template <int N>
|
|
static inline void coalesce_dimensions(const int64_t& ndim,
|
|
std::array<int64_t*, N>* strides_array,
|
|
std::vector<int64_t>* stride_size,
|
|
std::vector<int64_t>* shape_) {
|
|
if (ndim <= 1) {
|
|
return;
|
|
}
|
|
|
|
// We can coalesce two adjacent dimensions if either dim has size 1 or if:
|
|
// shape[n] * stride[n] == stride[n + 1].
|
|
auto can_coalesce = [&](int dim0, int dim1) {
|
|
auto shape0 = (*shape_)[dim0];
|
|
auto shape1 = (*shape_)[dim1];
|
|
if (shape0 == 1 || shape1 == 1) {
|
|
return true;
|
|
}
|
|
for (int64_t i = 0; i < N; i++) {
|
|
auto& stride = (*strides_array)[i];
|
|
if (shape0 * stride[dim0] != stride[dim1]) {
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
};
|
|
|
|
// replace each operands stride at dim0 with its stride at dim1
|
|
auto replace_stride = [&](int dim0, int dim1) {
|
|
for (int64_t i = 0; i < N; i++) {
|
|
auto& stride = (*strides_array)[i];
|
|
stride[dim0] = stride[dim1];
|
|
}
|
|
};
|
|
|
|
int prev_dim = 0;
|
|
for (int64_t dim = 1; dim < ndim; dim++) {
|
|
if (can_coalesce(prev_dim, dim)) {
|
|
if ((*shape_)[prev_dim] == 1) {
|
|
replace_stride(prev_dim, dim);
|
|
}
|
|
(*shape_)[prev_dim] *= (*shape_)[dim];
|
|
} else {
|
|
prev_dim++;
|
|
if (prev_dim != dim) {
|
|
replace_stride(prev_dim, dim);
|
|
(*shape_)[prev_dim] = (*shape_)[dim];
|
|
}
|
|
}
|
|
}
|
|
(*shape_).resize(prev_dim + 1);
|
|
for (int64_t i = 0; i < N; i++) {
|
|
(*stride_size)[i] = shape_->size();
|
|
}
|
|
}
|
|
|
|
template <int N>
|
|
static inline void CopyStride(
|
|
const std::vector<int64_t>& output_dims, // value_tensor
|
|
const std::vector<int64_t>& output_strides,
|
|
const int64_t& output_elesize,
|
|
const std::vector<int64_t>& input_dims, // input_tensor
|
|
const std::vector<int64_t>& input_strides,
|
|
const int64_t& input_elesize,
|
|
std::vector<int64_t>* desired_shape,
|
|
std::array<int64_t*, N>* strides_array,
|
|
int64_t* numel,
|
|
std::array<std::vector<int64_t>, N>& strides_vec) { // NOLINT
|
|
int ndim = output_dims.size();
|
|
|
|
std::vector<int64_t> stride_size;
|
|
|
|
*desired_shape = compute_shapes({input_dims, output_dims});
|
|
strides_vec[0] = compute_strides(output_dims, // input_tensor
|
|
output_strides,
|
|
output_elesize,
|
|
ndim,
|
|
desired_shape,
|
|
&stride_size);
|
|
|
|
strides_vec[1] = compute_strides(input_dims, // value_tensor
|
|
input_strides,
|
|
input_elesize,
|
|
ndim,
|
|
desired_shape,
|
|
&stride_size);
|
|
|
|
for (size_t i = 0; i < N; i++) {
|
|
(*strides_array)[i] = strides_vec[i].data();
|
|
}
|
|
reorder_dimensions<N>(stride_size, desired_shape, strides_array);
|
|
|
|
coalesce_dimensions<N>(ndim, strides_array, &stride_size, desired_shape);
|
|
|
|
int64_t num = 1;
|
|
for (size_t i = 0; i < desired_shape->size(); i++) {
|
|
num *= (*desired_shape)[i];
|
|
}
|
|
*numel = num;
|
|
}
|
|
|
|
template <int N>
|
|
static inline void IndexPutStride(
|
|
const std::vector<int64_t>& output_dims, // input_tensor
|
|
const std::vector<int64_t>& output_strides,
|
|
const int64_t& output_elesize,
|
|
const std::vector<int64_t>& input_dims, // value_tensor
|
|
const std::vector<int64_t>& input_strides,
|
|
const int64_t& input_elesize,
|
|
const std::vector<int64_t>& index_dims, // index_tensor
|
|
const std::vector<int64_t>& index_strides,
|
|
const int64_t& index_elesize,
|
|
std::vector<int64_t>* desired_shape,
|
|
std::array<int64_t*, N>* strides_array,
|
|
int64_t* numel,
|
|
std::array<std::vector<int64_t>, N>& strides_vec) { // NOLINT
|
|
int ndim = output_dims.size();
|
|
|
|
std::vector<int64_t> stride_size;
|
|
|
|
*desired_shape = compute_shapes({input_dims, output_dims, index_dims});
|
|
strides_vec[0] = compute_strides(output_dims, // input_tensor
|
|
output_strides,
|
|
output_elesize,
|
|
ndim,
|
|
desired_shape,
|
|
&stride_size);
|
|
|
|
strides_vec[1] = compute_strides(input_dims, // value_tensor
|
|
input_strides,
|
|
input_elesize,
|
|
ndim,
|
|
desired_shape,
|
|
&stride_size);
|
|
|
|
strides_vec[2] = compute_strides(index_dims, // index_tensor
|
|
index_strides,
|
|
index_elesize,
|
|
ndim,
|
|
desired_shape,
|
|
&stride_size);
|
|
|
|
for (size_t i = 0; i < N; i++) {
|
|
(*strides_array)[i] = strides_vec[i].data();
|
|
}
|
|
reorder_dimensions<N>(stride_size, desired_shape, strides_array);
|
|
|
|
coalesce_dimensions<N>(ndim, strides_array, &stride_size, desired_shape);
|
|
|
|
int64_t num = 1;
|
|
for (size_t i = 0; i < desired_shape->size(); i++) {
|
|
num *= (*desired_shape)[i];
|
|
}
|
|
*numel = num;
|
|
}
|
|
|
|
template <int N>
|
|
static inline void IndexGetStride(
|
|
const std::vector<int64_t>& output_dims,
|
|
const std::vector<int64_t>& output_strides,
|
|
const int64_t& output_elesize,
|
|
const std::vector<int64_t>& input_dims,
|
|
const std::vector<int64_t>& input_strides,
|
|
const int64_t& input_elesize,
|
|
const std::vector<int64_t>& index_dims,
|
|
const std::vector<int64_t>& index_strides,
|
|
const int64_t& index_elesize,
|
|
std::vector<int64_t>* desired_shape,
|
|
std::array<int64_t*, N>* strides_array,
|
|
int64_t* numel,
|
|
std::array<std::vector<int64_t>, N>& strides_vec) { // NOLINT
|
|
int ndim = output_dims.size();
|
|
|
|
std::vector<int64_t> stride_size;
|
|
|
|
*desired_shape = compute_shapes({input_dims, output_dims, index_dims});
|
|
|
|
strides_vec[0] = compute_strides(input_dims,
|
|
input_strides,
|
|
input_elesize,
|
|
ndim,
|
|
desired_shape,
|
|
&stride_size);
|
|
|
|
strides_vec[1] = compute_strides(output_dims,
|
|
output_strides,
|
|
output_elesize,
|
|
ndim,
|
|
desired_shape,
|
|
&stride_size);
|
|
|
|
strides_vec[2] = compute_strides(index_dims,
|
|
index_strides,
|
|
index_elesize,
|
|
ndim,
|
|
desired_shape,
|
|
&stride_size);
|
|
|
|
for (size_t i = 0; i < N; i++) {
|
|
(*strides_array)[i] = strides_vec[i].data();
|
|
}
|
|
reorder_dimensions<N>(stride_size, desired_shape, strides_array);
|
|
|
|
allocate_or_resize_outputs<N>(
|
|
desired_shape, output_elesize, ndim, strides_array);
|
|
|
|
coalesce_dimensions<N>(ndim, strides_array, &stride_size, desired_shape);
|
|
|
|
int64_t num = 1;
|
|
for (size_t i = 0; i < desired_shape->size(); i++) {
|
|
num *= (*desired_shape)[i];
|
|
}
|
|
*numel = num;
|
|
}
|
|
|
|
static inline void cal_shape_stride(const std::vector<int64_t>& index_dims,
|
|
int64_t* num_indices,
|
|
std::vector<int64_t>* shape_tmp,
|
|
std::vector<int64_t>* stride_tmp) {
|
|
std::vector<int64_t> index_dims_;
|
|
std::vector<int64_t> index_stride_;
|
|
|
|
bool tmp_flag = false;
|
|
for (unsigned i = 0; i < index_dims.size(); i++) {
|
|
if (index_dims[i] == -1) {
|
|
if (!tmp_flag) {
|
|
*num_indices = i;
|
|
tmp_flag = true;
|
|
continue;
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
|
|
if (!tmp_flag) {
|
|
index_dims_.push_back(index_dims[i]);
|
|
} else {
|
|
shape_tmp->push_back(index_dims[i]);
|
|
}
|
|
}
|
|
|
|
int shape_size = shape_tmp->size();
|
|
stride_tmp->resize(shape_size);
|
|
if (shape_size > 0) {
|
|
(*stride_tmp)[shape_size - 1] = 1;
|
|
}
|
|
if (shape_size > 1) {
|
|
for (int i = shape_size - 2; i >= 0; i--) {
|
|
(*stride_tmp)[i] = (*stride_tmp)[i + 1] * (*shape_tmp)[i + 1];
|
|
}
|
|
}
|
|
}
|
|
|
|
template <int N>
|
|
static inline void ScatterAddStride(
|
|
const std::vector<int64_t>& output_dims,
|
|
const std::vector<int64_t>& output_strides,
|
|
const int64_t& output_elesize,
|
|
const std::vector<int64_t>& input_dims,
|
|
const std::vector<int64_t>& input_strides,
|
|
const int64_t& input_elesize,
|
|
const std::vector<int64_t>& index_dims,
|
|
const std::vector<int64_t>& index_strides,
|
|
const int64_t& index_elesize,
|
|
std::vector<int64_t>* desired_shape,
|
|
std::array<int64_t*, N>* strides_array,
|
|
int64_t* numel,
|
|
std::array<std::vector<int64_t>, N>& strides_vec) { // NOLINT
|
|
int ndim = output_dims.size();
|
|
|
|
std::vector<int64_t> stride_size;
|
|
|
|
*desired_shape = compute_shapes({input_dims, output_dims, index_dims});
|
|
|
|
strides_vec[0] = compute_strides(input_dims,
|
|
input_strides,
|
|
input_elesize,
|
|
ndim,
|
|
desired_shape,
|
|
&stride_size);
|
|
|
|
strides_vec[1] = compute_strides(output_dims,
|
|
output_strides,
|
|
output_elesize,
|
|
ndim,
|
|
desired_shape,
|
|
&stride_size);
|
|
|
|
strides_vec[2] = compute_strides(index_dims,
|
|
index_strides,
|
|
index_elesize,
|
|
ndim,
|
|
desired_shape,
|
|
&stride_size);
|
|
|
|
for (size_t i = 0; i < N; i++) {
|
|
(*strides_array)[i] = strides_vec[i].data();
|
|
}
|
|
|
|
reorder_dimensions<N>(stride_size, desired_shape, strides_array);
|
|
|
|
coalesce_dimensions<N>(ndim, strides_array, &stride_size, desired_shape);
|
|
|
|
int64_t num = 1;
|
|
for (size_t i = 0; i < desired_shape->size(); i++) {
|
|
num *= (*desired_shape)[i];
|
|
}
|
|
*numel = num;
|
|
}
|
|
|
|
static inline bool hasContiguousSubspace(const std::vector<DenseTensor>& tl) {
|
|
auto isDefined = [](const DenseTensor& tensor) {
|
|
return tensor.initialized();
|
|
};
|
|
auto isNull = [](const DenseTensor& tensor) { return !tensor.initialized(); };
|
|
|
|
auto start = std::find_if(tl.begin(), tl.end(), isDefined);
|
|
auto stop = std::find_if(tl.rbegin(), tl.rend(), isDefined);
|
|
auto it = std::find_if(start, stop.base(), isNull);
|
|
return it == stop.base();
|
|
}
|
|
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
|
|
static inline std::vector<DenseTensor> expandTensors(
|
|
const GPUContext& dev_ctx, const std::vector<const DenseTensor*>& indices) {
|
|
std::vector<DenseTensor> result;
|
|
for (const auto& index : indices) {
|
|
if (index == nullptr) {
|
|
result.emplace_back();
|
|
continue;
|
|
}
|
|
|
|
if (index->dtype() == DataType::BOOL) {
|
|
DenseTensor bool_2_idx;
|
|
phi::NonZeroKernel<bool, GPUContext>(dev_ctx, *index, &bool_2_idx);
|
|
|
|
for (int j = 0; j < index->dims().size(); ++j) {
|
|
DenseTensor sliced_tensor;
|
|
phi::SliceKernel<int64_t, GPUContext>(
|
|
dev_ctx, bool_2_idx, {1}, {j}, {j + 1}, {1}, {}, &sliced_tensor);
|
|
result.emplace_back(sliced_tensor);
|
|
}
|
|
} else {
|
|
result.emplace_back(*index);
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
|
|
static inline std::vector<DenseTensor> expand_outplace(
|
|
const GPUContext& dev_ctx, const std::vector<DenseTensor>& to_expand) {
|
|
bool first = true;
|
|
DDim target_shape;
|
|
for (size_t i = 0; i < to_expand.size(); ++i) {
|
|
if (!to_expand[i].initialized()) continue;
|
|
if (first) {
|
|
target_shape = to_expand[i].dims();
|
|
first = false;
|
|
} else {
|
|
target_shape = InferSizeSymdimvector(target_shape, to_expand[i].dims());
|
|
}
|
|
}
|
|
|
|
std::vector<DenseTensor> result(to_expand.size());
|
|
for (size_t i = 0; i < to_expand.size(); ++i) {
|
|
if (!to_expand[i].initialized()) continue;
|
|
if (to_expand[i].dims() == target_shape) {
|
|
result[i] = to_expand[i];
|
|
} else {
|
|
phi::ExpandKernel<float, GPUContext>(
|
|
dev_ctx,
|
|
to_expand[i],
|
|
phi::IntArray(vectorize<int64_t>(target_shape)),
|
|
&result[i]);
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
|
|
template <typename T>
|
|
inline std::tuple<DenseTensor, std::vector<DenseTensor>, std::vector<int64_t>>
|
|
transposeToFrontAndInvPerm(const GPUContext& dev_ctx,
|
|
const DenseTensor& self,
|
|
const std::vector<DenseTensor>& indices) {
|
|
std::vector<int> dims;
|
|
std::vector<int64_t> inv_perm;
|
|
std::vector<DenseTensor> transposed_indices;
|
|
dims.reserve(self.dims().size());
|
|
inv_perm.resize(self.dims().size());
|
|
|
|
for (int i = 0; i < static_cast<int>(self.dims().size()); ++i) {
|
|
if (indices[i].initialized()) {
|
|
dims.push_back(i);
|
|
transposed_indices.emplace_back(indices[i]);
|
|
}
|
|
}
|
|
|
|
for (int i = 0; i < static_cast<int>(self.dims().size()); ++i) {
|
|
if (!indices[i].initialized()) {
|
|
dims.push_back(i);
|
|
transposed_indices.emplace_back();
|
|
}
|
|
}
|
|
|
|
for (int i = 0; i < static_cast<int>(self.dims().size()); ++i) {
|
|
inv_perm[dims[i]] = i;
|
|
}
|
|
|
|
DenseTensor transposed_self;
|
|
TransposeKernel<T, GPUContext>(dev_ctx, self, dims, &transposed_self);
|
|
|
|
return std::make_tuple(transposed_self, transposed_indices, inv_perm);
|
|
}
|
|
|
|
static inline std::vector<int64_t> computeLinearStride(
|
|
const DenseTensor& tensor) {
|
|
auto sizes = vectorize<int64_t>(tensor.dims());
|
|
std::vector<int64_t> stride(sizes.size());
|
|
if (stride.empty()) {
|
|
return stride;
|
|
}
|
|
stride.back() = 1;
|
|
std::partial_sum(sizes.rbegin(),
|
|
sizes.rend() - 1,
|
|
stride.rbegin() + 1,
|
|
std::multiplies<int64_t>());
|
|
return stride;
|
|
}
|
|
|
|
static inline DenseTensor wrapIndexOnce(const GPUContext& dev_ctx,
|
|
const DenseTensor& index,
|
|
const int64_t& dim,
|
|
const int64_t& dim_size,
|
|
bool check_range) {
|
|
DenseTensor dim_size_tensor;
|
|
dim_size_tensor.Resize(index.dims());
|
|
dev_ctx.Alloc<int64_t>(&dim_size_tensor);
|
|
|
|
auto* dim_size_data = dim_size_tensor.data<int64_t>();
|
|
auto numel = index.numel();
|
|
std::vector<int64_t> host_data(numel, dim_size);
|
|
const int64_t* stable_hd =
|
|
phi::backends::gpu::RestoreHostMemIfCapturingCUDAGraph(host_data.data(),
|
|
host_data.size());
|
|
phi::memory_utils::Copy(dev_ctx.GetPlace(),
|
|
dim_size_data,
|
|
CPUPlace(),
|
|
stable_hd,
|
|
numel * sizeof(int64_t),
|
|
dev_ctx.stream());
|
|
|
|
return phi::Remainder<int64_t>(dev_ctx, index, dim_size_tensor);
|
|
}
|
|
|
|
static inline std::tuple<DenseTensor, int64_t, int64_t, int64_t>
|
|
computeLinearIndex(const GPUContext& dev_ctx,
|
|
const DenseTensor& src,
|
|
const std::vector<DenseTensor>& indices,
|
|
bool check_range) {
|
|
std::vector<int64_t> strides = computeLinearStride(src);
|
|
DenseTensor linearIndex;
|
|
int64_t nElemBefore = 1, nElemAfter = 1, strideBefore = 0;
|
|
|
|
for (int64_t i = 0; i < src.dims().size(); ++i) {
|
|
if (indices[i].initialized()) {
|
|
auto wrapped_index =
|
|
wrapIndexOnce(dev_ctx, indices[i], i, src.dims()[i], check_range);
|
|
|
|
auto strides_tensor = phi::Full<int64_t, GPUContext>(
|
|
dev_ctx,
|
|
vectorize<int64_t>(wrapped_index.dims()),
|
|
phi::Scalar(strides[i]));
|
|
|
|
auto scaled_index = phi::Multiply<int64_t, GPUContext>(
|
|
dev_ctx, wrapped_index, strides_tensor);
|
|
|
|
if (linearIndex.initialized()) {
|
|
phi::AddKernel<int64_t, GPUContext>(
|
|
dev_ctx, linearIndex, scaled_index, &linearIndex);
|
|
} else {
|
|
linearIndex = scaled_index;
|
|
if (i > 0) {
|
|
strideBefore = src.strides()[i - 1];
|
|
}
|
|
}
|
|
} else if (linearIndex.initialized()) {
|
|
nElemAfter *= src.dims()[i];
|
|
} else {
|
|
nElemBefore *= src.dims()[i];
|
|
}
|
|
}
|
|
|
|
return std::make_tuple(
|
|
std::move(linearIndex), nElemBefore, strideBefore, nElemAfter);
|
|
}
|
|
|
|
template <typename T>
|
|
static inline std::tuple<DenseTensor,
|
|
DenseTensor,
|
|
int64_t,
|
|
int64_t,
|
|
int64_t,
|
|
std::vector<int64_t>>
|
|
makeLinearIndex(const GPUContext& dev_ctx,
|
|
const DenseTensor& self,
|
|
const std::vector<const DenseTensor*>& orig,
|
|
bool check_range) {
|
|
auto indices = expandTensors(dev_ctx, orig);
|
|
for (auto& idx : indices) {
|
|
if (idx.initialized() && idx.dtype() == DataType::INT32) {
|
|
idx = Cast<int32_t, GPUContext>(dev_ctx, idx, DataType::INT64);
|
|
}
|
|
}
|
|
indices = expand_outplace(dev_ctx, std::move(indices));
|
|
|
|
while (indices.size() < static_cast<size_t>(self.dims().size())) {
|
|
indices.emplace_back();
|
|
}
|
|
|
|
std::vector<int64_t> inverse_perm;
|
|
DenseTensor transposed_self = self;
|
|
std::vector<DenseTensor> transposed_indices;
|
|
std::vector<int64_t> inv_perm;
|
|
if (!hasContiguousSubspace(indices)) {
|
|
auto [tmp_self, tmp_indices, tmp_perm] =
|
|
transposeToFrontAndInvPerm<T>(dev_ctx, self, indices);
|
|
transposed_self = std::move(tmp_self);
|
|
transposed_indices = std::move(tmp_indices);
|
|
inv_perm = std::move(tmp_perm);
|
|
} else {
|
|
transposed_indices = indices;
|
|
}
|
|
|
|
auto [linear_index, n_elem_before, stride_before, n_elem_after] =
|
|
computeLinearIndex(
|
|
dev_ctx, transposed_self, transposed_indices, check_range);
|
|
|
|
return std::make_tuple(linear_index,
|
|
transposed_self,
|
|
n_elem_before,
|
|
stride_before,
|
|
n_elem_after,
|
|
inv_perm);
|
|
}
|
|
|
|
#endif
|
|
inline bool are_expandable(const std::vector<int64_t>& shape1,
|
|
const std::vector<int64_t>& shape2) {
|
|
size_t ndim1 = shape1.size();
|
|
size_t ndim2 = shape2.size();
|
|
size_t ndim = std::min(ndim1, ndim2);
|
|
|
|
for (int64_t i = static_cast<int64_t>(ndim) - 1; i >= 0; --i) {
|
|
auto dim1 = shape1[--ndim1];
|
|
auto dim2 = shape2[--ndim2];
|
|
if (dim1 == dim2 || dim1 == 1 || dim2 == 1) {
|
|
continue;
|
|
}
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
inline int64_t LargestIndex(const DenseTensor& tensor) {
|
|
int64_t result = 0;
|
|
const auto& dims = tensor.dims();
|
|
const auto& strides = tensor.strides();
|
|
|
|
for (int i = 0; i < dims.size(); ++i) {
|
|
result += (dims[i] - 1) * strides[i];
|
|
}
|
|
return result;
|
|
}
|
|
|
|
inline int GetNumBits(uint64_t max_val) {
|
|
if (max_val == 0) return 1;
|
|
|
|
int num_bits = 1;
|
|
while (max_val > 1) {
|
|
max_val >>= 1;
|
|
num_bits++;
|
|
}
|
|
return num_bits;
|
|
}
|
|
} // namespace funcs
|
|
} // namespace phi
|