82 lines
2.7 KiB
C++
82 lines
2.7 KiB
C++
// 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 "paddle/phi/kernels/funcs/strided_reshape_utils.h"
|
|
#include <array>
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/reshape_kernel.h"
|
|
|
|
namespace phi {
|
|
|
|
bool ReshapeStride(const DDim& old_dims,
|
|
const DDim& old_stride,
|
|
const DDim& new_dims,
|
|
DDim& new_stride) { // NOLINT
|
|
int64_t numel = product(old_dims);
|
|
if (numel < 0) {
|
|
std::array<int64_t, 2> tmp = {};
|
|
tmp[0] = 1;
|
|
tmp[1] = new_dims.size();
|
|
new_stride = DDim(tmp.data(), 2);
|
|
return true;
|
|
} else if (numel == 0) {
|
|
if (old_dims == new_dims) {
|
|
new_stride = old_stride;
|
|
} else {
|
|
new_stride = new_dims;
|
|
new_stride[new_dims.size() - 1] = 1;
|
|
for (int i = new_dims.size() - 2; i >= 0; i--) {
|
|
new_stride[i] = new_stride[i + 1] *
|
|
std::max(static_cast<int64_t>(1), new_dims[i + 1]);
|
|
}
|
|
}
|
|
return true;
|
|
} else {
|
|
int64_t old_numel = 1;
|
|
int64_t new_numel = 1;
|
|
int64_t old_stride_lastvalue = old_stride[old_stride.size() - 1];
|
|
int new_stride_index = new_dims.size() - 1;
|
|
new_stride = new_dims;
|
|
for (int old_dims_index = old_dims.size() - 1; old_dims_index >= 0;
|
|
old_dims_index--) {
|
|
old_numel *= old_dims[old_dims_index];
|
|
if ((old_dims_index == 0) || (old_dims[old_dims_index - 1] != 1 &&
|
|
old_stride[old_dims_index - 1] !=
|
|
old_numel * old_stride_lastvalue)) {
|
|
while (new_stride_index >= 0 &&
|
|
(new_numel < old_numel || new_dims[new_stride_index] == 1)) {
|
|
new_stride[new_stride_index] = new_numel * old_stride_lastvalue;
|
|
new_numel *= new_dims[new_stride_index];
|
|
new_stride_index--;
|
|
}
|
|
if (new_numel != old_numel) {
|
|
return false;
|
|
}
|
|
if (old_dims_index > 0) {
|
|
old_numel = 1;
|
|
new_numel = 1;
|
|
old_stride_lastvalue = old_stride[old_dims_index - 1];
|
|
}
|
|
}
|
|
}
|
|
if (new_stride_index != -1) {
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
} // namespace phi
|