Skip to content

Commit 9920cfd

Browse files
peri044peri044
committed
chore: rebase with master
Signed-off-by: Dheeraj Peri <[email protected]>
2 parents d4f4717 + d8c9a0c commit 9920cfd

File tree

583 files changed

+176546
-7402
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

583 files changed

+176546
-7402
lines changed

‎.circleci/config.yml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -428,7 +428,7 @@ commands:
428428
- run:
429429
name: Test torch
430430
command: |
431-
python3 -c "import torch; print(torch.cuda.is_available()); print(torch.cuda.device_count())"
431+
python3 -c "import torch; print(torch.cuda.is_available()); print(torch.cuda.device_count()); print(torch.__version__)"
432432
433433
- run:
434434
name: Get torch-tensorrt version information
@@ -988,7 +988,7 @@ parameters:
988988
# Nightly platform config
989989
torch-build:
990990
type: string
991-
default: "1.13.0"
991+
default: "1.14.0.dev20221114+cu117"
992992
torch-build-index:
993993
type: string
994994
default: "https://download.pytorch.org/whl/nightly/cu117"

‎.pre-commit-config.yaml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ repos:
66
- id: check-yaml
77
- id: trailing-whitespace
88
- id: check-added-large-files
9+
args:
10+
- --maxkb=1000
911
- id: check-vcs-permalinks
1012
- id: check-merge-conflict
1113
- id: mixed-line-ending

‎README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# Torch-TensorRT
22

