[Eernie TRT]: add slice op and add emb eltwise layernorm fp16 support (#23723)
* refine ernie trt dynamic shape support 1. add slice op converter 2. add emb eltwise layernorm fp16 support test=develop * fix dynamic shape test ut test=develop * fix comments. test=develop * fix comments test=developrevert-22778-infer_var_type
Zhaolong Xing
5 years ago
committed by
GitHub
parent
2b896c1f6b
commit
133f1fc123
@ -0,0 +1,69 @@
|
||||
/* Copyright (c) 2018 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/fluid/inference/tensorrt/convert/op_converter.h"
|
||||
#include "paddle/fluid/inference/tensorrt/plugin/slice_op_plugin.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
|
||||
class SliceOpConverter : public OpConverter {
|
||||
public:
|
||||
void operator()(const framework::proto::OpDesc& op,
|
||||
const framework::Scope& scope, bool test_mode) override {
|
||||
// This OP is implemented by trt dynamic shpae plugin.
|
||||
// Dynamic shape plugin requires TRT version greater than 6.0.
|
||||
#if IS_TRT_VERSION_GE(6000)
|
||||
VLOG(4) << "convert slice op to tensorrt layer";
|
||||
framework::OpDesc op_desc(op, nullptr);
|
||||
// Declare inputs
|
||||
auto* input = engine_->GetITensor(op_desc.Input("Input")[0]);
|
||||
|
||||
std::vector<int> axes =
|
||||
boost::get<std::vector<int>>(op_desc.GetAttr("axes"));
|
||||
std::vector<int> starts =
|
||||
boost::get<std::vector<int>>(op_desc.GetAttr("starts"));
|
||||
std::vector<int> ends =
|
||||
boost::get<std::vector<int>>(op_desc.GetAttr("ends"));
|
||||
|
||||
nvinfer1::ILayer* layer = nullptr;
|
||||
if (engine_->with_dynamic_shape()) {
|
||||
bool ban_fp16 = engine_->disable_trt_plugin_fp16();
|
||||
plugin::SlicePluginDynamic* plugin =
|
||||
new plugin::SlicePluginDynamic(starts, ends, ends, ban_fp16);
|
||||
layer = engine_->AddPluginV2(&input, 1, plugin);
|
||||
} else {
|
||||
PADDLE_THROW(platform::errors::Fatal(
|
||||
"You are running the Ernie(Bert) model in static"
|
||||
"shape mode, which is not supported for the time being.\n"
|
||||
"You can use the config.SetTRTDynamicShapeInfo(...) interface"
|
||||
" to set the shape information to run the dynamic shape mode."));
|
||||
}
|
||||
|
||||
auto output_name = op_desc.Output("Out")[0];
|
||||
RreplenishLayerAndOutput(layer, "skip_layernorm", {output_name}, test_mode);
|
||||
#else
|
||||
PADDLE_THROW(platform::errors::Fatal(
|
||||
"You are running the TRT Dynamic Shape mode, need to confirm that "
|
||||
"your TRT version is no less than 6.0"));
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace tensorrt
|
||||
} // namespace inference
|
||||
} // namespace paddle
|
||||
|
||||
REGISTER_TRT_OP_CONVERTER(slice, SliceOpConverter);
|
||||
@ -0,0 +1,208 @@
|
||||
// Copyright (c) 2018 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 <cuda_runtime.h>
|
||||
#include <stdio.h>
|
||||
#include <cassert>
|
||||
#include <cub/cub.cuh> // NOLINT
|
||||
#include <vector>
|
||||
#include "glog/logging.h"
|
||||
#include "paddle/fluid/inference/tensorrt/plugin/slice_op_plugin.h"
|
||||
#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin_factory.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
namespace plugin {
|
||||
|
||||
// Dynamic Plugin below.
|
||||
#if IS_TRT_VERSION_GE(6000)
|
||||
|
||||
template <typename T>
|
||||
__global__ void SliceKernel(int num, int dims, const T *input,
|
||||
const int *offsets_info, T *output) {
|
||||
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
extern __shared__ int shared_data[];
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
for (int i = 0; i < dims * 3; i++) {
|
||||
shared_data[i] = offsets_info[i];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (idx < num) {
|
||||
int t_idx = idx;
|
||||
int in_idx = 0;
|
||||
for (int i = dims - 1; i >= 0; i--) {
|
||||
// output_shape
|
||||
auto t = t_idx % shared_data[i * 3 + 1];
|
||||
// out offset
|
||||
auto s = t + shared_data[i * 3];
|
||||
// input_seg_offset
|
||||
in_idx = in_idx + shared_data[i * 3 + 2] * s;
|
||||
t_idx = t_idx / shared_data[i * 3 + 1];
|
||||
}
|
||||
output[idx] = input[in_idx];
|
||||
}
|
||||
}
|
||||
|
||||
int SlicePluginDynamic::initialize() { return 0; }
|
||||
|
||||
size_t SlicePluginDynamic::getSerializationSize() const { return 0; }
|
||||
|
||||
void SlicePluginDynamic::serialize(void *buffer) const {}
|
||||
|
||||
nvinfer1::DimsExprs SlicePluginDynamic::getOutputDimensions(
|
||||
int output_index, const nvinfer1::DimsExprs *inputs, int nb_inputs,
|
||||
nvinfer1::IExprBuilder &expr_builder) {
|
||||
auto in_dims = inputs[0];
|
||||
nvinfer1::DimsExprs ret;
|
||||
for (int i = 0; i < ret.nbDims; i++) {
|
||||
ret.d[i] = in_dims.d[i];
|
||||
}
|
||||
// start, ends should greater 0
|
||||
for (size_t i = 0; i < axes_.size(); i++) {
|
||||
int start = starts_[i];
|
||||
int end = ends_[i];
|
||||
ret.d[axes_[i]] = expr_builder.constant(end - start);
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
|
||||
bool SlicePluginDynamic::supportsFormatCombination(
|
||||
int pos, const nvinfer1::PluginTensorDesc *in_out, int nb_inputs,
|
||||
int nb_outputs) {
|
||||
PADDLE_ENFORCE_NOT_NULL(
|
||||
in_out, platform::errors::InvalidArgument(
|
||||
"The input of swish plugin shoule not be nullptr."));
|
||||
|
||||
PADDLE_ENFORCE_LT(
|
||||
pos, nb_inputs + nb_outputs,
|
||||
platform::errors::InvalidArgument("The pos(%d) should be less than the "
|
||||
"num(%d) of the input and the output.",
|
||||
pos, nb_inputs + nb_outputs));
|
||||
|
||||
const nvinfer1::PluginTensorDesc &in = in_out[pos];
|
||||
if (pos == 0) {
|
||||
#ifdef SUPPORTS_CUDA_FP16
|
||||
if (ban_fp16_) {
|
||||
return (in.type == nvinfer1::DataType::kFLOAT) &&
|
||||
(in.format == nvinfer1::TensorFormat::kLINEAR);
|
||||
} else {
|
||||
return (in.type == nvinfer1::DataType::kFLOAT ||
|
||||
in.type == nvinfer1::DataType::kHALF) &&
|
||||
(in.format == nvinfer1::TensorFormat::kLINEAR);
|
||||
}
|
||||
#else
|
||||
return (in.type == nvinfer1::DataType::kFLOAT) &&
|
||||
(in.format == nvinfer1::TensorFormat::kLINEAR);
|
||||
#endif
|
||||
}
|
||||
const nvinfer1::PluginTensorDesc &prev = in_out[pos - 1];
|
||||
// output
|
||||
return in.type == prev.type && in.format == prev.format;
|
||||
}
|
||||
|
||||
nvinfer1::DataType SlicePluginDynamic::getOutputDataType(
|
||||
int index, const nvinfer1::DataType *input_types, int nb_inputs) const {
|
||||
PADDLE_ENFORCE_EQ(index, 0, platform::errors::InvalidArgument(
|
||||
"The Slice Plugin only has one input, so the "
|
||||
"index value should be 0, but get %d.",
|
||||
index));
|
||||
PADDLE_ENFORCE_EQ((input_types[0] == nvinfer1::DataType::kFLOAT ||
|
||||
input_types[0] == nvinfer1::DataType::kHALF),
|
||||
true, platform::errors::InvalidArgument(
|
||||
"The input type should be half or float"));
|
||||
return input_types[0];
|
||||
}
|
||||
|
||||
int SlicePluginDynamic::enqueue(const nvinfer1::PluginTensorDesc *input_desc,
|
||||
const nvinfer1::PluginTensorDesc *output_desc,
|
||||
const void *const *inputs, void *const *outputs,
|
||||
void *workspace, cudaStream_t stream) {
|
||||
auto input_dims = input_desc[0].dims;
|
||||
auto out_dims = output_desc[0].dims;
|
||||
auto num_dims = input_dims.nbDims;
|
||||
size_t out_num = ProductDim(out_dims);
|
||||
|
||||
std::vector<int> seg_offsets;
|
||||
std::vector<int> offsets;
|
||||
std::vector<int> extends;
|
||||
|
||||
offsets.reserve(num_dims);
|
||||
extends.reserve(num_dims);
|
||||
seg_offsets.reserve(num_dims);
|
||||
|
||||
seg_offsets[num_dims - 1] = 1;
|
||||
for (int i = num_dims - 2; i >= 0; i--) {
|
||||
seg_offsets[i] = input_dims.d[i + 1] * seg_offsets[i + 1];
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < num_dims; ++i) {
|
||||
offsets[i] = 0;
|
||||
extends[i] = out_dims.d[i];
|
||||
}
|
||||
for (size_t i = 0; i < axes_.size(); ++i) {
|
||||
offsets[axes_[i]] = starts_[i];
|
||||
}
|
||||
|
||||
std::vector<int> offset_info;
|
||||
for (size_t i = 0; i < num_dims; ++i) {
|
||||
offset_info.push_back(offsets[i]);
|
||||
offset_info.push_back(extends[i]);
|
||||
offset_info.push_back(seg_offsets[i]);
|
||||
}
|
||||
|
||||
framework::Tensor offset_temp_tensor;
|
||||
|
||||
int device_id;
|
||||
cudaGetDevice(&device_id);
|
||||
offset_temp_tensor.Resize({3 * num_dims});
|
||||
auto *offset_temp_data =
|
||||
offset_temp_tensor.mutable_data<int>(platform::CUDAPlace(device_id));
|
||||
|
||||
cudaMemcpyAsync(offset_temp_data, offset_info.data(),
|
||||
sizeof(int) * 3 * num_dims, cudaMemcpyHostToDevice, stream);
|
||||
|
||||
int threads = 256;
|
||||
int blocks = (out_num + threads - 1) / threads;
|
||||
auto input_type = input_desc[0].type;
|
||||
if (input_type == nvinfer1::DataType::kFLOAT) {
|
||||
const float *input1 = static_cast<const float *>(inputs[0]);
|
||||
float *output = static_cast<float *>(outputs[0]);
|
||||
SliceKernel<float><<<blocks, threads, 3 * num_dims * sizeof(int), stream>>>(
|
||||
out_num, num_dims, input1, offset_temp_data, output);
|
||||
} else if (input_type == nvinfer1::DataType::kHALF) {
|
||||
#ifdef SUPPORTS_CUDA_FP16
|
||||
const half *input1 = static_cast<const half *>(inputs[0]);
|
||||
half *output = static_cast<half *>(outputs[0]);
|
||||
SliceKernel<half><<<blocks, threads, 3 * num_dims * sizeof(int), stream>>>(
|
||||
out_num, num_dims, input1, offset_temp_data, output);
|
||||
#else
|
||||
PADDLE_THROW(platform::errors::Fatal(
|
||||
"The cuda archs you specific should greater than 600."));
|
||||
#endif
|
||||
} else {
|
||||
PADDLE_THROW(platform::errors::Fatal(
|
||||
"The Slice TRT Plugin's input type should be float or half."));
|
||||
}
|
||||
return cudaGetLastError() != cudaSuccess;
|
||||
}
|
||||
#endif
|
||||
|
||||
} // namespace plugin
|
||||
} // namespace tensorrt
|
||||
} // namespace inference
|
||||
} // namespace paddle
|
||||
@ -0,0 +1,89 @@
|
||||
// Copyright (c) 2018 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 <algorithm>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/fluid/inference/tensorrt/engine.h"
|
||||
#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
namespace plugin {
|
||||
|
||||
#if IS_TRT_VERSION_GE(6000)
|
||||
class SlicePluginDynamic : public DynamicPluginTensorRT {
|
||||
public:
|
||||
explicit SlicePluginDynamic(std::vector<int> starts, std::vector<int> ends,
|
||||
std::vector<int> axes, bool ban_fp16)
|
||||
: starts_(starts), ends_(ends), axes_(axes), ban_fp16_(ban_fp16) {}
|
||||
SlicePluginDynamic(void const* serialData, size_t serialLength) {}
|
||||
nvinfer1::IPluginV2DynamicExt* clone() const override {
|
||||
return new SlicePluginDynamic(starts_, ends_, axes_, ban_fp16_);
|
||||
}
|
||||
|
||||
const char* getPluginType() const override { return "slice_plugin"; }
|
||||
int getNbOutputs() const override { return 1; }
|
||||
int initialize() override;
|
||||
|
||||
size_t getSerializationSize() const override;
|
||||
void serialize(void* buffer) const override;
|
||||
|
||||
nvinfer1::DimsExprs getOutputDimensions(
|
||||
int output_index, const nvinfer1::DimsExprs* inputs, int nb_inputs,
|
||||
nvinfer1::IExprBuilder& expr_builder) override;
|
||||
|
||||
bool supportsFormatCombination(int pos,
|
||||
const nvinfer1::PluginTensorDesc* inOut,
|
||||
int nbInputs, int nbOutputs) override;
|
||||
|
||||
void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
|
||||
int nbInputs,
|
||||
const nvinfer1::DynamicPluginTensorDesc* out,
|
||||
int nbOutputs) override {}
|
||||
|
||||
size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
|
||||
int nbInputs,
|
||||
const nvinfer1::PluginTensorDesc* outputs,
|
||||
int nbOutputs) const override {
|
||||
return 0;
|
||||
}
|
||||
|
||||
int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
|
||||
const nvinfer1::PluginTensorDesc* outputDesc,
|
||||
const void* const* inputs, void* const* outputs, void* workspace,
|
||||
cudaStream_t stream) override;
|
||||
nvinfer1::DataType getOutputDataType(int index,
|
||||
const nvinfer1::DataType* inputTypes,
|
||||
int nbInputs) const override;
|
||||
|
||||
void destroy() override { delete this; }
|
||||
|
||||
private:
|
||||
std::vector<int> starts_;
|
||||
std::vector<int> ends_;
|
||||
std::vector<int> axes_;
|
||||
|
||||
bool ban_fp16_{false};
|
||||
};
|
||||
#endif
|
||||
|
||||
} // namespace plugin
|
||||
} // namespace tensorrt
|
||||
} // namespace inference
|
||||
} // namespace paddle
|
||||
Loading…
Reference in new issue