wehub-resource-sync
655 lines
25 KiB
C++
655 lines
25 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/common/errors.h"
|
|
#include "paddle/phi/backends/onednn/onednn_reuse.h"
|
|
#include "paddle/phi/core/compat/convert_utils.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
#include "paddle/phi/core/expect.h"
|
|
#include "paddle/phi/core/utils/data_type.h"
|
|
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
|
|
namespace phi::fusion {
|
|
|
|
using funcs::CreateKey;
|
|
using funcs::OneDNNGetDataType;
|
|
using funcs::OneDNNMemDesc;
|
|
using funcs::RNNReorderType;
|
|
using phi::OneDNNContext;
|
|
using OneDNNMemoryFormat = dnnl::memory::format_tag;
|
|
|
|
template <typename T, typename T_out = T>
|
|
class GRUOneDNNHandler : public funcs::OneDNNHandlerT<T, dnnl::gru_forward> {
|
|
public:
|
|
GRUOneDNNHandler(const OneDNNContext& dev_ctx,
|
|
const dnnl::engine onednn_engine,
|
|
phi::Place cpu_place UNUSED,
|
|
const DenseTensor* input,
|
|
const DenseTensor* weight_h,
|
|
const DenseTensor* h0,
|
|
const bool is_reverse,
|
|
const float scale_data,
|
|
const float shift_data,
|
|
const std::string& gate_activation,
|
|
const std::string& activation,
|
|
const std::vector<float>& scale_weights,
|
|
const int64_t N,
|
|
const int64_t Ti,
|
|
const int64_t IC,
|
|
const int64_t OC)
|
|
: funcs::OneDNNHandlerT<T, dnnl::gru_forward>(
|
|
dev_ctx,
|
|
dev_ctx.GetEngine(),
|
|
cpu_place,
|
|
CreateKey(dev_ctx,
|
|
dev_ctx.GetInputsName("X")[0] +
|
|
dev_ctx.GetInputsName("WeightH")[0],
|
|
OneDNNGetDataType<T>(),
|
|
Ti)),
|
|
N(N),
|
|
Ti(Ti),
|
|
IC(IC),
|
|
OC(OC),
|
|
G(3) {
|
|
std::string unique_name =
|
|
dev_ctx.GetInputsName("X")[0] + dev_ctx.GetInputsName("WeightH")[0];
|
|
// Create memory key without Ti because weights, bias and h0 memories
|
|
// do not depend on Ti size but primitive and input/output memory do
|
|
memory_key_ = funcs::ExtendKeyWithThreadInfoIfNeeded(
|
|
dev_ctx, CreateKey(dev_ctx, unique_name, OneDNNGetDataType<T>()));
|
|
// Is it int8 kernel
|
|
const bool is_INT8 = std::is_same<T, uint8_t>::value;
|
|
if (is_INT8) {
|
|
const int weights_scale_mask =
|
|
0 +
|
|
(1 << 3) // bit, indicating the unique scales for `g` dim in `ldigo`
|
|
+
|
|
(1 << 4); // bit, indicating the unique scales for `o` dim in `ldigo`
|
|
|
|
attr_.set_rnn_data_qparams(scale_data, shift_data);
|
|
attr_.set_rnn_weights_qparams(weights_scale_mask, scale_weights);
|
|
}
|
|
|
|
if (unlikely(!this->isCached())) {
|
|
// oneDNN kernel has hardcoded activation functions
|
|
PADDLE_ENFORCE_EQ(
|
|
gate_activation,
|
|
"sigmoid",
|
|
common::errors::Unimplemented(
|
|
"oneDNN fusion_gru supports only sigmoid as a gate activation."));
|
|
PADDLE_ENFORCE_EQ(
|
|
activation,
|
|
"tanh",
|
|
common::errors::Unimplemented(
|
|
"oneDNN fusion_gru supports only tanh as an activation."));
|
|
|
|
// Weights for int8 kernel are of a type s8
|
|
const auto weights_dt =
|
|
is_INT8 ? dnnl::memory::data_type::s8 : OneDNNGetDataType<T>();
|
|
|
|
// oneDNN RNN dimensions
|
|
const int64_t D = 1; // Directions
|
|
const int64_t L = 1; // Layers (PP supports only 1 stacked layer)
|
|
const int64_t G = 3; // Number of Gates, 3 for GRU
|
|
|
|
// Create memory descriptors
|
|
auto input_md = OneDNNMemDesc(
|
|
{Ti, N, IC}, OneDNNGetDataType<T>(), OneDNNMemoryFormat::any);
|
|
auto weight_x_md =
|
|
OneDNNMemDesc({L, D, IC, G, OC}, weights_dt, OneDNNMemoryFormat::any);
|
|
auto weight_h_md =
|
|
OneDNNMemDesc({L, D, OC, G, OC}, weights_dt, OneDNNMemoryFormat::any);
|
|
auto bias_md = OneDNNMemDesc(
|
|
{L, D, G, OC}, OneDNNGetDataType<float>(), OneDNNMemoryFormat::ldgo);
|
|
auto hidden_md = OneDNNMemDesc(
|
|
{Ti, N, OC}, OneDNNGetDataType<T_out>(), OneDNNMemoryFormat::any);
|
|
auto h0_md = OneDNNMemDesc(
|
|
{L, D, N, OC}, OneDNNGetDataType<T>(), OneDNNMemoryFormat::ldnc);
|
|
|
|
// Create GRU oneDNN primitive
|
|
const auto direction =
|
|
is_reverse ? dnnl::rnn_direction::unidirectional_right2left
|
|
: dnnl::rnn_direction::unidirectional_left2right;
|
|
|
|
this->AcquireForwardPrimitiveDescriptor(
|
|
this->attr_,
|
|
dnnl::prop_kind::forward_inference,
|
|
direction,
|
|
input_md,
|
|
h0_md,
|
|
weight_x_md,
|
|
weight_h_md,
|
|
bias_md,
|
|
hidden_md,
|
|
dnnl::memory::desc());
|
|
}
|
|
}
|
|
|
|
bool is_NTC() { return this->is_NTC(this->fwd_pd_->dst_desc()); }
|
|
|
|
bool is_NTC(const dnnl::memory::desc& md) {
|
|
auto ntc_md = dnnl::memory::desc(
|
|
md.get_dims(), md.get_data_type(), dnnl::memory::format_tag::ntc);
|
|
return md == ntc_md;
|
|
}
|
|
|
|
void reorderRNNdata(void* input_data,
|
|
void* output_data,
|
|
std::vector<size_t> lod,
|
|
const bool is_reverse,
|
|
RNNReorderType reorder_type) {
|
|
switch (reorder_type) {
|
|
// Reorder input memory [WORDS, C] + LoD -> [N, T, C]
|
|
case RNNReorderType::PP_NTC: {
|
|
auto* input_data_iter = reinterpret_cast<T*>(input_data);
|
|
auto* output_data_iter = reinterpret_cast<T*>(output_data);
|
|
for (int n = 0; n < N; ++n) {
|
|
const auto num_elements = (lod[n + 1] - lod[n]) * IC;
|
|
const auto offset = is_reverse ? (Ti * IC - num_elements) : 0;
|
|
memcpy(output_data_iter + n * Ti * IC + offset,
|
|
input_data_iter,
|
|
sizeof(T) * num_elements);
|
|
input_data_iter += num_elements;
|
|
}
|
|
} break;
|
|
// Reorder input memory [WORDS, C] + LoD -> [T, N, C]
|
|
case RNNReorderType::PP_TNC: {
|
|
auto* input_data_iter = reinterpret_cast<T*>(input_data);
|
|
auto* output_data_iter = reinterpret_cast<T*>(output_data);
|
|
for (int n = 0; n < N; ++n) {
|
|
const auto num_elements = (lod[n + 1] - lod[n]);
|
|
const auto offset = is_reverse ? (Ti - num_elements) : 0;
|
|
for (size_t t = 0; t < num_elements; ++t) {
|
|
memcpy(output_data_iter + (t + offset) * N * IC + n * IC,
|
|
input_data_iter,
|
|
sizeof(T) * IC);
|
|
input_data_iter += IC;
|
|
}
|
|
}
|
|
} break;
|
|
// Reorder output values to PP format [N, T, C] -> [WORDS, C]
|
|
case RNNReorderType::NTC_PP: {
|
|
auto* input_data_iter = reinterpret_cast<T_out*>(input_data);
|
|
auto* output_data_iter = reinterpret_cast<T_out*>(output_data);
|
|
for (int n = 0; n < N; ++n) {
|
|
const auto num_elements = (lod[n + 1] - lod[n]) * OC;
|
|
const auto offset = is_reverse ? (Ti * OC - num_elements) : 0;
|
|
memcpy(output_data_iter,
|
|
input_data_iter + n * Ti * OC + offset,
|
|
sizeof(T_out) * num_elements);
|
|
output_data_iter += num_elements;
|
|
}
|
|
} break;
|
|
// Reorder output values to PP format [T, N, C] -> [WORDS, C]
|
|
case RNNReorderType::TNC_PP: {
|
|
auto* input_data_iter = reinterpret_cast<T_out*>(input_data);
|
|
auto* output_data_iter = reinterpret_cast<T_out*>(output_data);
|
|
for (int n = 0; n < N; ++n) {
|
|
const auto num_elements = lod[n + 1] - lod[n];
|
|
const auto offset = is_reverse ? (Ti - num_elements) : 0;
|
|
for (size_t t = 0; t < num_elements; ++t) {
|
|
memcpy(output_data_iter,
|
|
input_data_iter + (t + offset) * N * OC + n * OC,
|
|
sizeof(T_out) * OC);
|
|
output_data_iter += OC;
|
|
}
|
|
}
|
|
} break;
|
|
}
|
|
}
|
|
|
|
std::shared_ptr<dnnl::memory> AcquireInputMemoryWithReorder(
|
|
const DenseTensor* input, const bool is_reverse) {
|
|
const auto name = this->key_ + "@input_mem";
|
|
auto memory_p =
|
|
std::static_pointer_cast<dnnl::memory>(this->dev_ctx_.GetBlob(name));
|
|
|
|
if (!memory_p) {
|
|
memory_p = std::make_shared<dnnl::memory>(this->fwd_pd_->src_desc(),
|
|
this->engine_);
|
|
this->dev_ctx_.SetBlob(name, memory_p);
|
|
}
|
|
|
|
const auto& input_lod = input->lod()[0];
|
|
auto* x_data = funcs::to_void_cast(input->data<T>());
|
|
|
|
auto* x_onednn_data = memory_p->get_data_handle();
|
|
memset(x_onednn_data, 0, sizeof(T) * N * Ti * IC);
|
|
|
|
if (is_NTC(this->fwd_pd_->src_desc())) {
|
|
reorderRNNdata(
|
|
x_data, x_onednn_data, input_lod, is_reverse, RNNReorderType::PP_NTC);
|
|
} else {
|
|
reorderRNNdata(
|
|
x_data, x_onednn_data, input_lod, is_reverse, RNNReorderType::PP_TNC);
|
|
}
|
|
return memory_p;
|
|
}
|
|
|
|
std::shared_ptr<dnnl::memory> AcquireOutputMemory() {
|
|
const auto name = this->key_ + "@output_mem";
|
|
auto memory_p =
|
|
std::static_pointer_cast<dnnl::memory>(this->dev_ctx_.GetBlob(name));
|
|
|
|
if (!memory_p) {
|
|
memory_p = std::make_shared<dnnl::memory>(this->fwd_pd_->dst_desc(),
|
|
this->engine_);
|
|
this->dev_ctx_.SetBlob(name, memory_p);
|
|
}
|
|
return memory_p;
|
|
}
|
|
|
|
// H0 is for now persistable
|
|
template <typename U>
|
|
std::shared_ptr<dnnl::memory> AcquireH0Memory(const DenseTensor* h0) {
|
|
const std::string h0_key = memory_key_ + "@h0";
|
|
auto memory_p =
|
|
std::static_pointer_cast<dnnl::memory>(this->dev_ctx_.GetBlob(h0_key));
|
|
|
|
if (!memory_p) {
|
|
auto user_h0_memory = dnnl::memory();
|
|
if (h0) {
|
|
user_h0_memory = dnnl::memory(
|
|
{{1, 1, N, OC}, OneDNNGetDataType<U>(), OneDNNMemoryFormat::ldnc},
|
|
this->engine_,
|
|
funcs::to_void_cast(h0->data<U>()));
|
|
} else {
|
|
user_h0_memory = dnnl::memory(
|
|
{{1, 1, N, OC}, OneDNNGetDataType<U>(), OneDNNMemoryFormat::ldnc},
|
|
this->engine_);
|
|
memset(user_h0_memory.get_data_handle(), 0, sizeof(U) * N * OC);
|
|
}
|
|
memory_p = std::make_shared<dnnl::memory>(this->fwd_pd_->src_iter_desc(),
|
|
this->engine_);
|
|
|
|
auto& astream = phi::OneDNNContext::tls().get_stream();
|
|
dnnl::reorder(user_h0_memory, *memory_p, attr_)
|
|
.execute(astream, user_h0_memory, *memory_p);
|
|
|
|
this->dev_ctx_.SetBlob(h0_key, memory_p);
|
|
}
|
|
return memory_p;
|
|
}
|
|
|
|
template <typename U>
|
|
std::shared_ptr<dnnl::memory> AcquireWeightXMemory(
|
|
const DenseTensor* weight_x, const bool origin_mode) {
|
|
const std::string wx_key = this->memory_key_ + "@weight_x";
|
|
auto memory_p =
|
|
std::static_pointer_cast<dnnl::memory>(this->dev_ctx_.GetBlob(wx_key));
|
|
|
|
if (!memory_p) {
|
|
auto user_md = OneDNNMemDesc({1, 1, this->IC, this->G, this->OC},
|
|
OneDNNGetDataType<U>(),
|
|
OneDNNMemoryFormat::ldigo);
|
|
auto user_memory = dnnl::memory(user_md, this->engine_);
|
|
|
|
auto* weight_x_data = reinterpret_cast<U*>(user_memory.get_data_handle());
|
|
memcpy(weight_x_data,
|
|
weight_x->data<U>(),
|
|
sizeof(U) * this->IC * this->G * this->OC);
|
|
|
|
if (origin_mode == false) {
|
|
for (int64_t i = 0; i < this->IC; ++i) {
|
|
for (int64_t j = 0; j < this->OC; ++j) {
|
|
U minus_one(-1.0f);
|
|
weight_x_data[j] = minus_one * weight_x_data[j];
|
|
}
|
|
weight_x_data += 3 * this->OC;
|
|
}
|
|
}
|
|
|
|
memory_p = std::make_shared<dnnl::memory>(
|
|
this->fwd_pd_->weights_layer_desc(), this->engine_);
|
|
|
|
auto& astream = OneDNNContext::tls().get_stream();
|
|
dnnl::reorder(user_memory, *memory_p, this->attr_)
|
|
.execute(astream, user_memory, *memory_p);
|
|
|
|
this->dev_ctx_.SetBlob(wx_key, memory_p);
|
|
}
|
|
return memory_p;
|
|
}
|
|
|
|
template <typename U>
|
|
std::shared_ptr<dnnl::memory> AcquireWeightHMemory(
|
|
const DenseTensor* weight_h, const bool origin_mode) {
|
|
const std::string wh_key = this->memory_key_ + "@weight_h";
|
|
auto memory_p =
|
|
std::static_pointer_cast<dnnl::memory>(this->dev_ctx_.GetBlob(wh_key));
|
|
|
|
if (!memory_p) {
|
|
auto user_md = OneDNNMemDesc({1, 1, this->OC, this->G, this->OC},
|
|
OneDNNGetDataType<U>(),
|
|
OneDNNMemoryFormat::ldigo);
|
|
auto user_memory = dnnl::memory(user_md, this->engine_);
|
|
|
|
// Reorder weights_h from PP format [OC, 2OC] + [OC, OC] to
|
|
// oneDNN format [OC, 3OC]
|
|
auto* weight_h_data = reinterpret_cast<U*>(user_memory.get_data_handle());
|
|
auto* user_weight_h_data = weight_h->data<U>();
|
|
|
|
auto src1_iter = user_weight_h_data;
|
|
auto src2_iter = user_weight_h_data + 2 * this->OC * this->OC;
|
|
|
|
for (int64_t c = 0; c < this->OC; ++c) {
|
|
memcpy(weight_h_data, src1_iter, 2 * this->OC * sizeof(U));
|
|
memcpy(weight_h_data + 2 * this->OC, src2_iter, this->OC * sizeof(U));
|
|
|
|
src1_iter += 2 * this->OC;
|
|
src2_iter += this->OC;
|
|
weight_h_data += 3 * this->OC;
|
|
}
|
|
|
|
weight_h_data = reinterpret_cast<U*>(user_memory.get_data_handle());
|
|
|
|
if (origin_mode == false) {
|
|
for (int64_t i = 0; i < this->OC; ++i) {
|
|
for (int64_t j = 0; j < this->OC; ++j) {
|
|
U minus_one(-1.0f);
|
|
weight_h_data[j] = minus_one * weight_h_data[j];
|
|
}
|
|
weight_h_data += 3 * this->OC;
|
|
}
|
|
}
|
|
|
|
memory_p = std::make_shared<dnnl::memory>(
|
|
this->fwd_pd_->weights_iter_desc(), this->engine_);
|
|
|
|
auto& astream = OneDNNContext::tls().get_stream();
|
|
dnnl::reorder(user_memory, *memory_p, this->attr_)
|
|
.execute(astream, user_memory, *memory_p);
|
|
|
|
this->dev_ctx_.SetBlob(wh_key, memory_p);
|
|
}
|
|
return memory_p;
|
|
}
|
|
|
|
std::shared_ptr<dnnl::memory> AcquireBiasMemory(const DenseTensor* bias,
|
|
const bool origin_mode) {
|
|
const std::string bias_key = this->memory_key_ + "@bias";
|
|
auto memory_p = std::static_pointer_cast<dnnl::memory>(
|
|
this->dev_ctx_.GetBlob(bias_key));
|
|
|
|
if (!memory_p) {
|
|
memory_p = std::make_shared<dnnl::memory>(this->fwd_pd_->bias_desc(),
|
|
this->engine_);
|
|
auto* bias_data = reinterpret_cast<float*>(memory_p->get_data_handle());
|
|
if (bias) {
|
|
const float* user_bias_data =
|
|
bias->data<float>(); // Bias in oneDNN is always float
|
|
memcpy(bias_data, user_bias_data, sizeof(float) * this->G * this->OC);
|
|
} else {
|
|
// oneDNN always need bias memory, if it's not provided in PP, let
|
|
// oneDNN allocate memory and set it to 0
|
|
memset(bias_data, 0, sizeof(float) * this->G * this->OC);
|
|
}
|
|
|
|
if (origin_mode == false && bias) {
|
|
for (int64_t i = 0; i < this->OC; ++i) {
|
|
bias_data[i] *= -1;
|
|
}
|
|
}
|
|
this->dev_ctx_.SetBlob(bias_key, memory_p);
|
|
}
|
|
return memory_p;
|
|
}
|
|
|
|
protected:
|
|
// RNN dimensions
|
|
// N - Batch Size
|
|
// Ti - Max sentence length
|
|
// IC - Input Channels
|
|
// OC - Output Channels
|
|
// G - Number of gates
|
|
const int64_t N, Ti, IC, OC, G;
|
|
|
|
// Memory size of weights, bias and h0 does not depend
|
|
// on Ti size, thus we need another key to cache them
|
|
std::string memory_key_;
|
|
dnnl::primitive_attr attr_;
|
|
};
|
|
|
|
template <typename T, typename Tout = T>
|
|
void RunKernel(const phi::OneDNNContext& dev_ctx,
|
|
const DenseTensor& x,
|
|
const optional<DenseTensor>& h0,
|
|
const DenseTensor& weight_x,
|
|
const DenseTensor& weight_h,
|
|
const optional<DenseTensor>& bias,
|
|
const std::string& activation,
|
|
const std::string& gate_activation,
|
|
const bool is_reverse,
|
|
const bool use_seq,
|
|
const bool origin_mode,
|
|
const float scale_data,
|
|
const float shift_data,
|
|
const std::vector<float>& scale_weights,
|
|
DenseTensor* reordered_h0,
|
|
DenseTensor* xx,
|
|
DenseTensor* batched_input,
|
|
DenseTensor* batched_out,
|
|
DenseTensor* hidden) {
|
|
const auto& onednn_engine = dev_ctx.GetEngine();
|
|
|
|
auto x_dims = x.dims();
|
|
auto x_mat_dims = (x_dims.size() == 3 && x_dims[1] == 1)
|
|
? common::flatten_to_2d(x_dims, 1)
|
|
: x_dims;
|
|
|
|
// Get tensor dimensions
|
|
const auto x_mat_dims_vec = vectorize(x_mat_dims);
|
|
const auto weight_h_dims = vectorize(weight_h.dims());
|
|
const auto& input_lod = x.lod()[0];
|
|
|
|
// Calculate RNN dimensions
|
|
const int64_t N = static_cast<int64_t>(input_lod.size() -
|
|
1); // Number of sentences (batches)
|
|
const int64_t Ti = // Max length of the sentence in a batch
|
|
[&input_lod]() {
|
|
size_t res = 0;
|
|
for (size_t i = 0; i < (input_lod.size() - 1); ++i) {
|
|
res = std::max(res, input_lod[i + 1] - input_lod[i]);
|
|
}
|
|
return res;
|
|
}();
|
|
const int64_t IC = x_mat_dims_vec[1]; // Input channels
|
|
const int64_t OC = weight_h_dims[0]; // Output channels
|
|
|
|
GRUOneDNNHandler<T, Tout> handler(dev_ctx,
|
|
onednn_engine,
|
|
dev_ctx.GetPlace(),
|
|
&x,
|
|
&weight_h,
|
|
h0.get_ptr(),
|
|
is_reverse,
|
|
scale_data,
|
|
shift_data,
|
|
gate_activation,
|
|
activation,
|
|
scale_weights,
|
|
N,
|
|
Ti,
|
|
IC,
|
|
OC);
|
|
auto input_memory_p = handler.AcquireInputMemoryWithReorder(&x, is_reverse);
|
|
|
|
std::shared_ptr<dnnl::memory> h0_memory_p, weight_h_memory_p,
|
|
weight_x_memory_p;
|
|
|
|
if (weight_h.dtype() == DataType::FLOAT32) {
|
|
h0_memory_p = handler.template AcquireH0Memory<float>(h0.get_ptr());
|
|
weight_x_memory_p =
|
|
handler.template AcquireWeightXMemory<float>(&weight_x, origin_mode);
|
|
weight_h_memory_p =
|
|
handler.template AcquireWeightHMemory<float>(&weight_h, origin_mode);
|
|
} else if (weight_h.dtype() == DataType::BFLOAT16) {
|
|
h0_memory_p = handler.template AcquireH0Memory<phi::bfloat16>(h0.get_ptr());
|
|
weight_x_memory_p = handler.template AcquireWeightXMemory<phi::bfloat16>(
|
|
&weight_x, origin_mode);
|
|
weight_h_memory_p = handler.template AcquireWeightHMemory<phi::bfloat16>(
|
|
&weight_h, origin_mode);
|
|
} else {
|
|
h0_memory_p = handler.template AcquireH0Memory<uint8_t>(h0.get_ptr());
|
|
weight_x_memory_p =
|
|
handler.template AcquireWeightXMemory<int8_t>(&weight_x, origin_mode);
|
|
weight_h_memory_p =
|
|
handler.template AcquireWeightHMemory<int8_t>(&weight_h, origin_mode);
|
|
}
|
|
|
|
auto bias_memory_p = handler.AcquireBiasMemory(bias.get_ptr(), origin_mode);
|
|
auto hidden_onednn_memory_p = handler.AcquireOutputMemory();
|
|
|
|
std::unordered_map<int, dnnl::memory> gru_args = {
|
|
{DNNL_ARG_SRC_LAYER, *input_memory_p},
|
|
{DNNL_ARG_SRC_ITER, *h0_memory_p},
|
|
{DNNL_ARG_WEIGHTS_LAYER, *weight_x_memory_p},
|
|
{DNNL_ARG_WEIGHTS_ITER, *weight_h_memory_p},
|
|
{DNNL_ARG_BIAS, *bias_memory_p},
|
|
{DNNL_ARG_DST_LAYER, *hidden_onednn_memory_p}};
|
|
|
|
auto gru_forward_p = handler.AcquireForwardPrimitive();
|
|
|
|
auto& astream = OneDNNContext::tls().get_stream();
|
|
gru_forward_p->execute(astream, gru_args);
|
|
astream.wait();
|
|
|
|
auto* hidden_onednn_data = hidden_onednn_memory_p->get_data_handle();
|
|
auto* hidden_tmp_data = dev_ctx.template Alloc<Tout>(hidden);
|
|
auto* hidden_data = funcs::to_void_cast(hidden_tmp_data);
|
|
if (handler.is_NTC()) {
|
|
handler.reorderRNNdata(hidden_onednn_data,
|
|
hidden_data,
|
|
input_lod,
|
|
is_reverse,
|
|
RNNReorderType::NTC_PP);
|
|
} else {
|
|
handler.reorderRNNdata(hidden_onednn_data,
|
|
hidden_data,
|
|
input_lod,
|
|
is_reverse,
|
|
RNNReorderType::TNC_PP);
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FusionGRUKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const optional<DenseTensor>& h0,
|
|
const DenseTensor& weight_x,
|
|
const DenseTensor& weight_h,
|
|
const optional<DenseTensor>& bias,
|
|
const std::string& activation,
|
|
const std::string& gate_activation,
|
|
const bool is_reverse,
|
|
const bool use_seq,
|
|
const bool origin_mode,
|
|
const bool force_fp32_output,
|
|
DenseTensor* reordered_h0,
|
|
DenseTensor* xx,
|
|
DenseTensor* batched_input,
|
|
DenseTensor* batched_out,
|
|
DenseTensor* hidden) {
|
|
const std::string mkldnn_data_type =
|
|
dev_ctx.HasDnnAttr("mkldnn_data_type")
|
|
? PADDLE_GET_CONST(std::string,
|
|
dev_ctx.GetDnnAttr("mkldnn_data_type"))
|
|
: "float32";
|
|
const std::string onednn_data_type =
|
|
(dev_ctx.HasDnnAttr("onednn_data_type") &&
|
|
PADDLE_GET_CONST(std::string, dev_ctx.GetDnnAttr("onednn_data_type")) !=
|
|
"")
|
|
? PADDLE_GET_CONST(std::string,
|
|
dev_ctx.GetDnnAttr("onednn_data_type"))
|
|
: mkldnn_data_type;
|
|
std::vector<std::string> onednn_data_type_list = {
|
|
"float32", "int8", "bfloat16"};
|
|
PADDLE_ENFORCE_EQ(std::find(onednn_data_type_list.begin(),
|
|
onednn_data_type_list.end(),
|
|
onednn_data_type) != onednn_data_type_list.end(),
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"The onednn_data_type should be [float32, "
|
|
"int8, bfloat16], but found %s.",
|
|
onednn_data_type.c_str()));
|
|
const float scale_data =
|
|
dev_ctx.HasDnnAttr("Scale_data")
|
|
? PADDLE_GET_CONST(float, dev_ctx.GetDnnAttr("Scale_data"))
|
|
: 1.0f;
|
|
const float shift_data =
|
|
dev_ctx.HasDnnAttr("Shift_data")
|
|
? PADDLE_GET_CONST(float, dev_ctx.GetDnnAttr("Shift_data"))
|
|
: 1.0f;
|
|
std::vector<float> tmp_scale_weights = {1.0f};
|
|
const std::vector<float> scale_weights =
|
|
dev_ctx.HasDnnAttr("Scale_weights")
|
|
? PADDLE_GET_CONST(std::vector<float>,
|
|
dev_ctx.GetDnnAttr("Scale_weights"))
|
|
: tmp_scale_weights;
|
|
const bool is_bf16 = std::is_same<T, phi::bfloat16>::value;
|
|
// BF16 does not support force output
|
|
if (!is_bf16 && force_fp32_output) { // NOLINT
|
|
RunKernel<T, float>(dev_ctx,
|
|
x,
|
|
h0,
|
|
weight_x,
|
|
weight_h,
|
|
bias,
|
|
activation,
|
|
gate_activation,
|
|
is_reverse,
|
|
use_seq,
|
|
origin_mode,
|
|
scale_data,
|
|
shift_data,
|
|
scale_weights,
|
|
reordered_h0,
|
|
xx,
|
|
batched_input,
|
|
batched_out,
|
|
hidden);
|
|
} else {
|
|
RunKernel<T>(dev_ctx,
|
|
x,
|
|
h0,
|
|
weight_x,
|
|
weight_h,
|
|
bias,
|
|
activation,
|
|
gate_activation,
|
|
is_reverse,
|
|
use_seq,
|
|
origin_mode,
|
|
scale_data,
|
|
shift_data,
|
|
scale_weights,
|
|
reordered_h0,
|
|
xx,
|
|
batched_input,
|
|
batched_out,
|
|
hidden);
|
|
}
|
|
}
|
|
|
|
} // namespace phi::fusion
|
|
|
|
PD_REGISTER_KERNEL(fusion_gru,
|
|
OneDNN,
|
|
ONEDNN,
|
|
phi::fusion::FusionGRUKernel,
|
|
float,
|
|
phi::bfloat16,
|
|
uint8_t) {}
|