Skip to content

Commit 4198e19

Browse files
ruoqianguoperi044
ruoqianguo
authored and
peri044
committed
adding support for gelu converter
Signed-off-by: Ruoqian Guo <[email protected]>
1 parent c53ccac commit 4198e19

File tree

5 files changed

+101
-41
lines changed

5 files changed

+101
-41
lines changed

‎core/conversion/conversionctx/BUILD

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ cc_library(
1717
],
1818
deps = [
1919
"@tensorrt//:nvinfer",
20+
"@tensorrt//:nvinferplugin",
2021
"//core/util:prelude",
2122
#"//core/plugins:trtorch_plugins",
2223
] + select({

‎core/conversion/conversionctx/ConversionCtx.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
#include <iostream>
22
#include <sstream>
3+
4+
#include "NvInferPlugin.h"
5+
#include "NvInferPluginUtils.h"
6+
37
#include <utility>
48
// #include "core/plugins/plugin_prelude.h"
59
#include "core/conversion/conversionctx/ConversionCtx.h"

‎core/conversion/conversionctx/ConversionCtx.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,9 @@ struct ConversionCtx {
6868
// copy of the values
6969
std::vector<void*> builder_resources;
7070

71+
// Registry of official tensorrt plugin layers.
72+
std::unordered_map<std::string, nvinfer1::IPluginCreator*> mPluginRegistry;
73+
7174
std::unordered_map<const torch::jit::Value*, nvinfer1::ITensor*> value_tensor_map;
7275
std::unordered_map<const torch::jit::Value*, torch::jit::IValue> evaluated_value_map;
7376
};

‎core/conversion/converters/impl/activation.cpp

Lines changed: 68 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -42,24 +42,25 @@ convert(tanh, kTANH);
4242

4343
auto acthardtanh TRTORCH_UNUSED =
4444
RegisterNodeConversionPatterns()
45-
.pattern({"aten::hardtanh(Tensor self, Scalar min_val=-1, Scalar max_val=1) -> (Tensor)",
46-
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
47-
auto in = args[0].ITensorOrFreeze(ctx);
48-
auto min = args[1].unwrapToDouble();
49-
auto max = args[2].unwrapToDouble();
45+
.pattern(
46+
{"aten::hardtanh(Tensor self, Scalar min_val=-1, Scalar max_val=1) -> (Tensor)",
47+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
48+
auto in = args[0].ITensorOrFreeze(ctx);
49+
auto min = args[1].unwrapToDouble();
50+
auto max = args[2].unwrapToDouble();
5051

51-
auto new_layer = ctx->net->addActivation(*in, nvinfer1::ActivationType::kCLIP);
52-
TRTORCH_CHECK(new_layer, "Unable to create layer for aten::hardtanh");
52+
auto new_layer = ctx->net->addActivation(*in, nvinfer1::ActivationType::kCLIP);
53+
TRTORCH_CHECK(new_layer, "Unable to create layer for aten::hardtanh");
5354

54-
new_layer->setAlpha(min);
55-
new_layer->setBeta(max);
55+
new_layer->setAlpha(min);
56+
new_layer->setBeta(max);
5657

57-
new_layer->setName(util::node_info(n).c_str());
58-
auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
58+
new_layer->setName(util::node_info(n).c_str());
59+
auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
5960

60-
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
61-
return true;
62-
}})
61+
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
62+
return true;
63+
}})
6364
.pattern({// TODO: Remove after functionalization
6465
"aten::hardtanh_(Tensor(a!) self, Scalar min_val=-1, Scalar max_val=1) -> (Tensor(a!))",
6566
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
@@ -125,34 +126,35 @@ auto acthardtanh TRTORCH_UNUSED =
125126
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
126127
return true;
127128
}})
128-
.pattern({"aten::leaky_relu(Tensor self, Scalar negative_slope=0.01) -> (Tensor)",
129-
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
130-
auto self = args[0].ITensorOrFreeze(ctx);
131-
auto negative_slopeScalar = args[1].unwrapToScalar().to<float>();
129+
.pattern(
130+
{"aten::leaky_relu(Tensor self, Scalar negative_slope=0.01) -> (Tensor)",
131+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
132+
auto self = args[0].ITensorOrFreeze(ctx);
133+
auto negative_slopeScalar = args[1].unwrapToScalar().to<float>();
132134

133-
auto new_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kLEAKY_RELU);
134-
new_layer->setAlpha(negative_slopeScalar);
135+
auto new_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kLEAKY_RELU);
136+
new_layer->setAlpha(negative_slopeScalar);
135137

136-
new_layer->setName(util::node_info(n).c_str());
137-
auto out_tensor = new_layer->getOutput(0);
138-
out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
139-
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
140-
return true;
141-
}})
142-
.pattern({"aten::leaky_relu_(Tensor(a!) self, Scalar negative_slope=0.01) -> Tensor(a!)",
143-
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
144-
auto self = args[0].ITensorOrFreeze(ctx);
145-
auto negative_slopeScalar = args[1].unwrapToScalar().to<float>();
146-
147-
auto new_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kLEAKY_RELU);
148-
new_layer->setAlpha(negative_slopeScalar);
138+
new_layer->setName(util::node_info(n).c_str());
139+
auto out_tensor = new_layer->getOutput(0);
140+
out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
141+
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
142+
return true;
143+
}})
144+
.pattern(
145+
{"aten::leaky_relu_(Tensor(a!) self, Scalar negative_slope=0.01) -> Tensor(a!)",
146+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
147+
auto self = args[0].ITensorOrFreeze(ctx);
148+
auto negative_slopeScalar = args[1].unwrapToScalar().to<float>();
149149

150-
new_layer->setName(util::node_info(n).c_str());
151-
auto out_tensor = new_layer->getOutput(0);
152-
out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
153-
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
154-
return true;
155-
}})
150+
auto new_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kLEAKY_RELU);
151+
new_layer->setAlpha(negative_slopeScalar);
152+
new_layer->setName(util::node_info(n).c_str());
153+
auto out_tensor = new_layer->getOutput(0);
154+
out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
155+
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
156+
return true;
157+
}})
156158
.pattern({"aten::elu(Tensor self, Scalar alpha=1, Scalar scale=1, Scalar input_scale=1) -> (Tensor)",
157159
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
158160
auto in = args[0].ITensorOrFreeze(ctx);
@@ -167,7 +169,33 @@ auto acthardtanh TRTORCH_UNUSED =
167169
auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
168170
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
169171
return true;
170-
}});
172+
}})
173+
.pattern(
174+
{"aten::gelu(Tensor self) -> (Tensor)",
175+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
176+
auto in = args[0].ITensorOrFreeze(ctx);
177+
nvinfer1::DataType type = in->getType();
178+
TRTORCH_CHECK(
179+
type == nvinfer1::DataType::kFLOAT || type == nvinfer1::DataType::kHALF,
180+
"gelu only supports kFLOAT and kHALF");
181+
std::string pluginName = "CustomGeluPluginDynamic";
182+
nvinfer1::PluginFieldCollection fc;
183+
std::vector<nvinfer1::PluginField> f;
184+
int type_id = 0; // Integer encoding the DataType (0: FP32, 1: FP16)
185+
if (type == nvinfer1::DataType::kHALF)
186+
type_id = 1;
187+
f.emplace_back(nvinfer1::PluginField("type_id", &type_id, nvinfer1::PluginFieldType::kINT32, 1));
188+
fc.nbFields = f.size();
189+
fc.fields = f.data();
190+
nvinfer1::IPluginV2* pluginV2 = ctx->mPluginRegistry.at(pluginName)->createPlugin("gelu", &fc);
191+
TRTORCH_CHECK(pluginV2, "Unable to create gelu plugin from TensorRT plugin registry" << *n);
192+
auto new_layer = ctx->net->addPluginV2(reinterpret_cast<nvinfer1::ITensor* const*>(&in), 1, *pluginV2);
193+
new_layer->setName("gelu");
194+
auto out_tensor = new_layer->getOutput(0);
195+
out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
196+
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
197+
return true;
198+
}});
171199

