Files
2026-07-13 12:40:42 +08:00

544 lines
18 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
#include "paddle/phi/backends/cpu/cpu_context.h"
#include "paddle/phi/kernels/funcs/eigen/common.h"
namespace phi {
template <typename T>
void Unnormalize(const CPUContext& dev_ctx,
DenseTensor* grid_slice,
const int max_val, // height-1 or width-1
bool align_corners) {
auto& place = *dev_ctx.eigen_device();
auto grid_slice_t = EigenTensor<T, 3>::From(*grid_slice);
if (align_corners) {
auto factor = static_cast<T>(max_val * 0.5);
grid_slice_t.device(place) = (grid_slice_t + static_cast<T>(1)) * factor;
} else {
auto factor = static_cast<T>((max_val + 1) * 0.5);
grid_slice_t.device(place) =
(grid_slice_t + static_cast<T>(1)) * factor - static_cast<T>(0.5);
}
}
template <typename T>
void Unnormalize3D(const CPUContext& dev_ctx,
DenseTensor* grid_slice,
const int max_val, // height-1 or width-1
bool align_corners) {
auto& place = *dev_ctx.eigen_device();
auto grid_slice_t = EigenTensor<T, 4>::From(*grid_slice);
if (!align_corners) {
auto factor = static_cast<T>((max_val + 1) * 0.5);
grid_slice_t.device(place) =
(grid_slice_t + static_cast<T>(1)) * factor - static_cast<T>(0.5);
} else {
auto factor = static_cast<T>(max_val * 0.5);
grid_slice_t.device(place) = (grid_slice_t + static_cast<T>(1)) * factor;
}
}
template <typename T>
inline bool IsInBound(T x, T y, T x_max, T y_max) {
if (x < 0 || x > x_max || y < 0 || y > y_max) {
return false;
}
return true;
}
template <typename T>
inline bool IsInBound3D(T x, T y, T z, T x_max, T y_max, T z_max) {
if (x < 0 || x > x_max || y < 0 || y > y_max || z < 0 || z > z_max) {
return false;
}
return true;
}
template <typename T>
void GetGridPointValue(const DenseTensor& input,
DenseTensor* output,
const DenseTensor& x,
const DenseTensor& y) {
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t n = input.dims()[0];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t c = input.dims()[1];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t in_h = input.dims()[2];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t in_w = input.dims()[3];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t out_h = x.dims()[1];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t out_w = x.dims()[2];
auto x_t = EigenTensor<T, 3>::From(x);
auto y_t = EigenTensor<T, 3>::From(y);
auto output_t = EigenTensor<T, 4>::From(*output).setConstant((T)0);
auto input_t = EigenTensor<T, 4>::From(input);
for (int i = 0; i < n; i++) {
for (int k = 0; k < out_h; k++) {
for (int l = 0; l < out_w; l++) {
if (IsInBound<int>(static_cast<int>(x_t(i, k, l)),
static_cast<int>(y_t(i, k, l)),
(in_w - 1),
(in_h - 1))) {
for (int j = 0; j < c; j++) {
output_t(i, j, k, l) = input_t(i,
j,
static_cast<int>(y_t(i, k, l)),
static_cast<int>(x_t(i, k, l)));
}
}
}
}
}
}
template <typename T>
void GetGridPointValue_nearest(const DenseTensor& input,
DenseTensor* output,
const DenseTensor& x,
const DenseTensor& y) {
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t n = input.dims()[0];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t c = input.dims()[1];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t in_h = input.dims()[2];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t in_w = input.dims()[3];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t out_h = x.dims()[1];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t out_w = x.dims()[2];
auto x_t = EigenTensor<T, 3>::From(x);
auto y_t = EigenTensor<T, 3>::From(y);
auto output_t = EigenTensor<T, 4>::From(*output).setConstant((T)0);
auto input_t = EigenTensor<T, 4>::From(input);
for (int i = 0; i < n; i++) {
for (int k = 0; k < out_h; k++) {
for (int l = 0; l < out_w; l++) {
if (IsInBound<int>(static_cast<int>(std::nearbyint(x_t(i, k, l))),
static_cast<int>(std::nearbyint(y_t(i, k, l))),
(in_w - 1),
(in_h - 1))) {
for (int j = 0; j < c; j++) {
output_t(i, j, k, l) =
input_t(i,
j,
static_cast<int>(std::nearbyint(y_t(i, k, l))),
static_cast<int>(std::nearbyint(x_t(i, k, l))));
}
}
}
}
}
}
template <typename T>
void AllNeighbors(const CPUContext& dev_ctx,
const DenseTensor& input,
DenseTensor* grid_x,
DenseTensor* grid_y,
DenseTensor* x_w,
DenseTensor* x_e,
DenseTensor* y_n,
DenseTensor* y_s, // positions
DenseTensor* d_w,
DenseTensor* d_e,
DenseTensor* d_n,
DenseTensor* d_s, // distance
DenseTensor* v_wn,
DenseTensor* v_en,
DenseTensor* v_ws,
DenseTensor* v_es) { // values
auto& place = *dev_ctx.eigen_device();
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t c = input.dims()[1];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t n = grid_x->dims()[0];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t out_h = grid_x->dims()[1];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t out_w = grid_x->dims()[2];
// calculate coords of 4 corner points
x_w->Resize({n, out_h, out_w});
x_e->Resize({n, out_h, out_w});
y_n->Resize({n, out_h, out_w});
y_s->Resize({n, out_h, out_w});
dev_ctx.Alloc<T>(x_w);
dev_ctx.Alloc<T>(x_e);
dev_ctx.Alloc<T>(y_n);
dev_ctx.Alloc<T>(y_s);
auto x_w_t = EigenTensor<T, 3>::From(*x_w);
auto x_e_t = EigenTensor<T, 3>::From(*x_e);
auto y_n_t = EigenTensor<T, 3>::From(*y_n);
auto y_s_t = EigenTensor<T, 3>::From(*y_s);
auto grid_x_t = EigenTensor<T, 3>::From(*grid_x);
auto grid_y_t = EigenTensor<T, 3>::From(*grid_y);
x_w_t.device(place) = grid_x_t.floor();
x_e_t.device(place) = x_w_t + static_cast<T>(1);
y_n_t.device(place) = grid_y_t.floor();
y_s_t.device(place) = y_n_t + static_cast<T>(1);
// calculate distances to 4 sides
d_w->Resize({n, out_h, out_w});
d_e->Resize({n, out_h, out_w});
d_n->Resize({n, out_h, out_w});
d_s->Resize({n, out_h, out_w});
dev_ctx.Alloc<T>(d_w);
dev_ctx.Alloc<T>(d_e);
dev_ctx.Alloc<T>(d_n);
dev_ctx.Alloc<T>(d_s);
auto d_w_t = EigenTensor<T, 3>::From(*d_w);
auto d_e_t = EigenTensor<T, 3>::From(*d_e);
auto d_n_t = EigenTensor<T, 3>::From(*d_n);
auto d_s_t = EigenTensor<T, 3>::From(*d_s);
d_w_t.device(place) = grid_x_t - x_w_t;
d_e_t.device(place) = x_e_t - grid_x_t;
d_n_t.device(place) = grid_y_t - y_n_t;
d_s_t.device(place) = y_s_t - grid_y_t;
// calc 4 corner points value
v_wn->Resize({n, c, out_h, out_w});
v_en->Resize({n, c, out_h, out_w});
v_ws->Resize({n, c, out_h, out_w});
v_es->Resize({n, c, out_h, out_w});
dev_ctx.Alloc<T>(v_wn);
dev_ctx.Alloc<T>(v_en);
dev_ctx.Alloc<T>(v_ws);
dev_ctx.Alloc<T>(v_es);
GetGridPointValue<T>(input, v_wn, *x_w, *y_n);
GetGridPointValue<T>(input, v_en, *x_e, *y_n);
GetGridPointValue<T>(input, v_ws, *x_w, *y_s);
GetGridPointValue<T>(input, v_es, *x_e, *y_s);
}
template <typename T>
void Get3DGridPointValue(const DenseTensor& input,
DenseTensor* output,
const DenseTensor& x,
const DenseTensor& y,
const DenseTensor& z) {
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t n = input.dims()[0];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t c = input.dims()[1];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t in_d = input.dims()[2];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t in_h = input.dims()[3];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t in_w = input.dims()[4];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t out_d = x.dims()[1];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t out_h = x.dims()[2];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t out_w = x.dims()[3];
auto x_t = EigenTensor<T, 4>::From(x);
auto y_t = EigenTensor<T, 4>::From(y);
auto z_t = EigenTensor<T, 4>::From(z);
auto output_t =
EigenTensor<T, 5>::From(*output).setConstant(static_cast<T>(0.0));
auto input_t = EigenTensor<T, 5>::From(input);
for (int i = 0; i < n; i++) {
for (int m = 0; m < out_d; m++) {
for (int k = 0; k < out_h; k++) {
for (int l = 0; l < out_w; l++) {
if (IsInBound3D<int>(static_cast<int>(x_t(i, m, k, l)),
static_cast<int>(y_t(i, m, k, l)),
static_cast<int>(z_t(i, m, k, l)),
(in_w - 1),
(in_h - 1),
(in_d - 1))) {
for (int j = 0; j < c; j++) {
output_t(i, j, m, k, l) =
input_t(i,
j,
static_cast<int>(z_t(i, m, k, l)),
static_cast<int>(y_t(i, m, k, l)),
static_cast<int>(x_t(i, m, k, l)));
}
}
}
}
}
}
}
template <typename T>
void Get3DGridPointValue_nearest(const DenseTensor& input,
DenseTensor* output,
const DenseTensor& x,
const DenseTensor& y,
const DenseTensor& z) {
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t n = input.dims()[0];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t c = input.dims()[1];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t in_d = input.dims()[2];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t in_h = input.dims()[3];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t in_w = input.dims()[4];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t out_d = x.dims()[1];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t out_h = x.dims()[2];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t out_w = x.dims()[3];
auto x_t = EigenTensor<T, 4>::From(x);
auto y_t = EigenTensor<T, 4>::From(y);
auto z_t = EigenTensor<T, 4>::From(z);
auto output_t =
EigenTensor<T, 5>::From(*output).setConstant(static_cast<T>(0.0));
auto input_t = EigenTensor<T, 5>::From(input);
for (int i = 0; i < n; i++) {
for (int m = 0; m < out_d; m++) {
for (int k = 0; k < out_h; k++) {
for (int l = 0; l < out_w; l++) {
if (IsInBound3D<int>(
static_cast<int>(std::nearbyint(x_t(i, m, k, l))),
static_cast<int>(std::nearbyint(y_t(i, m, k, l))),
static_cast<int>(std::nearbyint(z_t(i, m, k, l))),
(in_w - 1),
(in_h - 1),
(in_d - 1))) {
for (int j = 0; j < c; j++) {
output_t(i, j, m, k, l) =
input_t(i,
j,
static_cast<int>(std::nearbyint(z_t(i, m, k, l))),
static_cast<int>(std::nearbyint(y_t(i, m, k, l))),
static_cast<int>(std::nearbyint(x_t(i, m, k, l))));
}
}
}
}
}
}
}
template <typename T>
void All3DNeighbors(const CPUContext& dev_ctx,
const DenseTensor& input,
DenseTensor* grid_x,
DenseTensor* grid_y,
DenseTensor* grid_z,
DenseTensor* x_w,
DenseTensor* x_e,
DenseTensor* y_n,
DenseTensor* y_s,
DenseTensor* z_t,
DenseTensor* z_b, // positions
DenseTensor* d_w,
DenseTensor* d_e,
DenseTensor* d_n,
DenseTensor* d_s,
DenseTensor* d_t,
DenseTensor* d_b, // distance
DenseTensor* v_twn,
DenseTensor* v_ten,
DenseTensor* v_tws,
DenseTensor* v_tes,
DenseTensor* v_bwn,
DenseTensor* v_ben,
DenseTensor* v_bws,
DenseTensor* v_bes) { // values
auto& place = *dev_ctx.eigen_device();
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t c = input.dims()[1];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t n = grid_x->dims()[0];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t out_d = grid_x->dims()[1];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t out_h = grid_x->dims()[2];
// TODO(large-tensor): downstream functors may still use int; guard until
// upgraded.
int64_t out_w = grid_x->dims()[3];
// calculate coords of 6 corner points
x_w->Resize({n, out_d, out_h, out_w});
x_e->Resize({n, out_d, out_h, out_w});
y_n->Resize({n, out_d, out_h, out_w});
y_s->Resize({n, out_d, out_h, out_w});
z_t->Resize({n, out_d, out_h, out_w});
z_b->Resize({n, out_d, out_h, out_w});
dev_ctx.Alloc<T>(x_w);
dev_ctx.Alloc<T>(x_e);
dev_ctx.Alloc<T>(y_n);
dev_ctx.Alloc<T>(y_s);
dev_ctx.Alloc<T>(z_t);
dev_ctx.Alloc<T>(z_b);
auto x_w_t = EigenTensor<T, 4>::From(*x_w);
auto x_e_t = EigenTensor<T, 4>::From(*x_e);
auto y_n_t = EigenTensor<T, 4>::From(*y_n);
auto y_s_t = EigenTensor<T, 4>::From(*y_s);
auto z_t_t = EigenTensor<T, 4>::From(*z_t);
auto z_b_t = EigenTensor<T, 4>::From(*z_b);
auto grid_x_t = EigenTensor<T, 4>::From(*grid_x);
auto grid_y_t = EigenTensor<T, 4>::From(*grid_y);
auto grid_z_t = EigenTensor<T, 4>::From(*grid_z);
x_w_t.device(place) = grid_x_t.floor();
x_e_t.device(place) = x_w_t + static_cast<T>(1);
y_n_t.device(place) = grid_y_t.floor();
y_s_t.device(place) = y_n_t + static_cast<T>(1);
z_t_t.device(place) = grid_z_t.floor();
z_b_t.device(place) = z_t_t + static_cast<T>(1);
// calculate distances to 6 sides
d_w->Resize({n, out_d, out_h, out_w});
d_e->Resize({n, out_d, out_h, out_w});
d_n->Resize({n, out_d, out_h, out_w});
d_s->Resize({n, out_d, out_h, out_w});
d_t->Resize({n, out_d, out_h, out_w});
d_b->Resize({n, out_d, out_h, out_w});
dev_ctx.Alloc<T>(d_w);
dev_ctx.Alloc<T>(d_e);
dev_ctx.Alloc<T>(d_n);
dev_ctx.Alloc<T>(d_s);
dev_ctx.Alloc<T>(d_t);
dev_ctx.Alloc<T>(d_b);
auto d_w_t = EigenTensor<T, 4>::From(*d_w);
auto d_e_t = EigenTensor<T, 4>::From(*d_e);
auto d_n_t = EigenTensor<T, 4>::From(*d_n);
auto d_s_t = EigenTensor<T, 4>::From(*d_s);
auto d_t_t = EigenTensor<T, 4>::From(*d_t);
auto d_b_t = EigenTensor<T, 4>::From(*d_b);
d_w_t.device(place) = grid_x_t - x_w_t;
d_e_t.device(place) = x_e_t - grid_x_t;
d_n_t.device(place) = grid_y_t - y_n_t;
d_s_t.device(place) = y_s_t - grid_y_t;
d_t_t.device(place) = grid_z_t - z_t_t;
d_b_t.device(place) = z_b_t - grid_z_t;
// calc 8 corner points value
v_twn->Resize({n, c, out_d, out_h, out_w});
v_ten->Resize({n, c, out_d, out_h, out_w});
v_tws->Resize({n, c, out_d, out_h, out_w});
v_tes->Resize({n, c, out_d, out_h, out_w});
v_bwn->Resize({n, c, out_d, out_h, out_w});
v_ben->Resize({n, c, out_d, out_h, out_w});
v_bws->Resize({n, c, out_d, out_h, out_w});
v_bes->Resize({n, c, out_d, out_h, out_w});
dev_ctx.Alloc<T>(v_twn);
dev_ctx.Alloc<T>(v_ten);
dev_ctx.Alloc<T>(v_tws);
dev_ctx.Alloc<T>(v_tes);
dev_ctx.Alloc<T>(v_bwn);
dev_ctx.Alloc<T>(v_ben);
dev_ctx.Alloc<T>(v_bws);
dev_ctx.Alloc<T>(v_bes);
Get3DGridPointValue<T>(input, v_twn, *x_w, *y_n, *z_t);
Get3DGridPointValue<T>(input, v_ten, *x_e, *y_n, *z_t);
Get3DGridPointValue<T>(input, v_tws, *x_w, *y_s, *z_t);
Get3DGridPointValue<T>(input, v_tes, *x_e, *y_s, *z_t);
Get3DGridPointValue<T>(input, v_bwn, *x_w, *y_n, *z_b);
Get3DGridPointValue<T>(input, v_ben, *x_e, *y_n, *z_b);
Get3DGridPointValue<T>(input, v_bws, *x_w, *y_s, *z_b);
Get3DGridPointValue<T>(input, v_bes, *x_e, *y_s, *z_b);
}
} // namespace phi