Skip to content

Commit 4a9c595

Browse files
peri044peri044
committed
modify gelu plugin calls
Signed-off-by: Dheeraj Peri <[email protected]>
1 parent 4198e19 commit 4a9c595

File tree

11 files changed

+170
-167
lines changed

11 files changed

+170
-167
lines changed

‎BUILD

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ pkg_tar(
1818
"//core/conversion/var:include",
1919
"//core/conversion/tensorcontainer:include",
2020
"//core/conversion/evaluators:include",
21-
"//core/plugins:trtorch_plugins",
21+
"//core/plugins:include",
2222
"//core/runtime:include",
2323
"//core/lowering:include",
2424
"//core/lowering/passes:include",

‎core/conversion/conversionctx/ConversionCtx.h

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -68,9 +68,6 @@ 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-
7471
std::unordered_map<const torch::jit::Value*, nvinfer1::ITensor*> value_tensor_map;
7572
std::unordered_map<const torch::jit::Value*, torch::jit::IValue> evaluated_value_map;
7673
};

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

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

4343
auto acthardtanh TRTORCH_UNUSED =
4444
RegisterNodeConversionPatterns()
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();
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();
5150

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

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

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

61-
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
62-
return true;
63-
}})
60+
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
61+
return true;
62+
}})
6463
.pattern({// TODO: Remove after functionalization
6564
"aten::hardtanh_(Tensor(a!) self, Scalar min_val=-1, Scalar max_val=1) -> (Tensor(a!))",
6665
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
@@ -126,35 +125,33 @@ auto acthardtanh TRTORCH_UNUSED =
126125
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
127126
return true;
128127
}})
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>();
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>();
134132

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

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>();
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>();
149146

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-
}})
147+
auto new_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kLEAKY_RELU);
148+
new_layer->setAlpha(negative_slopeScalar);
149+
new_layer->setName(util::node_info(n).c_str());
150+
auto out_tensor = new_layer->getOutput(0);
151+
out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
152+
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
153+
return true;
154+
}})
158155
.pattern({"aten::elu(Tensor self, Scalar alpha=1, Scalar scale=1, Scalar input_scale=1) -> (Tensor)",
159156
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
160157
auto in = args[0].ITensorOrFreeze(ctx);
@@ -169,33 +166,36 @@ auto acthardtanh TRTORCH_UNUSED =
169166
auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
170167
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
171168
return true;
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-
}});
169+
}})
170+
.pattern({"aten::gelu(Tensor self) -> (Tensor)",
171+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
172+
auto in = args[0].ITensorOrFreeze(ctx);
173+
nvinfer1::DataType type = in->getType();
174+
TRTORCH_CHECK(
175+
type == nvinfer1::DataType::kFLOAT || type == nvinfer1::DataType::kHALF,
176+
"gelu only supports kFLOAT and kHALF");
177+
std::string pluginName = "CustomGeluPluginDynamic";
178+
nvinfer1::PluginFieldCollection fc;
179+
std::vector<nvinfer1::PluginField> f;
180+
int type_id = ctx->settings.op_precision == nvinfer1::DataType::kFLOAT
181+
? 0
182+
: 1; // Integer encoding the DataType (0: FP32, 1: FP16)
183+
f.emplace_back(nvinfer1::PluginField("type_id", &type_id, nvinfer1::PluginFieldType::kINT32, 1));
184+
fc.nbFields = f.size();
185+
fc.fields = f.data();
186+
187+
auto creator = getPluginRegistry()->getPluginCreator("CustomGeluPluginDynamic", "1", "");
188+
auto gelu_plugin = creator->createPlugin("gelu", &fc);
189+
190+
TRTORCH_CHECK(gelu_plugin, "Unable to create gelu plugin from TensorRT plugin registry" << *n);
191+
auto new_layer =
192+
ctx->net->addPluginV2(reinterpret_cast<nvinfer1::ITensor* const*>(&in), 1, *gelu_plugin);
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+
}});
199199

200200
} // namespace
201201
} // namespace impl

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

Lines changed: 55 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -17,75 +17,73 @@ static inline at::Tensor repeat_if_defined(const at::Tensor& t, int64_t repeat)
1717
return t;
1818
}
1919