33
[![Documentation](https://img.shields.io/badge/docs-master-brightgreen)](https://nvidia.github.io/Torch-TensorRT/)
4+
[![CircleCI](https://circleci.com/gh/pytorch/TensorRT.svg?style=svg)](https://app.circleci.com/pipelines/github/pytorch/TensorRT)
45

56
> Ahead of Time (AOT) compiling for PyTorch JIT and FX
67

‎WORKSPACE

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -56,17 +56,17 @@ new_local_repository(
5656
http_archive(
5757
name = "libtorch",
5858
build_file = "@//third_party/libtorch:BUILD",
59-
sha256 = "0a013dceedb252f4965b666a2ad772d962135597db5889bd5d43644697c17dbc",
59+
sha256 = "b565c662435fd58ec295fa0791388ea52ad0f5fd33517b2d7c0fdcc91b6db531",
6060
strip_prefix = "libtorch",
61-
urls = ["https://download.pytorch.org/libtorch/cu117/libtorch-cxx11-abi-shared-with-deps-1.13.0%2Bcu117.zip"],
61+
urls = ["https://download.pytorch.org/libtorch/nightly/cu116/libtorch-cxx11-abi-shared-with-deps-1.14.0.dev20221114%2Bcu117.zip"],
6262
)
6363

6464
http_archive(
6565
name = "libtorch_pre_cxx11_abi",
6666
build_file = "@//third_party/libtorch:BUILD",
67-
sha256 = "cdbd43985ad9d5886793d5dc455d665cf3fd4b4617ef1094479678ff210ed0af",
67+
sha256 = "fbb37446c33b05c1e26256c09f6ffb46cea1f6ff9ee2ad5b79b146d09023b0c1",
6868
strip_prefix = "libtorch",
69-
urls = ["https://download.pytorch.org/libtorch/cu117/libtorch-shared-with-deps-1.13.0%2Bcu117.zip"],
69+
urls = ["https://download.pytorch.org/libtorch/nightly/cu116/libtorch-shared-with-deps-1.14.0.dev20221114%2Bcu117.zip"],
7070
)
7171

7272
# Download these tarballs manually from the NVIDIA website
@@ -76,20 +76,20 @@ http_archive(
7676
http_archive(
7777
name = "cudnn",
7878
build_file = "@//third_party/cudnn/archive:BUILD",
79-
sha256 = "ec96d2376d81fca42bdd3d4c3d705a99b29a065bab57f920561c763e29c67d01",
80-
strip_prefix = "cudnn-linux-x86_64-8.4.1.50_cuda11.6-archive",
79+
sha256 = "5454a6fd94f008728caae9adad993c4e85ef36302e26bce43bea7d458a5e7b6d",
80+
strip_prefix = "cudnn-linux-x86_64-8.5.0.96_cuda11-archive",
8181
urls = [
82-
"https://developer.nvidia.com/compute/cudnn/secure/8.4.1/local_installers/11.6/cudnn-linux-x86_64-8.4.1.50_cuda11.6-archive.tar.xz",
82+
"https://developer.nvidia.com/compute/cudnn/secure/8.5.0/local_installers/11.7/cudnn-linux-x86_64-8.5.0.96_cuda11-archive.tar.xz",
8383
],
8484
)
8585

8686
http_archive(
8787
name = "tensorrt",
8888
build_file = "@//third_party/tensorrt/archive:BUILD",
89-
sha256 = "8d7c2085c1639dcc73875048c23598a8526ce3089136876e31d90258e49e4f61",
90-
strip_prefix = "TensorRT-8.4.3.1",
89+
sha256 = "39cc7f077057d1363794e8ff51c4cf21a5dbeccf1116b0020ba0dae0f3063076",
90+
strip_prefix = "TensorRT-8.5.1.7",
9191
urls = [
92-
"https://developer.nvidia.com/compute/machine-learning/tensorrt/secure/8.4.3/tars/tensorrt-8.4.3.1.linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz",
92+
"https://developer.nvidia.com/compute/machine-learning/tensorrt/secure/8.5.1/tars/TensorRT-8.5.1.7.Linux.x86_64-gnu.cuda-11.8.cudnn8.6.tar.gz",
9393
],
9494
)
9595

‎core/conversion/converters/converter_util.cpp

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,8 @@ nvinfer1::ITensor* addPadding(
1313
nvinfer1::ITensor* tensor,
1414
int nDim,
1515
bool trailing,
16-
bool use_zeros) {
16+
bool use_zeros,
17+
const std::string& name) {
1718
const auto dims = tensor->getDimensions();
1819

1920
if (dims.nbDims < nDim) {
@@ -27,7 +28,11 @@ nvinfer1::ITensor* addPadding(
2728
TORCHTRT_CHECK(shuffle_layer, "Unable to create shuffle layer");
2829
shuffle_layer->setReshapeDimensions(newDims);
2930
shuffle_layer->setZeroIsPlaceholder(use_zeros);
30-
shuffle_layer->setName((util::node_info(n) + " [Reshape to " + util::toStr(newDims) + ']').c_str());
31+
if (name.size()) {
32+
shuffle_layer->setName(name.c_str());
33+
} else {
34+
shuffle_layer->setName((util::node_info(n) + " [Reshape to " + util::toStr(newDims) + ']').c_str());
35+
}
3136
return shuffle_layer->getOutput(0);
3237
} else {
3338
return tensor;
@@ -40,7 +45,8 @@ nvinfer1::ITensor* addUnpadding(
4045
nvinfer1::ITensor* tensor,
4146
int nDim,
4247
bool trailing,
43-
bool use_zeros) {
48+
bool use_zeros,
49+
const std::string& name) {
4450
const auto dims = tensor->getDimensions();
4551
if (dims.nbDims > nDim) {
4652
auto newDims = dims;
@@ -52,7 +58,11 @@ nvinfer1::ITensor* addUnpadding(
5258
TORCHTRT_CHECK(shuffle_layer, "Unable to create shuffle layer");
5359
shuffle_layer->setReshapeDimensions(newDims);
5460
shuffle_layer->setZeroIsPlaceholder(use_zeros);
55-
shuffle_layer->setName((util::node_info(n) + " [Reshape to " + util::toStr(newDims) + "]").c_str());
61+
if (name.size()) {
62+
shuffle_layer->setName(name.c_str());
63+
} else {
64+
shuffle_layer->setName((util::node_info(n) + " [Reshape to " + util::toStr(newDims) + ']').c_str());
65+
}
5666
return shuffle_layer->getOutput(0);
5767
} else {
5868
return tensor;
@@ -195,7 +205,11 @@ nvinfer1::ITensor* applyIdentityOp(ConversionCtx* ctx, nvinfer1::ITensor* tensor
195205
return id_out_tensor;
196206
}
197207

198-
nvinfer1::ITensor* castITensor(ConversionCtx* ctx, nvinfer1::ITensor* tensor, nvinfer1::DataType dtype) {
208+
nvinfer1::ITensor* castITensor(
209+
ConversionCtx* ctx,
210+
nvinfer1::ITensor* tensor,
211+
nvinfer1::DataType dtype,
212+
const std::string& layer_name_prefix) {
199213
if (tensor->getType() != dtype) {
200214
std::ostringstream tensor_id;
201215
tensor_id << reinterpret_cast<int*>(tensor);
@@ -209,6 +223,9 @@ nvinfer1::ITensor* castITensor(ConversionCtx* ctx, nvinfer1::ITensor* tensor, nv
209223
LOG_DEBUG(ctx->logger, "Casting ITensor " << tensor_id.str() << " from " << tensor->getType() << " to " << dtype);
210224

211225
std::stringstream ss;
226+
if (layer_name_prefix.size()) {
227+
ss << layer_name_prefix << " ";
228+
}
212229
ss << "[Cast ITensor " << tensor_id.str() << " from " << tensor->getType() << " to " << dtype << "]";
213230
id_layer->setName(ss.str().c_str());
214231
return casted_tensor;

‎core/conversion/converters/converter_util.h

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,8 @@ nvinfer1::ITensor* addPadding(
2222
nvinfer1::ITensor* tensor,
2323
int nDim,
2424
bool trailing = true,
25-
bool use_zeros = true);
25+
bool use_zeros = true,
26+
const std::string& name = "");
2627

2728
// If nDim < tensor size, adds shuffle layer to un-pad tensor (at the end if trailing) and returns (nDim-dimensional)
2829
// shuffle layer's output Otherwise, does nothing and passes tensor through. use _zeros controls whether we should be
@@ -33,7 +34,8 @@ nvinfer1::ITensor* addUnpadding(
3334
nvinfer1::ITensor* tensor,
3435
int nDim,
3536
bool trailing = true,
36-
bool use_zeros = true);
37+
bool use_zeros = true,
38+
const std::string& name = "");
3739

3840
// TODO: Change add_elementwise schema to output nvinfer1::ITensor* instead of nvinfer1::ILayer*,
3941
// for consistency with other utils. Need to change schema and usage in all calling contexts
@@ -54,7 +56,11 @@ nvinfer1::ITensor* add_abs(
5456
nvinfer1::ITensor* applyIdentityOp(ConversionCtx* ctx, nvinfer1::ITensor* tensor, const std::string& name);
5557

5658
// If an ITensor is of a type not dtype, add an Identity layer to cast it to dtype
57-
nvinfer1::ITensor* castITensor(ConversionCtx* ctx, nvinfer1::ITensor* tensor, nvinfer1::DataType dtype);
59+
nvinfer1::ITensor* castITensor(
60+
ConversionCtx* ctx,
61+
nvinfer1::ITensor* tensor,
62+
nvinfer1::DataType dtype,
63+
const std::string& layer_name_prefix = "");
5864

5965
// Freeze an at::Tensor in a IConstant layer
6066
nvinfer1::ITensor* tensor_to_const(ConversionCtx* ctx, at::Tensor t, const std::string& name = std::string());

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ auto cast_registrations TORCHTRT_UNUSED =
2626
} else {
2727
trt_dtype = util::ScalarTypeToTRTDataType(static_cast<at::ScalarType>(output_dtype));
2828
}
29-
auto casted_itensor = castITensor(ctx, self, trt_dtype);
29+
auto casted_itensor = castITensor(ctx, self, trt_dtype, util::node_info(n));
3030
auto output = ctx->AssociateValueAndTensor(n->outputs()[0], casted_itensor);
3131
LOG_DEBUG("[aten::to.dtype] Output tensor shape: " << output->getDimensions());
3232

@@ -48,7 +48,7 @@ auto cast_registrations TORCHTRT_UNUSED =
4848
} else {
4949
trt_dtype = util::ScalarTypeToTRTDataType(static_cast<at::ScalarType>(output_dtype));
5050
}
51-
auto casted_itensor = castITensor(ctx, self, trt_dtype);
51+
auto casted_itensor = castITensor(ctx, self, trt_dtype, util::node_info(n));
5252
auto output = ctx->AssociateValueAndTensor(n->outputs()[0], casted_itensor);
5353
LOG_DEBUG("[aten::to.device] Output tensor shape: " << output->getDimensions());
5454

@@ -59,7 +59,7 @@ auto cast_registrations TORCHTRT_UNUSED =
5959
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
6060
auto self = args[0].ITensorOrFreeze(ctx);
6161
nvinfer1::DataType other_dtype = args[1].ITensorOrFreeze(ctx)->getType();
62-
auto casted_itensor = castITensor(ctx, self, other_dtype);
62+
auto casted_itensor = castITensor(ctx, self, other_dtype, util::node_info(n));
6363
auto output = ctx->AssociateValueAndTensor(n->outputs()[0], casted_itensor);
6464
LOG_DEBUG("[aten::to.other] Output tensor shape: " << output->getDimensions());
6565

@@ -77,7 +77,7 @@ auto cast_registrations TORCHTRT_UNUSED =
7777

7878
auto output_dtype = args[2].unwrapToScalar().to<int64_t>();
7979
auto trt_dtype = util::ScalarTypeToTRTDataType(static_cast<at::ScalarType>(output_dtype));
80-
auto casted_itensor = castITensor(ctx, self, trt_dtype);
80+
auto casted_itensor = castITensor(ctx, self, trt_dtype, util::node_info(n));
8181
auto output = ctx->AssociateValueAndTensor(n->outputs()[0], casted_itensor);
8282
LOG_DEBUG("[aten::to.prim_Device] Output tensor shape: " << output->getDimensions());
8383

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -325,7 +325,8 @@ auto element_wise_registrations TORCHTRT_UNUSED =
325325
add_elementwise(ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n));
326326
} else if (rounding_mode == "trunc") {
327327
// trunc = floor(abs(div)) * sign(div)
328-
auto tmp_div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, "tmp_div");
328+
auto tmp_div = add_elementwise(
329+
ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n) + "_tmp_div");
329330
auto abs = add_abs(ctx, n, tmp_div->getOutput(0), util::node_info(n) + "_absolute_val");
330331

331332
// In this case, we allow the floor unary on non-TRT Unary types, as it is needed for this

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,13 @@ auto reduce_registrations TORCHTRT_UNUSED =
7272
auto in_dims = util::toVec(in_tensor->getDimensions());
7373
LOG_WARNING("Sum Converter disregards dtype");
7474

75+
if (in_tensor->getType() == nvinfer1::DataType::kBOOL) {
76+
LOG_DEBUG(
77+
"Found type " << in_tensor->getType() << " in aten::sum, casting to "
78+
<< nvinfer1::DataType::kINT32 << " for compatibility.");
79+
in_tensor = castITensor(ctx, in_tensor, nvinfer1::DataType::kINT32);
80+
}
81+
7582
uint32_t axis_mask = (uint32_t)(((uint64_t)1 << in_dims.size()) - 1);
7683

7784
auto sum_layer = ctx->net->addReduce(*in_tensor, nvinfer1::ReduceOperation::kSUM, axis_mask, false);
@@ -113,6 +120,14 @@ auto reduce_registrations TORCHTRT_UNUSED =
113120
LOG_DEBUG("Keep dims: " << keepdim);
114121

115122
LOG_WARNING("Sum converter disregards dtype");
123+
124+
if (in_tensor->getType() == nvinfer1::DataType::kBOOL) {
125+
LOG_DEBUG(
126+
"Found type " << in_tensor->getType() << " in aten::sum, casting to "
127+
<< nvinfer1::DataType::kINT32 << " for compatibility.");
128+
in_tensor = castITensor(ctx, in_tensor, nvinfer1::DataType::kINT32);
129+
}
130+
116131
auto sum_layer = ctx->net->addReduce(*in_tensor, nvinfer1::ReduceOperation::kSUM, axis_mask, keepdim);
117132

118133
TORCHTRT_CHECK(sum_layer, "Unable to create sum layer from node: " << *n);

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

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,15 +39,29 @@ auto topk_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().patte
3939

4040
LOG_DEBUG("Output topk reduce dim: " << dim);
4141

42+
// The topk layer requires at least 2 input dimensions
43+
auto nbDims = self->getDimensions().nbDims;
44+
if (nbDims == 1) {
45+
self = addPadding(ctx, n, self, 2, true, true);
46+
}
47+
4248
auto TopKOperation = largest ? (nvinfer1::TopKOperation::kMAX) : (nvinfer1::TopKOperation::kMIN);
4349

4450
auto new_layer = ctx->net->addTopK(*self, TopKOperation, k, shiftDim);
4551

4652
TORCHTRT_CHECK(new_layer, "Unable to create topk layer from node: " << *n);
4753

48-
auto out0 = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
49-
auto out1 = ctx->AssociateValueAndTensor(n->outputs()[1], new_layer->getOutput(1));
54+
auto values = new_layer->getOutput(0);
55+
auto indices = new_layer->getOutput(1);
56+
57+
// If we expanded the input, squeeze the outputs
58+
if (nbDims == 1) {
59+
values = addUnpadding(ctx, n, values, 1, true, true, util::node_info(n) + "_squeeze_values");
60+
indices = addUnpadding(ctx, n, indices, 1, true, true, util::node_info(n) + "_squeeze_indices");
61+
}
5062

63+
auto out0 = ctx->AssociateValueAndTensor(n->outputs()[0], values);
64+
auto out1 = ctx->AssociateValueAndTensor(n->outputs()[1], indices);
5165
LOG_DEBUG("Output tensor(0) shape: " << out0->getDimensions());
5266
LOG_DEBUG("Output tensor(1) shape: " << out1->getDimensions());
5367

‎cpp/CMakeLists.txt

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ target_sources(${lib_name}
2424

2525
target_link_libraries(${lib_name}
2626
PUBLIC
27-
torch
27+
"${TORCH_LIBRARIES}"
2828
TensorRT::nvinfer
2929
core
3030
)
@@ -71,7 +71,7 @@ target_sources(${torchtrt_lib_name}
7171
target_link_libraries(${torchtrt_lib_name}
7272
PUBLIC
7373
TensorRT::TensorRT
74-
torch
74+
"${TORCH_LIBRARIES}"
7575
PRIVATE
7676
torch_tensorrt
7777
core
@@ -111,7 +111,7 @@ target_sources(${runtime_lib_name}
111111
target_link_libraries(${runtime_lib_name}
112112
PUBLIC
113113
TensorRT::TensorRT
114-
torch
114+
"${TORCH_LIBRARIES}"
115115
PRIVATE
116116
core_runtime
117117
core_plugins

0 commit comments

Comments
 (0)