wehub-resource-sync
717 lines
21 KiB
C++
717 lines
21 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.
|
|
|
|
#pragma once
|
|
#include <vector>
|
|
|
|
#include "paddle/common/ddim.h"
|
|
#include "paddle/fluid/framework/details/op_registry.h"
|
|
#include "paddle/fluid/framework/operator.h"
|
|
#include "paddle/fluid/primitive/base/lazy_tensor.h"
|
|
#include "paddle/fluid/primitive/primitive/primitive.h"
|
|
#include "paddle/phi/api/include/tensor.h"
|
|
#include "paddle/phi/kernels/funcs/common_infer_shape_functions.h"
|
|
|
|
namespace paddle {
|
|
class Tensor;
|
|
namespace primitive {
|
|
template <typename T>
|
|
static Tensor get_slice(const Tensor& x, int64_t idx) {
|
|
return slice<T>(x, {0}, {idx}, {idx + 1}, {1}, {});
|
|
}
|
|
|
|
template <typename T>
|
|
static Tensor get_slice_vec(const Tensor& x,
|
|
int64_t start_idx,
|
|
int64_t end_idx) {
|
|
return slice<T>(x, {0}, {start_idx}, {end_idx}, {1}, {});
|
|
}
|
|
|
|
template <typename T>
|
|
void set_output(const Tensor& x_tmp, Tensor* x);
|
|
|
|
template <typename T>
|
|
void by_pass(const Tensor& x_tmp, Tensor* x);
|
|
|
|
// This function determine whether dtype is in [float16, bfloat16, uint16]
|
|
static bool is_half_dtype(const DataType& dtype) {
|
|
if (dtype == DataType::FLOAT16 || dtype == DataType::BFLOAT16 ||
|
|
dtype == DataType::UINT16) {
|
|
return true;
|
|
} else {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
// This function expands the dimension of origin Tensor based on the value of
|
|
// axis
|
|
static std::vector<int64_t> get_expand_dims(const Tensor& origin,
|
|
const std::vector<int64_t>& axis) {
|
|
std::vector<int64_t> result(origin.shape());
|
|
for (size_t i = 0; i < axis.size(); ++i) {
|
|
int64_t offset = axis[i];
|
|
if (offset < 0) {
|
|
offset += result.size() + 1;
|
|
}
|
|
|
|
PADDLE_ENFORCE_LE(
|
|
offset,
|
|
result.size(),
|
|
common::errors::OutOfRange("Your index [%lu] exceeds the number of "
|
|
"elements in origin_dims[%lu].",
|
|
offset,
|
|
result.size()));
|
|
result.insert(result.begin() + offset, 1);
|
|
}
|
|
return result;
|
|
}
|
|
|
|
// This function compute unsqueeze dims for reshape to replace unsqueeze.
|
|
static std::vector<int64_t> get_unsqueeze_dims(
|
|
const Tensor& origin, const std::vector<int64_t>& axis) {
|
|
auto sort_axis = axis;
|
|
std::sort(sort_axis.begin(), sort_axis.end());
|
|
auto origin_dims = origin.shape();
|
|
auto total_shape_size = origin_dims.size() + sort_axis.size();
|
|
std::vector<int64_t> result;
|
|
size_t j = 0, k = 0;
|
|
for (size_t i = 0; i < total_shape_size; ++i) {
|
|
if (j < sort_axis.size() && sort_axis[j] == int64_t(i)) {
|
|
result.push_back(1);
|
|
j++;
|
|
} else {
|
|
PADDLE_ENFORCE_LT(
|
|
k,
|
|
origin_dims.size(),
|
|
common::errors::OutOfRange("Your index [%lu] exceeds the number of "
|
|
"elements in origin_dims[%lu].",
|
|
k,
|
|
origin_dims.size()));
|
|
result.push_back(origin_dims[k]);
|
|
k++;
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
|
|
// This function compute `dynamic` unsqueeze dims for reshape to replace
|
|
// unsqueeze. And should used only on `dynamic`.
|
|
template <typename T>
|
|
Tensor get_unsqueeze_dims(const Tensor& origin_shape,
|
|
const std::vector<int64_t>& axis) {
|
|
auto total_shape_size = origin_shape.numel() + axis.size();
|
|
const Tensor one = full<T>({1}, 1, origin_shape.dtype());
|
|
|
|
std::vector<Tensor> result(total_shape_size, one);
|
|
// to support axis not in increasing order.
|
|
std::vector<bool> is_set(total_shape_size, false);
|
|
|
|
for (size_t i = 0; i < axis.size(); ++i) {
|
|
PADDLE_ENFORCE_LT(
|
|
axis[i],
|
|
total_shape_size,
|
|
common::errors::OutOfRange("Your index [%lu] exceeds the number of "
|
|
"elements in origin_dims[%lu].",
|
|
axis[i],
|
|
total_shape_size));
|
|
is_set[axis[i]] = true;
|
|
}
|
|
|
|
size_t j = 0;
|
|
for (size_t i = 0; i < total_shape_size; ++i) {
|
|
if (is_set[i]) {
|
|
continue;
|
|
}
|
|
result[i] = get_slice<T>(origin_shape, int64_t(j));
|
|
is_set[i] = true;
|
|
++j;
|
|
}
|
|
return concat<T>(result);
|
|
}
|
|
|
|
// This function compute unsqueeze dims for reshape to replace unsqueeze.
|
|
static std::vector<int64_t> get_squeeze_dims(const Tensor& origin,
|
|
const std::vector<int64_t>& axis) {
|
|
auto origin_dims = origin.shape();
|
|
auto total_shape_size = origin_dims.size();
|
|
std::vector<int64_t> result;
|
|
for (size_t i = 0; i < total_shape_size; ++i) {
|
|
if (origin_dims[i] != 1) {
|
|
result.push_back(origin_dims[i]);
|
|
} else if (origin_dims[i] == 1 &&
|
|
std::find(axis.begin(), axis.end(), int64_t(i)) == axis.end()) {
|
|
result.push_back(1);
|
|
} else {
|
|
continue;
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
|
|
static std::vector<int64_t> process_dims(const Tensor& origin,
|
|
const std::vector<int64_t>& axis) {
|
|
auto origin_dims = origin.shape();
|
|
auto total_shape_size = origin_dims.size();
|
|
std::vector<int64_t> result;
|
|
auto axis_size = axis.size();
|
|
if (axis_size == 0) {
|
|
for (size_t i = 0; i < total_shape_size; ++i) {
|
|
result.push_back(i);
|
|
}
|
|
} else {
|
|
for (size_t i = 0; i < axis_size; ++i) {
|
|
if (axis[i] < 0) {
|
|
result.push_back(axis[i] + total_shape_size);
|
|
} else {
|
|
result.push_back(axis[i]);
|
|
}
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
|
|
// These method don't need to be specified
|
|
// These method only handle the static shape case
|
|
static phi::DDim get_reduce_dims_from_out(const phi::DDim& dout_dims,
|
|
const phi::DDim& in_dims) {
|
|
bool has_dynamic_shape = false;
|
|
for (int i = 0; i < dout_dims.size(); i++) {
|
|
if (dout_dims[i] == -1) {
|
|
has_dynamic_shape = true;
|
|
break;
|
|
}
|
|
}
|
|
PADDLE_ENFORCE_EQ(
|
|
has_dynamic_shape,
|
|
false,
|
|
common::errors::InvalidArgument(
|
|
"Function get_reduce_dims_from_out() only use in static shape case, "
|
|
"but the input [dout_dims] have the dynamic shape."));
|
|
|
|
for (int i = 0; i < in_dims.size(); i++) {
|
|
if (in_dims[i] == -1) {
|
|
has_dynamic_shape = true;
|
|
break;
|
|
}
|
|
}
|
|
PADDLE_ENFORCE_EQ(
|
|
has_dynamic_shape,
|
|
false,
|
|
common::errors::InvalidArgument(
|
|
"Function get_reduce_dims_from_out() only use in static shape case, "
|
|
"but the input [in_dims] have the dynamic shape."));
|
|
|
|
int bat = dout_dims.size() - in_dims.size();
|
|
std::vector<int64_t> result;
|
|
for (int i = 0; i < bat; ++i) {
|
|
result.push_back(i);
|
|
}
|
|
for (int i = 0; i < in_dims.size(); ++i) {
|
|
if (in_dims[i] == 1 && dout_dims[i + bat] != 1) {
|
|
result.push_back(i + bat);
|
|
} else {
|
|
PADDLE_ENFORCE_EQ(
|
|
in_dims[i],
|
|
dout_dims[i + bat],
|
|
common::errors::InvalidArgument(
|
|
"ReduceDims dimension mismatch. Operands could "
|
|
"not be broadcast together with the shape of dout = [%s] and "
|
|
"the shape of in_dims = [%s]. Received [%d] in X is not equal to "
|
|
"[%d] in Y at i:%d.",
|
|
dout_dims,
|
|
in_dims,
|
|
dout_dims[i + bat],
|
|
in_dims[i],
|
|
i));
|
|
}
|
|
}
|
|
return common::make_ddim(result);
|
|
}
|
|
|
|
static phi::DDim get_reduce_dims(const phi::DDim& x_dims,
|
|
const phi::DDim& y_dims) {
|
|
auto out_dims = phi::funcs::BroadcastTwoDims(x_dims, y_dims);
|
|
return get_reduce_dims_from_out(out_dims, x_dims);
|
|
}
|
|
|
|
void SetEmptyGrad(const std::vector<std::vector<Tensor>>& outputs,
|
|
const std::vector<std::vector<bool>>& stop_gradients);
|
|
|
|
std::vector<std::vector<Tensor>> ConstructVjpResultByStopGradients(
|
|
const std::vector<std::vector<Tensor>>& outputs,
|
|
const std::vector<std::vector<bool>>& stop_gradients);
|
|
|
|
static bool find_value(const std::vector<int64_t>& vec, int64_t value) {
|
|
if (std::find(vec.begin(), vec.end(), value) != vec.end()) {
|
|
return true;
|
|
} else {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
static bool has_dynamic_shape(const std::vector<int64_t>& shape) {
|
|
return std::find(shape.begin(), shape.end(), -1) != shape.end();
|
|
}
|
|
|
|
static bool has_dynamic_shape(const std::vector<int64_t>& shape,
|
|
const std::vector<int64_t>& axis) {
|
|
bool flag = false;
|
|
const int64_t rank = shape.size();
|
|
for (int64_t idx : axis) {
|
|
if (idx < 0) idx += rank;
|
|
PADDLE_ENFORCE_LT(
|
|
idx,
|
|
rank,
|
|
::common::errors::PreconditionNotMet(
|
|
"Required idx < shape.size(), but received %d.", idx));
|
|
if (shape[idx] == -1) {
|
|
flag = true;
|
|
break;
|
|
}
|
|
}
|
|
return flag;
|
|
}
|
|
|
|
template <typename T>
|
|
Tensor ConvertToMT(const Tensor& x) {
|
|
bool need_cast = x.dtype() == DataType::FLOAT16 ||
|
|
x.dtype() == DataType::BFLOAT16 ||
|
|
x.dtype() == DataType::UINT16;
|
|
if (need_cast) {
|
|
return cast<T>(x, DataType::FLOAT32);
|
|
}
|
|
return x;
|
|
}
|
|
|
|
template <typename T>
|
|
Tensor ConvertToOrig(const Tensor& out, DataType input_dtype) {
|
|
bool need_cast = out.dtype() != input_dtype;
|
|
if (need_cast) {
|
|
return cast<T>(out, input_dtype);
|
|
}
|
|
return out;
|
|
}
|
|
|
|
class LayerNormDecompHelper {
|
|
public:
|
|
LayerNormDecompHelper(const Tensor& x,
|
|
const paddle::optional<Tensor>& scale,
|
|
const paddle::optional<Tensor>& bias,
|
|
int begin_norm_axis) {
|
|
auto x_dims = x.dims();
|
|
x_rank_ = x_dims.size();
|
|
begin_norm_axis_ = begin_norm_axis;
|
|
if (begin_norm_axis_ < 0) {
|
|
begin_norm_axis_ += x_rank_;
|
|
}
|
|
|
|
scale_need_reshape_ = (begin_norm_axis + 1 != x_rank_);
|
|
static_norm_shape_ = true;
|
|
|
|
for (int i = begin_norm_axis; i < x_rank_; ++i) {
|
|
if (x_dims[i] < 0) {
|
|
static_norm_shape_ = false;
|
|
normalized_numel_ = -1;
|
|
break;
|
|
}
|
|
|
|
normalized_shape_.push_back(x_dims[i]);
|
|
|
|
normalized_numel_ *= x_dims[i];
|
|
}
|
|
|
|
if (!static_norm_shape_) {
|
|
// try get static norm numel from scale for bias
|
|
normalized_numel_ = -1;
|
|
if (scale.get_ptr()) {
|
|
normalized_numel_ = scale->dims()[0];
|
|
} else if (bias.get_ptr()) {
|
|
normalized_numel_ = bias->dims()[0];
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
Tensor Process(const Tensor& s, const Tensor& x) {
|
|
if (!scale_need_reshape_) {
|
|
return s;
|
|
}
|
|
|
|
if (static_norm_shape_) {
|
|
return reshape<T>(s, normalized_shape_);
|
|
} else {
|
|
return backend::reshape<T>(
|
|
s, get_slice_vec<T>(shape64<T>(x), begin_norm_axis_, x_rank_));
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
Tensor GetNormalizedNumel(const Tensor& x) {
|
|
if (normalized_numel_ != -1) {
|
|
return full_scalar<T>(normalized_numel_, x.dtype());
|
|
} else {
|
|
auto x_shape = shape64<T>(x);
|
|
auto numel = get_slice<T>(x_shape, begin_norm_axis_);
|
|
for (int64_t i = begin_norm_axis_ + 1; i < x_rank_; ++i) {
|
|
numel = numel * get_slice<T>(x_shape, i);
|
|
}
|
|
|
|
return cast<T>(numel, x.type());
|
|
}
|
|
}
|
|
|
|
private:
|
|
std::vector<int64_t> normalized_shape_;
|
|
bool scale_need_reshape_;
|
|
bool static_norm_shape_;
|
|
int64_t x_rank_;
|
|
int64_t normalized_numel_{1};
|
|
int begin_norm_axis_;
|
|
};
|
|
|
|
template <typename T>
|
|
class BatchNormDecompHelper {
|
|
public:
|
|
BatchNormDecompHelper(const Tensor& x,
|
|
const paddle::optional<Tensor>& scale,
|
|
const paddle::optional<Tensor>& bias,
|
|
const std::string& data_format) {
|
|
auto x_dims = phi::vectorize(x.dims());
|
|
x_rank_ = x_dims.size();
|
|
|
|
if (data_format == "NCHW") {
|
|
channel_axis_ = 1;
|
|
reduce_axis_.push_back(0);
|
|
for (int64_t i = channel_axis_ + 1; i < x_rank_; ++i) {
|
|
reduce_axis_.push_back(i);
|
|
}
|
|
} else if (data_format == "NHWC") {
|
|
channel_axis_ = x_rank_ - 1;
|
|
for (int64_t i = 0; i < channel_axis_; ++i) {
|
|
reduce_axis_.push_back(i);
|
|
}
|
|
} else {
|
|
PADDLE_THROW(
|
|
common::errors::Unimplemented("Only support NCHW and NHWC format."));
|
|
}
|
|
|
|
scale_bias_new_shape_.push_back(0);
|
|
for (int64_t i = channel_axis_ + 1; i < x_rank_; ++i) {
|
|
scale_bias_new_shape_.push_back(1);
|
|
}
|
|
|
|
// int64_t channel_dim = x_dims[channel_axis_];
|
|
// if ((channel_dim < 0) && scale) {
|
|
// channel_dim = scale->dims()[0];
|
|
// }
|
|
// if ((channel_dim < 0) && bias) {
|
|
// channel_dim = bias->dims()[0];
|
|
// }
|
|
}
|
|
|
|
const std::vector<int64_t>& GetReduceAxis() const { return reduce_axis_; }
|
|
|
|
const std::vector<int64_t>& GetScaleBiasNewShape() const {
|
|
return scale_bias_new_shape_;
|
|
}
|
|
|
|
Tensor GetNHW(const Tensor& x) {
|
|
auto x_dims = x.dims();
|
|
|
|
bool static_nhw = true;
|
|
int64_t nhw_numel = 1;
|
|
for (int64_t i = 0; i < x_rank_; ++i) {
|
|
if (i == channel_axis_) {
|
|
continue;
|
|
}
|
|
if (x_dims[i] < 0) {
|
|
static_nhw = false;
|
|
break;
|
|
}
|
|
nhw_numel *= x_dims[i];
|
|
}
|
|
|
|
if (static_nhw) {
|
|
return full_scalar<T>(nhw_numel, x.dtype());
|
|
} else {
|
|
auto x_shape = shape64<T>(x);
|
|
auto nhw = get_slice<T>(x_shape, 0);
|
|
for (int64_t i = 1; i < x_rank_; ++i) {
|
|
if (i == channel_axis_) {
|
|
continue;
|
|
}
|
|
|
|
nhw = nhw * get_slice<T>(x_shape, i);
|
|
}
|
|
|
|
return cast<T>(nhw, x.dtype());
|
|
}
|
|
}
|
|
|
|
private:
|
|
std::vector<int64_t> reduce_axis_;
|
|
std::vector<int64_t> scale_bias_new_shape_;
|
|
int64_t channel_axis_;
|
|
int64_t x_rank_;
|
|
};
|
|
|
|
template <typename T>
|
|
class InstanceNormDecompHelper {
|
|
public:
|
|
explicit InstanceNormDecompHelper(const Tensor& x) {
|
|
x_dims_ = phi::vectorize(x.dims());
|
|
x_rank_ = x_dims_.size();
|
|
|
|
for (int64_t i = 2; i < x_rank_; ++i) {
|
|
reduce_axis_.push_back(i);
|
|
n_plus_reduce_axis_.push_back(i);
|
|
}
|
|
n_plus_reduce_axis_.push_back(0);
|
|
}
|
|
|
|
Tensor GetHW(const Tensor& x) {
|
|
auto dims = phi::vectorize(x.dims());
|
|
int64_t rank = dims.size();
|
|
if (has_dynamic_shape(x.shape())) {
|
|
Tensor x_shape = shape64<T>(x);
|
|
auto hw = full_scalar<T>(1.0, x.dtype());
|
|
for (int64_t i = 2; i < rank; ++i) {
|
|
hw = hw * get_slice<T>(x_shape, i);
|
|
}
|
|
return cast<T>(hw, x.dtype());
|
|
} else {
|
|
int64_t hw = 1;
|
|
for (int64_t i = 2; i < rank; ++i) {
|
|
hw *= dims[i];
|
|
}
|
|
return full_scalar<T>(hw, x.dtype());
|
|
}
|
|
}
|
|
|
|
const std::vector<int64_t> GetReduceAxis() const { return reduce_axis_; }
|
|
const std::vector<int64_t> GetNPlusReduceAxis() const {
|
|
return n_plus_reduce_axis_;
|
|
}
|
|
const std::vector<int64_t>& GetDims() const { return x_dims_; }
|
|
|
|
private:
|
|
std::vector<int64_t> reduce_axis_;
|
|
std::vector<int64_t> n_plus_reduce_axis_;
|
|
std::vector<int64_t> x_dims_;
|
|
int64_t x_rank_;
|
|
};
|
|
|
|
template <typename T>
|
|
class GroupNormDecompHelper {
|
|
public:
|
|
GroupNormDecompHelper(const Tensor& x,
|
|
const paddle::optional<Tensor>& scale,
|
|
const paddle::optional<Tensor>& bias,
|
|
int64_t group_num,
|
|
const std::string& data_format) {
|
|
auto x_dims = phi::vectorize(x.dims());
|
|
x_rank_ = x_dims.size();
|
|
|
|
if (data_format == "NCHW") {
|
|
channel_axis_ = 1;
|
|
for (int64_t i = channel_axis_ + 1; i < x_rank_ + 1; ++i) {
|
|
reduce_axis_.push_back(i);
|
|
}
|
|
} else if (data_format == "NHWC") {
|
|
channel_axis_ = x_rank_ - 1;
|
|
for (int64_t i = 1; i < channel_axis_; ++i) {
|
|
reduce_axis_.push_back(i);
|
|
}
|
|
reduce_axis_.push_back(x_rank_);
|
|
} else {
|
|
PADDLE_THROW(
|
|
common::errors::Unimplemented("Only support NCHW and NHWC format."));
|
|
}
|
|
|
|
scale_bias_new_shape_.push_back(group_num);
|
|
scale_bias_new_shape_.push_back(-1);
|
|
for (int64_t i = channel_axis_ + 1; i < x_rank_; ++i) {
|
|
scale_bias_new_shape_.push_back(1);
|
|
}
|
|
|
|
int64_t channel_dim = x_dims[channel_axis_];
|
|
if ((channel_dim < 0) && scale) {
|
|
channel_dim = scale->dims()[0];
|
|
}
|
|
if ((channel_dim < 0) && bias) {
|
|
channel_dim = bias->dims()[0];
|
|
}
|
|
|
|
int unk_count = 0;
|
|
for (int64_t i = 0; i < x_rank_; ++i) {
|
|
if ((i != channel_axis_) && (x_dims[i] < 0)) {
|
|
unk_count++;
|
|
}
|
|
}
|
|
|
|
if (channel_dim > 0) {
|
|
// Can use vector<int64> as output shape
|
|
|
|
// case 1: axis is the last one
|
|
// case 2: from axis + 1 to end all positive
|
|
can_use_vector_int_as_output_shape_ =
|
|
(channel_axis_ + 1 == x_rank_) ||
|
|
std::find(x_dims.begin() + channel_axis_ + 1, x_dims.end(), -1) ==
|
|
x_dims.end();
|
|
|
|
// case 3: one ONE unk dim(-1) except axis
|
|
|
|
can_use_vector_int_as_output_shape_ =
|
|
can_use_vector_int_as_output_shape_ || (unk_count <= 1);
|
|
} else {
|
|
can_use_vector_int_as_output_shape_ = (unk_count == 0);
|
|
}
|
|
|
|
std::vector<int64_t> split_dim;
|
|
split_dim.push_back(group_num);
|
|
split_dim.push_back(channel_dim < 0 ? -1 : channel_dim / group_num);
|
|
|
|
if (can_use_vector_int_as_output_shape_) {
|
|
split_out_shape_.reserve(x_rank_ + 1);
|
|
for (int64_t i = 0; i < channel_axis_; ++i) {
|
|
split_out_shape_.push_back(0);
|
|
merge_out_shape_.push_back(0);
|
|
}
|
|
|
|
split_out_shape_.insert(
|
|
split_out_shape_.end(), split_dim.begin(), split_dim.end());
|
|
merge_out_shape_.push_back(channel_dim);
|
|
|
|
for (int64_t i = channel_axis_ + 1; i < x_rank_; ++i) {
|
|
split_out_shape_.push_back(x_dims[i]);
|
|
merge_out_shape_.push_back(x_dims[i]);
|
|
}
|
|
} else {
|
|
auto x_shape = shape64<T>(x);
|
|
if (channel_axis_ > 0) {
|
|
split_shape_tensor_.push_back(
|
|
get_slice_vec<T>(x_shape, 0, channel_axis_));
|
|
merge_shape_tensor_.push_back(
|
|
get_slice_vec<T>(x_shape, 0, channel_axis_));
|
|
}
|
|
|
|
split_shape_tensor_.push_back(
|
|
full<T>({1}, split_dim[0], DataType::INT64));
|
|
split_shape_tensor_.push_back(
|
|
full<T>({1}, split_dim[1], DataType::INT64));
|
|
|
|
merge_shape_tensor_.push_back(full<T>({1}, channel_dim, DataType::INT64));
|
|
|
|
if (channel_axis_ + 1 < x_rank_) {
|
|
split_shape_tensor_.push_back(
|
|
get_slice_vec<T>(x_shape, channel_axis_ + 1, x_rank_));
|
|
merge_shape_tensor_.push_back(
|
|
get_slice_vec<T>(x_shape, channel_axis_ + 1, x_rank_));
|
|
}
|
|
}
|
|
}
|
|
|
|
Tensor Split(const Tensor& s) {
|
|
if (can_use_vector_int_as_output_shape_) {
|
|
return reshape<T>(s, split_out_shape_);
|
|
} else {
|
|
return backend::reshape<T>(s, concat<T>(split_shape_tensor_, 0));
|
|
}
|
|
}
|
|
|
|
Tensor Merge(const Tensor& x) {
|
|
if (can_use_vector_int_as_output_shape_) {
|
|
return reshape<T>(x, merge_out_shape_);
|
|
} else {
|
|
return backend::reshape<T>(x, concat<T>(merge_shape_tensor_, 0));
|
|
}
|
|
}
|
|
|
|
const std::vector<int64_t>& GetReduceAxis() const { return reduce_axis_; }
|
|
|
|
std::vector<int64_t> GetMeanVarSqueezeAxis() const {
|
|
std::vector<int64_t> output;
|
|
|
|
for (int64_t i = 1; i < channel_axis_; ++i) {
|
|
output.push_back(1);
|
|
}
|
|
for (int64_t i = channel_axis_ + 1; i <= x_rank_; ++i) {
|
|
output.push_back(-1);
|
|
}
|
|
return output;
|
|
}
|
|
|
|
const std::vector<int64_t>& GetScaleBiasNewShape() const {
|
|
return scale_bias_new_shape_;
|
|
}
|
|
|
|
Tensor GetHW(const Tensor& x) {
|
|
auto x_dims = x.dims();
|
|
// process reduce axis
|
|
|
|
bool static_hw = true;
|
|
int64_t hwg_numel = 1;
|
|
|
|
for (size_t i = 0; i < reduce_axis_.size(); ++i) {
|
|
if (x_dims[reduce_axis_[i]] < 0) {
|
|
static_hw = false;
|
|
break;
|
|
}
|
|
hwg_numel *= x_dims[reduce_axis_[i]];
|
|
}
|
|
|
|
if (static_hw) {
|
|
return full_scalar<T>(hwg_numel, x.dtype());
|
|
} else {
|
|
auto x_shape = shape64<T>(x);
|
|
auto numel = get_slice<T>(x_shape, reduce_axis_.front());
|
|
for (size_t i = 1; i < reduce_axis_.size(); ++i) {
|
|
numel = numel * get_slice<T>(x_shape, reduce_axis_[i]);
|
|
}
|
|
|
|
return cast<T>(numel, x.dtype());
|
|
}
|
|
}
|
|
|
|
std::vector<int64_t> GetReduceAxisExceptChannel() const {
|
|
std::vector<int64_t> reduce_axis;
|
|
reduce_axis.reserve(x_rank_ - 1);
|
|
|
|
for (int64_t i = 0; i < x_rank_ + 1; ++i) {
|
|
if (i != channel_axis_ && i != channel_axis_ + 1) {
|
|
reduce_axis.push_back(i);
|
|
}
|
|
}
|
|
|
|
return reduce_axis;
|
|
}
|
|
|
|
private:
|
|
bool can_use_vector_int_as_output_shape_{false};
|
|
std::vector<int64_t> split_out_shape_;
|
|
std::vector<Tensor> split_shape_tensor_;
|
|
|
|
std::vector<int64_t> merge_out_shape_;
|
|
std::vector<Tensor> merge_shape_tensor_;
|
|
std::vector<int64_t> reduce_axis_;
|
|
std::vector<int64_t> scale_bias_new_shape_;
|
|
|
|
int64_t group_num_;
|
|
int64_t channel_axis_;
|
|
int64_t x_rank_;
|
|
};
|
|
|
|
} // namespace primitive
|
|
} // namespace paddle
|