172200
} // namespace
173201
} // namespace impl

‎tests/core/conversion/converters/test_activation.cpp

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -198,4 +198,28 @@ TEST(Converters, ATenEluConvertsCorrectly) {
198198
auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
199199

200200
ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
201-
}
201+
}
202+
203+
TEST(Converters, ATenGELUConvertsCorrectly) {
204+
const auto graph = R"IR(
205+
graph(%0 : Tensor):
206+
%3 : Tensor = aten::gelu(%0)
207+
return (%3))IR";
208+
209+
auto g = std::make_shared<torch::jit::Graph>();
210+
torch::jit::parseIR(graph, &*g);
211+
212+
auto in = at::randint(-5, 5, {5}, {at::kCUDA});
213+
214+
auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
215+
auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
216+
217+
in = at::clone(in);
218+
params = trtorch::core::conversion::get_named_params(g->inputs(), {});
219+
auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
220+
221+
// The official tensorrt plugin applies the Gelu activation x * Phi(x), where Phi is the Gaussian cdf, approximated
222+
// by: 0.5 * (1 + tanh(sqrt(2 / M_PI) * (x + 0.044715 * x^3))) and the pytorch uses c10::cuda::compat::normcdf to
223+
// compute Phi(x). So there's a difference here.
224+
ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-4));
225+
}

0 commit comments

Comments
 (0)