Skip to content

Commit 90c5194

Browse files
inocsininocsin
committed
refactor: update functions in slice converter
Signed-off-by: inocsin <[email protected]>
1 parent 0ccee81 commit 90c5194

File tree

3 files changed

+59
-35
lines changed

3 files changed

+59
-35
lines changed

‎core/conversion/converters/converter_util.cpp

Lines changed: 42 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -204,104 +204,119 @@ nvinfer1::ITensor* clamp(
204204
ConversionCtx* ctx,
205205
nvinfer1::ITensor* x,
206206
nvinfer1::ITensor* lower_bound,
207-
nvinfer1::ITensor* upper_bound) {
208-
auto max_layer = ctx->net->addElementWise(*x, *lower_bound, nvinfer1::ElementWiseOperation::kMAX);
207+
nvinfer1::ITensor* upper_bound,
208+
std::string const& name) {
209+
210+
auto max_layer = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMAX, x, lower_bound, "max layer for " + name);
209211
TORCHTRT_CHECK(max_layer, "Unable to create max layer for clamp");
210212
LOG_DEBUG(ctx->logger, "Create " << max_layer->getName() << " for clamp");
211213
auto max_itensor = max_layer->getOutput(0);
212214

213-
auto min_layer = ctx->net->addElementWise(*max_itensor, *upper_bound, nvinfer1::ElementWiseOperation::kMIN);
215+
auto min_layer = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMIN, max_itensor, upper_bound, "min layer for " + name);
214216
TORCHTRT_CHECK(min_layer, "Unable to create min layer for clamp");
215217
LOG_DEBUG(ctx->logger, "Create " << min_layer->getName() << " for clamp");
216218
auto min_itensor = min_layer->getOutput(0);
217219
return min_itensor;
218220
}
219221

220222
// clamp x to [0, input_dim]
221-
nvinfer1::ITensor* clamp_to_input_dim(ConversionCtx* ctx, nvinfer1::ITensor* x, nvinfer1::ITensor* input_dim, int nbdims) {
222-
// auto nbdims = input_dim->getDimensions().d[0];
223+
nvinfer1::ITensor* clamp_to_input_dim(
224+
ConversionCtx* ctx,
225+
nvinfer1::ITensor* x,
226+
nvinfer1::ITensor* input_dim,
227+
int nbdims,
228+
std::string const& name) {
229+
223230
auto zero = torch::zeros({nbdims}).to(torch::kI32);
224231
auto zero_itensor = tensor_to_const(ctx, zero);
225232
auto one = torch::ones({nbdims}).to(torch::kI32);
226233
auto one_itensor = tensor_to_const(ctx, one);
227234

228-
auto upper_bound_layer = ctx->net->addElementWise(*input_dim, *one_itensor, nvinfer1::ElementWiseOperation::kSUB);
235+
auto upper_bound_layer = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, input_dim, one_itensor, "sub layer for " + name);
229236
TORCHTRT_CHECK(upper_bound_layer, "Unable to create sub layer for clamp to inputDim");
230237
LOG_DEBUG(ctx->logger, "Create " << upper_bound_layer->getName() << " for clamp to inputDim");
231238
auto upper_bound = upper_bound_layer->getOutput(0);
232239

233-
auto max_layer = ctx->net->addElementWise(*x, *zero_itensor, nvinfer1::ElementWiseOperation::kMAX);
240+
auto max_layer = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMAX, x, zero_itensor, "max layer for " + name);
234241
TORCHTRT_CHECK(max_layer, "Unable to create max_layer for clamp to inputDim");
235242
LOG_DEBUG(ctx->logger, "Create " << max_layer->getName() << " for clamp to inputDim");
236243
auto max_itensor = max_layer->getOutput(0);
237244

238-
auto min_layer = ctx->net->addElementWise(*max_itensor, *upper_bound, nvinfer1::ElementWiseOperation::kMIN);
245+
auto min_layer = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMIN, max_itensor, upper_bound, "min layer for " + name);
239246
TORCHTRT_CHECK(min_layer, "Unable to create min_layer for clamp to inputDim");
240247
LOG_DEBUG(ctx->logger, "Create " << min_layer->getName() << " for clamp to inputDim");
241248
auto min_itensor = min_layer->getOutput(0);
242249
return min_itensor;
243250
}
244251

