wehub-resource-sync
143 lines
5.5 KiB
C++
143 lines
5.5 KiB
C++
// Copyright (c) 2024 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/layer_norm_kernel.h"
|
|
|
|
#include "paddle/phi/backends/onednn/onednn_reuse.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T>
|
|
class LayerNormOneDNNHandler
|
|
: public funcs::OneDNNHandlerNoCachingT<T,
|
|
dnnl::layer_normalization_forward> {
|
|
public:
|
|
LayerNormOneDNNHandler(const std::vector<int64_t>& dims,
|
|
const float& epsilon,
|
|
const dnnl::normalization_flags& flags,
|
|
const bool& is_test,
|
|
const DenseTensor* x,
|
|
const dnnl::engine engine,
|
|
Place cpu_place)
|
|
: funcs::OneDNNHandlerNoCachingT<T, dnnl::layer_normalization_forward>(
|
|
engine, cpu_place) {
|
|
const auto fwd_prop_kind = is_test ? dnnl::prop_kind::forward_inference
|
|
: dnnl::prop_kind::forward_training;
|
|
|
|
this->AcquireForwardPrimitiveDescriptor(
|
|
fwd_prop_kind, x->mem_desc(), x->mem_desc(), epsilon, flags);
|
|
}
|
|
|
|
std::tuple<std::shared_ptr<dnnl::memory>, std::shared_ptr<dnnl::memory>>
|
|
AcquireScaleShiftMemory(const DenseTensor* scale, const DenseTensor* shift) {
|
|
auto scale_memory = this->AcquireMemoryFromPrimitive(
|
|
this->fwd_pd_->weights_desc(),
|
|
funcs::to_void_cast<float>(scale->data<float>()));
|
|
auto shift_memory = this->AcquireMemoryFromPrimitive(
|
|
this->fwd_pd_->weights_desc(),
|
|
funcs::to_void_cast<float>(shift->data<float>()));
|
|
|
|
return std::make_tuple(scale_memory, shift_memory);
|
|
}
|
|
|
|
std::shared_ptr<dnnl::memory> AcquireMeanMemory(const OneDNNContext& dev_ctx,
|
|
DenseTensor* mean) {
|
|
float* mean_data = dev_ctx.template Alloc<float>(
|
|
mean, this->fwd_pd_->mean_desc().get_size());
|
|
return this->AcquireMemoryFromPrimitive(this->fwd_pd_->mean_desc(),
|
|
mean_data);
|
|
}
|
|
|
|
std::shared_ptr<dnnl::memory> AcquireVarianceMemory(
|
|
const OneDNNContext& dev_ctx, DenseTensor* variance) {
|
|
float* variance_data = dev_ctx.template Alloc<float>(
|
|
variance, this->fwd_pd_->variance_desc().get_size());
|
|
return this->AcquireMemoryFromPrimitive(this->fwd_pd_->variance_desc(),
|
|
variance_data);
|
|
}
|
|
};
|
|
|
|
template <typename T, typename Context>
|
|
void LayerNormKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const optional<DenseTensor>& scale_opt,
|
|
const optional<DenseTensor>& bias_opt,
|
|
double epsilon,
|
|
int begin_norm_axis,
|
|
DenseTensor* y,
|
|
DenseTensor* mean,
|
|
DenseTensor* var) {
|
|
bool is_test = dev_ctx.HasDnnAttr("is_test")
|
|
? PADDLE_GET_CONST(bool, dev_ctx.GetDnnAttr("is_test"))
|
|
: false;
|
|
|
|
const auto& onednn_engine = dev_ctx.GetEngine();
|
|
|
|
auto src_tz = vectorize(x.dims());
|
|
PADDLE_ENFORCE_EQ(begin_norm_axis,
|
|
(src_tz.size() - 1),
|
|
common::errors::InvalidArgument(
|
|
"MKL-DNN Layer Norm supports only last logical "
|
|
"axis:%d as begin_norm_axis.",
|
|
(src_tz.size() - 1)));
|
|
|
|
const bool with_scaleshift = (scale_opt && bias_opt);
|
|
dnnl::normalization_flags flags{};
|
|
|
|
if (with_scaleshift) {
|
|
flags |= dnnl::normalization_flags::use_scale |
|
|
dnnl::normalization_flags::use_shift;
|
|
}
|
|
|
|
LayerNormOneDNNHandler<T> handler(
|
|
src_tz, epsilon, flags, is_test, &x, onednn_engine, dev_ctx.GetPlace());
|
|
|
|
auto src_memory = handler.AcquireSrcMemory(&x);
|
|
auto dst_memory = handler.AcquireDstMemory(y);
|
|
|
|
auto layer_norm_p = handler.AcquireForwardPrimitive();
|
|
|
|
auto& astream = OneDNNContext::tls().get_stream();
|
|
std::unordered_map<int, dnnl::memory> args = {{DNNL_ARG_SRC, *src_memory},
|
|
{DNNL_ARG_DST, *dst_memory}};
|
|
|
|
if (!is_test) {
|
|
auto mean_memory = handler.AcquireMeanMemory(dev_ctx, mean);
|
|
auto variance_memory = handler.AcquireVarianceMemory(dev_ctx, var);
|
|
|
|
args.insert({DNNL_ARG_MEAN, *mean_memory});
|
|
args.insert({DNNL_ARG_VARIANCE, *variance_memory});
|
|
}
|
|
|
|
if (with_scaleshift) {
|
|
auto scaleshift_mems = handler.AcquireScaleShiftMemory(scale_opt.get_ptr(),
|
|
bias_opt.get_ptr());
|
|
args.insert({DNNL_ARG_SCALE, *(std::get<0>(scaleshift_mems))});
|
|
args.insert({DNNL_ARG_SHIFT, *(std::get<1>(scaleshift_mems))});
|
|
}
|
|
|
|
layer_norm_p->execute(astream, args);
|
|
astream.wait();
|
|
|
|
y->set_mem_desc(dst_memory->get_desc());
|
|
}
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(
|
|
layer_norm, OneDNN, ONEDNN, phi::LayerNormKernel, float, phi::bfloat16) {
|
|
kernel->OutputAt(1).SetDataType(phi::DataType::UNDEFINED);
|
|
kernel->OutputAt(2).SetDataType(phi::DataType::UNDEFINED);
|
|
}
|