wehub-resource-sync
216 lines
9.1 KiB
C++
216 lines
9.1 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/backends/onednn/onednn_reuse.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
|
|
namespace phi::fusion {
|
|
|
|
template <typename T, dnnl::algorithm BINARY_OP>
|
|
void FusedElementwiseKernel(const OneDNNContext& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& y,
|
|
const int axis,
|
|
const std::string& fuse_activation,
|
|
const float fuse_alpha,
|
|
const float fuse_beta,
|
|
const float fused_output_scale,
|
|
const std::vector<int>& fused_unsqueeze2_axes,
|
|
const float scale_x,
|
|
const float scale_y,
|
|
const float scale_out,
|
|
DenseTensor* out) {
|
|
const auto& onednn_engine = dev_ctx.GetEngine();
|
|
|
|
dnnl::post_ops post_operations;
|
|
funcs::AppendActivation(
|
|
dev_ctx, post_operations, fuse_activation, fuse_alpha, fuse_beta);
|
|
if (fused_output_scale != 1.0) {
|
|
// linear post op's formula is `alpha * dst + beta`. Here we only want to
|
|
// scale the output not shift it, so the beta is set to 0.0f.
|
|
post_operations.append_eltwise(
|
|
dnnl::algorithm::eltwise_linear, fused_output_scale, 0.0f);
|
|
}
|
|
|
|
auto* non_const_x = &x;
|
|
auto* non_const_y = &y;
|
|
|
|
funcs::BinaryOneDNNHandler<T> handler(BINARY_OP,
|
|
axis,
|
|
onednn_engine,
|
|
dev_ctx.GetPlace(),
|
|
non_const_x,
|
|
non_const_y,
|
|
out,
|
|
scale_x,
|
|
scale_y,
|
|
scale_out,
|
|
true,
|
|
post_operations);
|
|
|
|
// oneDNN's binary is optimized for broadcasting y into x, so in other case
|
|
// we have to swap tensors to achieve optimal performance
|
|
if (x.numel() < y.numel()) {
|
|
std::swap(non_const_x, non_const_y);
|
|
}
|
|
|
|
const auto src_x_memory =
|
|
handler.swin_case ? (x.numel() == y.numel()
|
|
? handler.AcquireExtendSrcMemory(non_const_x, 0)
|
|
: handler.AcquireSrcMemory(non_const_x))
|
|
: handler.AcquireSrcMemory(non_const_x);
|
|
|
|
const auto src_y_memory =
|
|
handler.swin_case ? (x.numel() == y.numel()
|
|
? handler.AcquireSecondSrcMemory(non_const_y)
|
|
: handler.AcquireExtendSrcMemory(non_const_y, 1))
|
|
: handler.AcquireSecondSrcMemory(non_const_y);
|
|
|
|
// For Inplace src and dst should be the same memory object.
|
|
// So x should share buffer with z. But UT mechanics is testing inplace
|
|
// execution for this op not checking that x can be broadcasted to match in
|
|
// shape y tensor.
|
|
// This is wrong as when x is to be broadcasted then z(out) will match the
|
|
// shape of y which is bigger than x. Hence if x is smaller in shape than z
|
|
// and they share a buffer (of shape x) then this buffer is not big enough
|
|
// to hold result of elementwise operation.
|
|
const bool reuse_x_memory = non_const_x->numel() == out->numel() &&
|
|
non_const_x->IsSharedBufferWith(*out);
|
|
std::shared_ptr<dnnl::memory> dst_memory;
|
|
|
|
if (reuse_x_memory) {
|
|
dst_memory = src_x_memory;
|
|
// NOTE(chenfeiyu): when the output reuses memory from other tensor rather
|
|
// than allocate its own, it's still need to take care of its data type.
|
|
// Unfortunately, paddle's operator only infers the output' shape, but not
|
|
// the data type. Alloc<T> takes care of allocation and data type
|
|
// normally, but if the memory is already allocated and there is no need
|
|
// to re-allocate, it just set the data type. So this it added there to
|
|
// get the right data type.
|
|
dev_ctx.template Alloc<T>(out);
|
|
} else {
|
|
dst_memory = handler.AcquireDstMemory(out);
|
|
}
|
|
|
|
const auto binary_prim = handler.AcquireForwardPrimitive();
|
|
|
|
auto& astream = OneDNNContext::tls().get_stream();
|
|
|
|
std::unordered_map<int, dnnl::memory> args = {{DNNL_ARG_SRC_0, *src_x_memory},
|
|
{DNNL_ARG_SRC_1, *src_y_memory},
|
|
{DNNL_ARG_DST, *dst_memory}};
|
|
|
|
if (handler.Has_SRC_0_Scale()) {
|
|
args.insert({DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC_0,
|
|
handler.Get_SRC_0_Scale_Memory()});
|
|
}
|
|
|
|
if (handler.Has_SRC_1_Scale()) {
|
|
args.insert({DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC_1,
|
|
handler.Get_SRC_1_Scale_Memory()});
|
|
}
|
|
|
|
binary_prim->execute(astream, args);
|
|
astream.wait();
|
|
|
|
auto out_md = dst_memory->get_desc();
|
|
|
|
if (handler.use_broadcasting_hack) {
|
|
auto dims = out_md.get_dims();
|
|
dims.insert(dims.begin(), non_const_x->dims()[0]);
|
|
dims[1] /= dims[0];
|
|
out_md = out_md.reshape(dims);
|
|
}
|
|
|
|
if (fused_unsqueeze2_axes.empty()) {
|
|
out->set_mem_desc(out_md);
|
|
} else {
|
|
funcs::SetOutMemDescWithUnsqueeze2FuseSupport(
|
|
fused_unsqueeze2_axes, out, out_md);
|
|
}
|
|
}
|
|
|
|
#define DEFINE_ONEDNN_ELEMENTWISE_KERNEL(name, algorithm) \
|
|
template <typename T, typename Context> \
|
|
void name##Kernel(const Context& dev_ctx, \
|
|
const DenseTensor& x, \
|
|
const DenseTensor& y, \
|
|
const int axis, \
|
|
const std::string& fuse_activation, \
|
|
const float fuse_alpha, \
|
|
const float fuse_beta, \
|
|
const float fused_output_scale, \
|
|
const std::vector<int>& fused_unsqueeze2_axes, \
|
|
const float scale_x, \
|
|
const float scale_y, \
|
|
const float scale_out, \
|
|
DenseTensor* out) { \
|
|
FusedElementwiseKernel<T, algorithm>(dev_ctx, \
|
|
x, \
|
|
y, \
|
|
axis, \
|
|
fuse_activation, \
|
|
fuse_alpha, \
|
|
fuse_beta, \
|
|
fused_output_scale, \
|
|
fused_unsqueeze2_axes, \
|
|
scale_x, \
|
|
scale_y, \
|
|
scale_out, \
|
|
out); \
|
|
}
|
|
|
|
DEFINE_ONEDNN_ELEMENTWISE_KERNEL(FusedAdd, dnnl::algorithm::binary_add)
|
|
DEFINE_ONEDNN_ELEMENTWISE_KERNEL(FusedSubtract, dnnl::algorithm::binary_sub)
|
|
DEFINE_ONEDNN_ELEMENTWISE_KERNEL(FusedMultiply, dnnl::algorithm::binary_mul)
|
|
DEFINE_ONEDNN_ELEMENTWISE_KERNEL(FusedDivide, dnnl::algorithm::binary_div)
|
|
|
|
} // namespace phi::fusion
|
|
|
|
PD_REGISTER_KERNEL(fused_elementwise_add,
|
|
OneDNN,
|
|
ONEDNN,
|
|
phi::fusion::FusedAddKernel,
|
|
float,
|
|
phi::bfloat16,
|
|
int8_t,
|
|
uint8_t) {}
|
|
|
|
PD_REGISTER_KERNEL(fused_elementwise_sub,
|
|
OneDNN,
|
|
ONEDNN,
|
|
phi::fusion::FusedSubtractKernel,
|
|
float,
|
|
phi::bfloat16,
|
|
int8_t,
|
|
uint8_t) {}
|
|
|
|
PD_REGISTER_KERNEL(fused_elementwise_mul,
|
|
OneDNN,
|
|
ONEDNN,
|
|
phi::fusion::FusedMultiplyKernel,
|
|
float,
|
|
phi::bfloat16,
|
|
int8_t,
|
|
uint8_t) {}
|
|
|
|
PD_REGISTER_KERNEL(fused_elementwise_div,
|
|
OneDNN,
|
|
ONEDNN,
|
|
phi::fusion::FusedDivideKernel,
|
|
float,
|
|
phi::bfloat16,
|
|
int8_t,
|
|
uint8_t) {}
|