245252
// return indices < 0 ? inputDims + indices : indices
246-
nvinfer1::ITensor* bump_if_negative(ConversionCtx* ctx, nvinfer1::ITensor* input_dim, nvinfer1::ITensor* indices, int nbdims) {
253+
nvinfer1::ITensor* normalize_indices(
254+
ConversionCtx* ctx,
255+
nvinfer1::ITensor* input_dim,
256+
nvinfer1::ITensor* indices,
257+
int nbdims,
258+
std::string const& name) {
259+
247260
auto zero = torch::zeros({nbdims}).to(torch::kI32);
248261
auto neg = -torch::ones({nbdims}).to(torch::kI32);
249262
auto zero_itensor = tensor_to_const(ctx, zero);
250263
auto neg_itensor = tensor_to_const(ctx, neg);
251264
// find the indices that = -1
252-
auto signs = clamp(ctx, indices, neg_itensor, zero_itensor);
265+
auto signs = clamp(ctx, indices, neg_itensor, zero_itensor, "clamp layer for " + name);
253266

254267
// get the inputDim value where indices == -1, else 0
255-
auto mul = ctx->net->addElementWise(*signs, *input_dim, nvinfer1::ElementWiseOperation::kPROD);
256-
TORCHTRT_CHECK(mul, "Unable to create mul layer in bump_if_negative");
257-
LOG_DEBUG(ctx->logger, "Create " << mul->getName() << " for bump_if_negative");
268+
auto mul = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, signs, input_dim, "prod layer for " + name);
269+
TORCHTRT_CHECK(mul, "Unable to create mul layer in normalize_indices");
270+
LOG_DEBUG(ctx->logger, "Create " << mul->getName() << " for normalize_indices");
258271
auto mul_itensor = mul->getOutput(0);
259272

260273
// add the inputDim value to indices where indices == -1
261-
auto sub = ctx->net->addElementWise(*indices, *mul_itensor, nvinfer1::ElementWiseOperation::kSUB);
262-
TORCHTRT_CHECK(sub, "Unable to create sub layer in bump_if_negative");
263-
LOG_DEBUG(ctx->logger, "Create " << sub->getName() << " for bump_if_negative");
274+
auto sub = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, indices, mul_itensor, "sub layer for " + name);
275+
TORCHTRT_CHECK(sub, "Unable to create sub layer in normalize_indices");
276+
LOG_DEBUG(ctx->logger, "Create " << sub->getName() << " for normalize_indices");
264277
auto sub_itensor = sub->getOutput(0);
265278
return sub_itensor;
266279
}
267280

268-
std::vector<nvinfer1::ITensor*> update_start_and_end(
281+
std::vector<nvinfer1::ITensor*> normalize_start_and_end(
269282
ConversionCtx* ctx,
270283
nvinfer1::ITensor* in_shape,
271284
nvinfer1::ITensor* in_start,
272285
nvinfer1::ITensor* in_end,
273-
int nbdims) {
274-
auto start = bump_if_negative(ctx, in_shape, in_start, nbdims);
275-
auto out_start = clamp_to_input_dim(ctx, start, in_shape, nbdims);
276-
auto end = bump_if_negative(ctx, in_shape, in_end, nbdims);
277-
auto out_end = clamp_to_input_dim(ctx, end, in_shape, nbdims);
286+
int nbdims,
287+
std::string const& name) {
288+
auto start = normalize_indices(ctx, in_shape, in_start, nbdims, "normalize start of " + name);
289+
auto out_start = clamp_to_input_dim(ctx, start, in_shape, nbdims, "clamp start to inputDim for " + name);
290+
auto end = normalize_indices(ctx, in_shape, in_end, nbdims, "normalize end of " + name);
291+
auto out_end = clamp_to_input_dim(ctx, end, in_shape, nbdims, "clamp end to inputDim for " + name);
278292
std::vector<nvinfer1::ITensor*> outputs;
279293
outputs.push_back(out_start);
280294
outputs.push_back(out_end);
281295
return outputs;
282296
}
283297

284298
// size = (end - start) / stride + 1, where range is [start, end], end is included
285-
nvinfer1::ITensor* calculate_output_size(
299+
nvinfer1::ITensor* get_slice_size(
286300
ConversionCtx* ctx,
287301
nvinfer1::ITensor* start,
288302
nvinfer1::ITensor* end,
289303
nvinfer1::ITensor* stride,
290-
int nbdims) {
304+
int nbdims,
305+
std::string const& name) {
291306
at::Tensor one_tensor = torch::ones({nbdims}).to(torch::kI32);
292307
auto one_itensor = tensor_to_const(ctx, one_tensor);
293308

294-
auto sub_layer = ctx->net->addElementWise(*end, *start, nvinfer1::ElementWiseOperation::kSUB);
309+
auto sub_layer = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, end, start, "get_slice_size sub layer for " + name);
295310
TORCHTRT_CHECK(sub_layer, "Unable to create sub layer in calculate_output_size");
296311
LOG_DEBUG(ctx->logger, "Create " << sub_layer->getName() << " for calculate_output_size");
297312
auto sub_itensor = sub_layer->getOutput(0);
298313

299-
auto div_layer = ctx->net->addElementWise(*sub_itensor, *stride, nvinfer1::ElementWiseOperation::kDIV);
314+
auto div_layer = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, sub_itensor, stride, "get_slice_size div layer for " + name);
300315
TORCHTRT_CHECK(div_layer, "Unable to create div layer in calculate_output_size");
301316
LOG_DEBUG(ctx->logger, "Create " << div_layer->getName() << " for calculate_output_size");
302317
auto div_itensor = div_layer->getOutput(0);
303318

304-
auto add_layer = ctx->net->addElementWise(*div_itensor, *one_itensor, nvinfer1::ElementWiseOperation::kSUM);
319+
auto add_layer = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, div_itensor, one_itensor, "get_slice_size sum layer for " + name);
305320
TORCHTRT_CHECK(add_layer, "Unable to create add layer in calculate_output_size");
306321
LOG_DEBUG(ctx->logger, "Create " << add_layer->getName() << " for calculate_output_size");
307322
auto size_itensor = add_layer->getOutput(0);

