wehub-resource-sync
178 lines
6.4 KiB
C++
178 lines
6.4 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/onednn/onednn_reuse.h"
|
|
|
|
namespace phi {
|
|
|
|
/*
|
|
Pad3D is done by using up to 7 reorders. Following example is done
|
|
on 2D data for simplicity, but it is straightforward to extend it to 3D case.
|
|
|
|
Let us consider following example:
|
|
|
|
N C H W L R T B
|
|
X_dims = (1, 1, 3, 3), paddings = (1, 2, 3, 4) in order Left, Right, Top, Bottom
|
|
|
|
We have to copy the X tensor into Out tensor, but except from that we have to
|
|
fill the rest of the memory with an additional padding. To avoid looping through
|
|
the whole Out memory two times, only these parts of Out memory that won't store
|
|
X's memory are filled with pad value. That behavior is achieved by using
|
|
oneDNN's submemory descriptors which allows us to set offsets for each dimension
|
|
and skip some parts of the memory. For 2D case up to 5 reorders will be used in
|
|
Pad3D kernel(if padding=0 reorder is skipped). In the following example i'th
|
|
number means, that this part of memory was filled by i'th reorder. 4'th reorder
|
|
is copying X memory into Out memory. i&j means that both i'th and j'th reorder
|
|
will set the padding at that location:
|
|
|
|
INDEX
|
|
| 0 1 2 3 4 5
|
|
|_______________________
|
|
0 |0&2 2 2 2 1&2 1&2
|
|
1 |0&2 2 2 2 1&2 1&2
|
|
I 2 |0&2 2 2 2 1&2 1&2
|
|
N 3 | 0 4 4 4 1 1
|
|
D 4 | 0 4 4 4 1 1
|
|
E 5 | 0 4 4 4 1 1
|
|
X 6 |0&3 3 3 3 1&3 1&3
|
|
7 |0&3 3 3 3 1&3 1&3
|
|
8 |0&3 3 3 3 1&3 1&3
|
|
9 |0&3 3 3 3 1&3 1&3
|
|
|
|
Since oneDNN's reorder cannot set the pad value to the memory by itself, we have
|
|
to prefill Out's memory and use it as a temporary buffer, which later is copied
|
|
into the rest of Out's memory. At the end last reorder is done which copies X
|
|
memory into Out memory.
|
|
|
|
*/
|
|
|
|
inline int64_t CalculateNumOfPrefillElems(
|
|
const std::vector<int64_t>& out_tz, const std::vector<int64_t>& paddings) {
|
|
int64_t max_elems = 0;
|
|
int64_t independent_dims = out_tz[0] * out_tz[1];
|
|
|
|
for (size_t i = 0; i < paddings.size() / 2; ++i) {
|
|
int64_t elems = std::max(paddings[2 * i], paddings[2 * i + 1]);
|
|
for (size_t j = 0; j < paddings.size() / 2; ++j) {
|
|
if (j != i) {
|
|
elems *= out_tz[out_tz.size() - 1 - j];
|
|
}
|
|
}
|
|
|
|
if (max_elems < elems) {
|
|
max_elems = elems;
|
|
}
|
|
}
|
|
return independent_dims * max_elems;
|
|
}
|
|
|
|
template <typename T>
|
|
void FillPartOfPadding(const dnnl::engine& onednn_engine,
|
|
T* prefilled_mem_ptr,
|
|
const std::shared_ptr<dnnl::memory>& out_mem_p,
|
|
const std::vector<int64_t>& chunk_tz,
|
|
const std::vector<int64_t>& offsets) {
|
|
auto& astream = OneDNNContext::tls().get_stream();
|
|
|
|
dnnl::memory::desc prefilled_mem_desc(
|
|
chunk_tz,
|
|
funcs::OneDNNGetDataType<T>(),
|
|
funcs::GetPlainOneDNNFormat(chunk_tz.size()));
|
|
dnnl::memory prefilled_mem(
|
|
prefilled_mem_desc, onednn_engine, prefilled_mem_ptr);
|
|
|
|
dnnl::memory::desc out_slice_md =
|
|
out_mem_p->get_desc().submemory_desc(chunk_tz, {offsets});
|
|
dnnl::memory out_slice_mem(
|
|
out_slice_md, onednn_engine, out_mem_p->get_data_handle());
|
|
|
|
auto reorder_p = dnnl::reorder(prefilled_mem, out_slice_mem);
|
|
reorder_p.execute(astream, prefilled_mem, out_slice_mem);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void PadOpKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const std::vector<int64_t>& paddings,
|
|
double pad_value,
|
|
DenseTensor* out) {
|
|
const auto& onednn_engine = dev_ctx.GetEngine();
|
|
auto& astream = OneDNNContext::tls().get_stream();
|
|
|
|
std::vector<int64_t> x_tz = vectorize(x.dims());
|
|
// due to the need of supporting NDHWC, inferring out shape
|
|
// must be done inside the kernel
|
|
std::vector<int64_t> out_tz(x_tz);
|
|
|
|
for (size_t i = 0; i < paddings.size() / 2; ++i) {
|
|
out_tz[out_tz.size() - 1 - i] += paddings[2 * i] + paddings[2 * i + 1];
|
|
}
|
|
out->Resize(out_tz);
|
|
|
|
funcs::ReorderOneDNNHandler reorder_handler(
|
|
x_tz, x.dtype(), funcs::ToOneDNNDataType(x.dtype()), onednn_engine);
|
|
|
|
auto reorder_src_memory_p = reorder_handler.AcquireSrcMemory(
|
|
x.mem_desc(), funcs::to_void_cast(x.data<T>()));
|
|
auto reorder_dst_memory_p = reorder_handler.AcquireDstMemory(
|
|
out,
|
|
out_tz,
|
|
funcs::GetPlainOneDNNFormat(out_tz.size()),
|
|
dev_ctx.GetPlace());
|
|
|
|
// to avoid allocating new temporary memory, Out's memory is used as a tmp
|
|
// buffer for storing a contiguous memory consisting of pad_value, which
|
|
// later is used as a SRC for reorders that are filling Out with padding
|
|
T* out_ptr = out->data<T>();
|
|
std::fill(out_ptr,
|
|
out_ptr + CalculateNumOfPrefillElems(out_tz, paddings),
|
|
pad_value);
|
|
|
|
// paddings are in order: left, right, top, bottom, front, back
|
|
for (size_t i = 0; i < paddings.size(); ++i) {
|
|
if (paddings[i] != 0) {
|
|
std::vector<int64_t> offsets(out_tz.size(), 0);
|
|
std::vector<int64_t> chunk_tz(out_tz.begin(), out_tz.end());
|
|
|
|
chunk_tz[out_tz.size() - 1 - i / 2] = paddings[i];
|
|
if (i % 2 == 1) {
|
|
offsets[out_tz.size() - 1 - i / 2] =
|
|
paddings[i - 1] + x_tz[out_tz.size() - 1 - i / 2];
|
|
}
|
|
|
|
FillPartOfPadding(
|
|
onednn_engine, out_ptr, reorder_dst_memory_p, chunk_tz, offsets);
|
|
}
|
|
}
|
|
astream.wait();
|
|
|
|
std::vector<int64_t> offsets(out_tz.size(), 0);
|
|
for (size_t i = 0; i < paddings.size() / 2; ++i) {
|
|
offsets[out_tz.size() - 1 - i] = paddings[2 * i];
|
|
}
|
|
|
|
auto slice_mem_p =
|
|
reorder_handler.AcquireSubmemory(x_tz, offsets, reorder_dst_memory_p);
|
|
|
|
auto reorder_p =
|
|
reorder_handler.AcquireReorder(slice_mem_p, reorder_src_memory_p);
|
|
reorder_p->execute(astream, *reorder_src_memory_p, *slice_mem_p);
|
|
astream.wait();
|
|
|
|
out->set_mem_desc(reorder_dst_memory_p->get_desc());
|
|
}
|
|
} // namespace phi
|