20-
auto batch_norm_registrations TRTORCH_UNUSED =
21-
RegisterNodeConversionPatterns()
22-
.pattern({
23-
R"SIG(aten::batch_norm(Tensor input, Tensor? gamma, Tensor? beta,
20+
auto batch_norm_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().pattern({
21+
R"SIG(aten::batch_norm(Tensor input, Tensor? gamma, Tensor? beta,
2422
Tensor? mean, Tensor? var,
2523
bool training, float momentum, float eps, bool cudnn_enabled) -> (Tensor))SIG",
26-
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
27-
auto input = args[0].ITensor(); // assumes non-static input Tensor
28-
auto orig_shape = input->getDimensions();
29-
auto shape = util::toVec(orig_shape);
30-
auto tensor_type = util::toATenDType(input->getType());
31-
auto options = torch::TensorOptions().dtype(tensor_type);
24+
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
25+
auto input = args[0].ITensor(); // assumes non-static input Tensor
26+
auto orig_shape = input->getDimensions();
27+
auto shape = util::toVec(orig_shape);
28+
auto tensor_type = util::toATenDType(input->getType());
29+
auto options = torch::TensorOptions().dtype(tensor_type);
3230

33-
torch::Tensor gamma, beta, mean, var;
31+
torch::Tensor gamma, beta, mean, var;
3432

35-
if (ctx->input_is_dynamic) {
36-
gamma = args[1].unwrapToTensor();
37-
beta = args[2].unwrapToTensor();
38-
mean = args[3].unwrapToTensor();
39-
var = args[4].unwrapToTensor();
40-
} else {
41-
gamma = args[1].unwrapToTensor(at::full({shape}, 1, {options}));
42-
beta = args[2].unwrapToTensor(at::full({shape}, 1, {options}));
43-
mean = args[3].unwrapToTensor(at::full({shape}, 0, {options}));
44-
var = args[4].unwrapToTensor(at::full({shape}, 0, {options}));
45-
}
33+
if (ctx->input_is_dynamic) {
34+
gamma = args[1].unwrapToTensor();
35+
beta = args[2].unwrapToTensor();
36+
mean = args[3].unwrapToTensor();
37+
var = args[4].unwrapToTensor();
38+
} else {
39+
gamma = args[1].unwrapToTensor(at::full({shape}, 1, {options}));
40+
beta = args[2].unwrapToTensor(at::full({shape}, 1, {options}));
41+
mean = args[3].unwrapToTensor(at::full({shape}, 0, {options}));
42+
var = args[4].unwrapToTensor(at::full({shape}, 0, {options}));
43+
}
4644

47-
auto eps = args[7].unwrapToDouble(1e-5f);
45+
auto eps = args[7].unwrapToDouble(1e-5f);
4846

49-
LOG_DEBUG("momentum disregarded");
50-
LOG_DEBUG("training disregarded");
51-
LOG_DEBUG("cudnn disregarded");
47+
LOG_DEBUG("momentum disregarded");
48+
LOG_DEBUG("training disregarded");
49+
LOG_DEBUG("cudnn disregarded");
5250

53-
auto should_unpack = util::toVec(orig_shape).size() < 4;
54-
if (should_unpack) {
55-
// expand spatial dims from 1D to 2D
56-
auto new_shape = util::toDimsTailPad(util::toVec(orig_shape), 4);
57-
LOG_DEBUG(
58-
"Input shape is less than 4D got: "
59-
<< orig_shape << ", inserting shuffle layer to reshape to 4D tensor shape: " << new_shape);
60-
auto in_shuffle = ctx->net->addShuffle(*input);
61-
in_shuffle->setReshapeDimensions(new_shape);
62-
in_shuffle->setName(std::string("[Reshape input to " + util::toStr(new_shape) + ']').c_str());
63-
input = in_shuffle->getOutput(0);
64-
}
51+
auto should_unpack = util::toVec(orig_shape).size() < 4;
52+
if (should_unpack) {
53+
// expand spatial dims from 1D to 2D
54+
auto new_shape = util::toDimsTailPad(util::toVec(orig_shape), 4);
55+
LOG_DEBUG(
56+
"Input shape is less than 4D got: "
57+
<< orig_shape << ", inserting shuffle layer to reshape to 4D tensor shape: " << new_shape);
58+
auto in_shuffle = ctx->net->addShuffle(*input);
59+
in_shuffle->setReshapeDimensions(new_shape);
60+
in_shuffle->setName(std::string("[Reshape input to " + util::toStr(new_shape) + ']').c_str());
61+
input = in_shuffle->getOutput(0);
62+
}
6563

66-
auto scale = gamma / torch::sqrt(var + eps);
67-
auto bias = beta - mean * scale;
64+
auto scale = gamma / torch::sqrt(var + eps);
65+
auto bias = beta - mean * scale;
6866

69-
auto scale_weights = Weights(ctx, scale);
70-
auto bias_weights = Weights(ctx, bias);
67+
auto scale_weights = Weights(ctx, scale);
68+
auto bias_weights = Weights(ctx, bias);
7169

72-
auto power = Weights(ctx, at::ones_like(scale));
73-
auto bn = ctx->net->addScaleNd(
74-
*input, nvinfer1::ScaleMode::kCHANNEL, bias_weights.data, scale_weights.data, power.data, 1);
75-
bn->setName(util::node_info(n).c_str());
76-
auto out_tensor = bn->getOutput(0);
70+
auto power = Weights(ctx, at::ones_like(scale));
71+
auto bn = ctx->net->addScaleNd(
72+
*input, nvinfer1::ScaleMode::kCHANNEL, bias_weights.data, scale_weights.data, power.data, 1);
73+
bn->setName(util::node_info(n).c_str());
74+
auto out_tensor = bn->getOutput(0);
7775

78-
if (should_unpack) {
79-
LOG_DEBUG("Inserting shuffle layer to reshape to back to original shape: " << orig_shape);
80-
auto out_shuffle = ctx->net->addShuffle(*out_tensor);
81-
out_shuffle->setReshapeDimensions(orig_shape);
82-
out_shuffle->setName(std::string("[Reshape output to " + util::toStr(orig_shape) + ']').c_str());
83-
out_tensor = out_shuffle->getOutput(0);
84-
}
76+
if (should_unpack) {
77+
LOG_DEBUG("Inserting shuffle layer to reshape to back to original shape: " << orig_shape);
78+
auto out_shuffle = ctx->net->addShuffle(*out_tensor);
79+
out_shuffle->setReshapeDimensions(orig_shape);
80+
out_shuffle->setName(std::string("[Reshape output to " + util::toStr(orig_shape) + ']').c_str());
81+
out_tensor = out_shuffle->getOutput(0);
82+
}
8583

86-
ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
87-
return true;
88-
}});
84+
ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
85+
return true;
86+
}});
8987

