Skip to content

Commit

Permalink
feat(trtorchc): Adding new support for dtypes and formats in
Browse files Browse the repository at this point in the history 
trtorchc

Signed-off-by: Naren Dasan <naren@narendasan.com>
Signed-off-by: Naren Dasan <narens@nvidia.com>
  • Loading branch information
@narendasan
narendasan committed Jul 21, 2021
1 parent de4f358 commit c39bf81
Show file tree
Hide file tree
Showing 5 changed files with 276 additions and 79 deletions.
4 changes: 0 additions & 4 deletions core/ir/Input.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -128,8 +128,6 @@ Input::Input(std::vector<int64_t> shape, nvinfer1::DataType dtype, nvinfer1::Ten
max = util::toDims(shape);
input_shape = util::toDims(shape);
input_is_dynamic = false;
format = nvinfer1::TensorFormat::kLINEAR;
dtype = dtype;

TRTORCH_CHECK(valid_input_dtype(dtype), "Unsupported input data type: " << dtype);
this->dtype = dtype;
Expand All @@ -156,8 +154,6 @@ Input::Input(std::vector<int64_t> min_shape, std::vector<int64_t> opt_shape, std
min = util::toDims(min_shape);
opt = util::toDims(opt_shape);
max = util::toDims(max_shape);
format = nvinfer1::TensorFormat::kLINEAR;
dtype = nvinfer1::DataType::kFLOAT;

std::vector<int64_t> dyn_shape;
for (size_t i = 0; i < opt_shape.size(); i++) {
Expand Down
13 changes: 11 additions & 2 deletions cpp/api/include/trtorch/trtorch.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#pragma once

#include <cuda_runtime.h>
#include <iostream>
#include <memory>
#include <string>
#include <vector>
Expand Down Expand Up @@ -66,10 +67,12 @@ struct TRTORCH_API CompileSpec {
kHalf,
/// INT8
kChar,
/// INT32
kInt32,
/// INT
kInt,
/// Bool
kBool,
/// Sentinel value
kUnknown
};

/**
Expand Down Expand Up @@ -139,6 +142,7 @@ struct TRTORCH_API CompileSpec {
}

private:
friend std::ostream& operator<<(std::ostream& os, const DataType& dtype);
Value value;
};

Expand Down Expand Up @@ -278,6 +282,8 @@ struct TRTORCH_API CompileSpec {
kContiguous,
/// Channel Last / NHWC
kChannelsLast,
/// Sentinel value
kUnknown,
};

/**
Expand Down Expand Up @@ -346,7 +352,9 @@ struct TRTORCH_API CompileSpec {
return value != other;
}


private:
friend std::ostream& operator<<(std::ostream& os, const TensorFormat& format);
Value value;
};

Expand Down Expand Up @@ -472,6 +480,7 @@ struct TRTORCH_API CompileSpec {

bool get_explicit_set_dtype() {return explicit_set_dtype;}
private:
friend std::ostream& operator<<(std::ostream& os, const Input& input);
bool input_is_dynamic;
bool explicit_set_dtype;
};
Expand Down
68 changes: 64 additions & 4 deletions cpp/api/src/compile_spec.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,74 @@

namespace trtorch {

std::ostream& operator<<(std::ostream& os, const CompileSpec::DataType& dtype) {
switch (dtype) {
case CompileSpec::DataType::kChar:
os << "char";
break;
case CompileSpec::DataType::kHalf:
os << "half";
break;
case CompileSpec::DataType::kInt:
os << "int";
break;
case CompileSpec::DataType::kBool:
os << "bool";
break;
case CompileSpec::DataType::kFloat:
os << "float";
break;
case CompileSpec::DataType::kUnknown:
default:
os << "unknown";
break;
}
return os;
}

std::ostream& operator<<(std::ostream& os, const CompileSpec::TensorFormat& format) {
switch (format) {
case CompileSpec::TensorFormat::kChannelsLast:
os << "channels last";
break;
case CompileSpec::TensorFormat::kContiguous:
os << "contiguous";
break;
case CompileSpec::TensorFormat::kUnknown:
default:
os << "unknown";
break;
}
return os;
}

std::ostream& operator<<(std::ostream& os, const CompileSpec::Input& input) {
auto vec_to_str = [](std::vector<int64_t> shape) -> std::string {
std::stringstream ss;
ss << '[';
for (auto i : shape) {
ss << i << ',';
}
ss << ']';
return ss.str();
};

if (!input.input_is_dynamic) {
os << "Input(shape: " << vec_to_str(input.shape) << ", dtype: " << input.dtype << ", format: " << input.format << ')';
} else {
os << "Input(shape: " << vec_to_str(input.shape) << ", min: " << vec_to_str(input.min_shape) << ", opt: " << vec_to_str(input.opt_shape) << ", max: " << vec_to_str(input.max_shape) << ", dtype: " << input.dtype << ", format: " << input.format << ')';
}
return os;
}


nvinfer1::DataType toTRTDataType(CompileSpec::DataType value) {
switch (value) {
case CompileSpec::DataType::kChar:
return nvinfer1::DataType::kINT8;
case CompileSpec::DataType::kHalf:
return nvinfer1::DataType::kHALF;
case CompileSpec::DataType::kInt32:
case CompileSpec::DataType::kInt:
return nvinfer1::DataType::kINT32;
case CompileSpec::DataType::kBool:
return nvinfer1::DataType::kBOOL;
Expand Down Expand Up @@ -47,7 +108,7 @@ CompileSpec::DataType::DataType(c10::ScalarType t) {
value = DataType::kChar;
break;
case at::kInt:
value = DataType::kInt32;
value = DataType::kInt;
break;
case at::kBool:
value = DataType::kBool;
Expand Down Expand Up @@ -250,7 +311,6 @@ core::CompileSpec to_internal_compile_spec(CompileSpec external) {
/* We want default behavior for types to match PyTorch, so in the case the user did not explicitly set the dtype for
inputs they will follow PyTorch convetions */
for (size_t i = 0; i < external.inputs.size(); i++) {
std::cout << "EXPLICIT " << external.inputs[i].get_explicit_set_dtype() << std::endl;
if (!external.inputs[i].get_explicit_set_dtype()) {
auto& precisions = internal.convert_info.engine_settings.enabled_precisions;
auto& internal_ins = internal.convert_info.inputs;
Expand All @@ -261,9 +321,9 @@ core::CompileSpec to_internal_compile_spec(CompileSpec external) {
} else {
internal_ins[i].dtype = nvinfer1::DataType::kFLOAT;
}
std::cout << "internal type: " << internal_ins[i].dtype;
}
}

internal.convert_info.engine_settings.disable_tf32 = external.disable_tf32;
internal.convert_info.engine_settings.refit = external.refit;
internal.convert_info.engine_settings.debug = external.debug;
Expand Down
32 changes: 21 additions & 11 deletions cpp/trtorchc/README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ to standard TorchScript. Load with `torch.jit.load()` and run like you would run

```
trtorchc [input_file_path] [output_file_path]
[input_shapes...] {OPTIONS}
[input_specs...] {OPTIONS}
TRTorch is a compiler for TorchScript, it will compile and optimize
TorchScript programs to run on NVIDIA GPUs using TensorRT
Expand All @@ -28,24 +28,29 @@ trtorchc [input_file_path] [output_file_path]
-w, --warnings Disables warnings generated during
compilation onto the console (warnings
are on by default)
--info Dumps info messages generated during
--i, --info Dumps info messages generated during
compilation onto the console
--build-debuggable-engine Creates a debuggable engine
--use-strict-types Restrict operating type to only use set
default operation precision
(op_precision)
operation precision
--allow-gpu-fallback (Only used when targeting DLA
(device-type)) Lets engine run layers on
GPU if they are not supported on DLA
-p[precision],
--default-op-precision=[precision]
Default operating precision for the
engine (Int8 requires a
--disable-tf32 Prevent Float32 layers from using the
TF32 data format
-p[precision...],
--enabled-precison=[precision...] (Repeatable) Enabling an operating
precision for kernels to use when
building the engine (Int8 requires a
calibration-cache argument) [ float |
float32 | f32 | half | float16 | f16 |
int8 | i8 ] (default: float)
-d[type], --device-type=[type] The type of device the engine should be
built for [ gpu | dla ] (default: gpu)
--gpu-id=[gpu_id] GPU id if running on multi-GPU platform
(defaults to 0)
--dla-core=[dla_core] DLACore id if running on available DLA
(defaults to 0)
--engine-capability=[capability] The type of device the engine should be
built for [ default | safe_gpu |
safe_dla ]
Expand All @@ -72,16 +77,21 @@ trtorchc [input_file_path] [output_file_path]
input_file_path Path to input TorchScript file
output_file_path Path for compiled TorchScript (or
TensorRT engine) file
input_shapes... Sizes for inputs to engine, can either
input_specs... Specs for inputs to engine, can either
be a single size or a range defined by
Min, Optimal, Max sizes, e.g.
"(N,..,C,H,W)"
"[(MIN_N,..,MIN_C,MIN_H,MIN_W);(OPT_N,..,OPT_C,OPT_H,OPT_W);(MAX_N,..,MAX_C,MAX_H,MAX_W)]"
"[(MIN_N,..,MIN_C,MIN_H,MIN_W);(OPT_N,..,OPT_C,OPT_H,OPT_W);(MAX_N,..,MAX_C,MAX_H,MAX_W)]".
Data Type and format can be specified by
adding an "@" followed by dtype and "%"
followed by format to the end of the
shape spec. e.g. "(3, 3, 32,
32)@f16%NHWC"
"--" can be used to terminate flag options and force all following
arguments to be treated as positional options
```

e.g.
```
trtorchc tests/modules/ssd_traced.jit.pt ssd_trt.ts "[(1,3,300,300); (1,3,512,512); (1, 3, 1024, 1024)]" -p f16
trtorchc tests/modules/ssd_traced.jit.pt ssd_trt.ts "[(1,3,300,300); (1,3,512,512); (1, 3, 1024, 1024)]@fp16%contiguous" -p f16
```
Loading

0 comments on commit c39bf81

Please sign in to comment.