‎core/conversion/converters/converter_util.h

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -55,23 +55,30 @@ nvinfer1::ITensor* clamp(
5555
ConversionCtx* ctx,
5656
nvinfer1::ITensor* x,
5757
nvinfer1::ITensor* lower_bound,
58-
nvinfer1::ITensor* upper_bound);
58+
nvinfer1::ITensor* upper_bound,
59+
std::string const& name);
5960

60-
nvinfer1::ITensor* bump_if_negative(ConversionCtx* ctx, nvinfer1::ITensor* input_dim, nvinfer1::ITensor* indices);
61+
nvinfer1::ITensor* normalize_indices(
62+
ConversionCtx* ctx,
63+
nvinfer1::ITensor* input_dim,
64+
nvinfer1::ITensor* indices,
65+
std::string const& name);
6166

62-
std::vector<nvinfer1::ITensor*> update_start_and_end(
67+
std::vector<nvinfer1::ITensor*> normalize_start_and_end(
6368
ConversionCtx* ctx,
6469
nvinfer1::ITensor* in_shape,
6570
nvinfer1::ITensor* in_start,
6671
nvinfer1::ITensor* in_end,
67-
int nbdims);
72+
int nbdims,
73+
std::string const& name);
6874

69-
nvinfer1::ITensor* calculate_output_size(
75+
nvinfer1::ITensor* get_slice_size(
7076
ConversionCtx* ctx,
7177
nvinfer1::ITensor* start,
7278
nvinfer1::ITensor* end,
7379
nvinfer1::ITensor* stride,
74-
int nbdims);
80+
int nbdims,
81+
std::string const& name);
7582

7683
} // namespace converters
7784
} // namespace conversion

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -314,6 +314,7 @@ auto select_registrations TORCHTRT_UNUSED =
314314
// add Shape Tensor
315315
auto ishape_layer = ctx->net->addShape(*in);
316316
auto ishape_tensor = ishape_layer->getOutput(0); // input shape
317+
std::string node_name = n->outputs()[0]->debugName().c_str();
317318

318319
int startIdx = 0;
319320
auto startIdxIVal = args[2].IValue();
@@ -384,18 +385,19 @@ auto select_registrations TORCHTRT_UNUSED =
384385
// update start and end
385386
nvinfer1::ITensor* out_start;
386387
nvinfer1::ITensor* out_end;
387-
auto start_end = update_start_and_end(ctx, ishape_tensor, start_itensor, end_itensor, nbdims);
388+
auto start_end = normalize_start_and_end(ctx, ishape_tensor, start_itensor, end_itensor, nbdims, node_name);
388389
out_start = start_end[0];
389390
out_end = start_end[1];
390391

391392
// calculate size
392-
auto size_itensor = calculate_output_size(ctx, out_start, out_end, stride_itensor, nbdims);
393+
auto size_itensor = get_slice_size(ctx, out_start, out_end, stride_itensor, nbdims, node_name);
393394

394395
// update slice layer
395396
slice_layer->setInput(1, *out_start); // start
396397
slice_layer->setInput(2, *size_itensor); // size, must be set if input is dynamic
397398
}
398399
auto slice_out = slice_layer->getOutput(0);
400+
399401
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], slice_out);
400402
LOG_DEBUG("Slice layer output shape: " << out->getDimensions());
401403

0 commit comments

Comments
 (0)