9088
} // namespace
9189
} // namespace impl

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
#include "NvInfer.h"
22
#include "NvInferRuntimeCommon.h"
33
#include "core/conversion/converters/converters.h"
4-
#include "core/plugins/impl/interpolate_plugin.h"
54
#include "core/util/prelude.h"
65
#include "torch/torch.h"
76

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

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -320,21 +320,25 @@ auto pooling_registrations TRTORCH_UNUSED =
320320
std::vector<nvinfer1::PluginField> f;
321321

322322
std::vector<int32_t> in_shape_casted(in_shape.begin(), in_shape.end());
323-
f.emplace_back(nvinfer1::PluginField("in_shape", in_shape_casted.data(), nvinfer1::PluginFieldType::kINT32, in_shape.size()));
323+
f.emplace_back(nvinfer1::PluginField(
324+
"in_shape", in_shape_casted.data(), nvinfer1::PluginFieldType::kINT32, in_shape.size()));
324325

325326
std::vector<int32_t> out_shape_casted(out_shape.begin(), out_shape.end());
326-
f.emplace_back(nvinfer1::PluginField("out_shape", out_shape_casted.data(), nvinfer1::PluginFieldType::kINT32, out_shape.size()));
327+
f.emplace_back(nvinfer1::PluginField(
328+
"out_shape", out_shape_casted.data(), nvinfer1::PluginFieldType::kINT32, out_shape.size()));
327329

328330
std::vector<int32_t> out_size_casted(out_size.begin(), out_size.end());
329-
f.emplace_back(nvinfer1::PluginField("out_size", out_size_casted.data(), nvinfer1::PluginFieldType::kINT32, out_size.size()));
331+
f.emplace_back(nvinfer1::PluginField(
332+
"out_size", out_size_casted.data(), nvinfer1::PluginFieldType::kINT32, out_size.size()));
330333

331334
f.emplace_back(nvinfer1::PluginField("scales", nullptr, nvinfer1::PluginFieldType::kFLOAT64, 0));
332335

333336
std::string name = "adaptive_pool2d";
334-
f.emplace_back(nvinfer1::PluginField("mode", &name, nvinfer1::PluginFieldType::kCHAR , 1));
337+
f.emplace_back(nvinfer1::PluginField("mode", &name, nvinfer1::PluginFieldType::kCHAR, 1));
335338

336339
int32_t align_corners = 0;
337-
f.emplace_back(nvinfer1::PluginField("align_corners", &align_corners, nvinfer1::PluginFieldType::kINT32, 1));
340+
f.emplace_back(
341+
nvinfer1::PluginField("align_corners", &align_corners, nvinfer1::PluginFieldType::kINT32, 1));
338342

339343
int32_t use_scales = 0;
340344
f.emplace_back(nvinfer1::PluginField("use_scales", &use_scales, nvinfer1::PluginFieldType::kINT32, 1));

0 commit comments

Comments
 (0)