refactor(//py): Redo Python tests for new API

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

core/compiler.cpp

@@ -297,46 +297,48 @@ void MapInputsAndDetermineDTypes(
   cfg.convert_info.inputs = std::move(ir::associate_specs_with_inputs(g, cfg.inputs, static_params));
 
   for (auto& in : g->inputs()) {
-    auto est_type_opt = first_use_type_map.find(in)->second;
-    ir::Input& spec = cfg.convert_info.inputs.find(in)->second;
-    if (est_type_opt && !spec.dtype_is_user_defined) {
-      // If we can calculate the type from the graph and the type was not defined by the user then use the calculated
-      // type
-      LOG_INFO(
-          "Since input type is not explicitly defined, infering using first tensor calculation\n  Found input "
-          << in->debugName() << " has type " << est_type_opt.value()
-          << ". If this is incorrect explicitly set dtype for input and file a bug");
-      spec.dtype = util::ScalarTypeToTRTDataType(est_type_opt.value());
-    } else if (!est_type_opt && !spec.dtype_is_user_defined) {
-      // If we cannot calculate the type and the user did not define the type, then default to FP32
-      LOG_WARNING(
-          "Cannot infer input type from calcuations in graph for input "
-          << in->debugName() << ". Assuming it is Float32. If not, specify input type explicity");
-      spec.dtype = nvinfer1::DataType::kFLOAT;
-    } else if (spec.dtype_is_user_defined && cfg.partition_info.enabled) {
-      if (!est_type_opt) {
-        LOG_INFO("Cannot infer input tensor dtype in graph, unable to verify user input dtype settings");
-      } else {
-        if (util::TRTDataTypeToScalarType(cfg.convert_info.inputs.find(in)->second.dtype) != est_type_opt.value()) {
-          std::stringstream ss;
-          ss << "For input " << in->debugName() << ", found user specified input dtype as ";
-          ss << cfg.convert_info.inputs.find(in)->second.dtype;
-          ss << ", however when inspecting the graph, the input type expected was inferred to be ";
-          ss << est_type_opt.value() << std::endl;
-          ss << "The compiler is going to use the user setting " << cfg.convert_info.inputs.find(in)->second.dtype;
-          ss << "\nThis conflict may cause an error at runtime due to partial compilation being enabled and therefore\n";
-          ss << "compatibility with PyTorch's data type convention is required.\n";
-          ss << "If you do indeed see errors at runtime either:\n";
-          ss << "- Remove the dtype spec for " << in->debugName() << std::endl;
-          ss << "- Disable partial compilation by setting require_full_compilation to True";
-          auto warn_str = ss.str();
-          LOG_WARNING(warn_str);
-          // Overwrite type map with user settings
-          first_use_type_map[in] = {util::TRTDataTypeToScalarType(cfg.convert_info.inputs.find(in)->second.dtype)};
+    if (static_params.find(in) == static_params.end()) {
+      ir::Input& spec = cfg.convert_info.inputs.find(in)->second;
+      auto est_type_opt = first_use_type_map.find(in)->second;
+      if (est_type_opt && !spec.dtype_is_user_defined) {
+        // If we can calculate the type from the graph and the type was not defined by the user then use the calculated
+        // type
+        LOG_INFO(
+            "Since input type is not explicitly defined, infering using first tensor calculation\n  Found input "
+            << in->debugName() << " has type " << est_type_opt.value()
+            << ". If this is incorrect explicitly set dtype for input and file a bug");
+        spec.dtype = util::ScalarTypeToTRTDataType(est_type_opt.value());
+      } else if (!est_type_opt && !spec.dtype_is_user_defined) {
+        // If we cannot calculate the type and the user did not define the type, then default to FP32
+        LOG_WARNING(
+            "Cannot infer input type from calcuations in graph for input "
+            << in->debugName() << ". Assuming it is Float32. If not, specify input type explicity");
+        spec.dtype = nvinfer1::DataType::kFLOAT;
+      } else if (spec.dtype_is_user_defined && cfg.partition_info.enabled) {
+        if (!est_type_opt) {
+          LOG_INFO("Cannot infer input tensor dtype in graph, unable to verify user input dtype settings");
+        } else {
+          if (util::TRTDataTypeToScalarType(cfg.convert_info.inputs.find(in)->second.dtype) != est_type_opt.value()) {
+            std::stringstream ss;
+            ss << "For input " << in->debugName() << ", found user specified input dtype as ";
+            ss << cfg.convert_info.inputs.find(in)->second.dtype;
+            ss << ", however when inspecting the graph, the input type expected was inferred to be ";
+            ss << est_type_opt.value() << std::endl;
+            ss << "The compiler is going to use the user setting " << cfg.convert_info.inputs.find(in)->second.dtype;
+            ss << "\nThis conflict may cause an error at runtime due to partial compilation being enabled and therefore\n";
+            ss << "compatibility with PyTorch's data type convention is required.\n";
+            ss << "If you do indeed see errors at runtime either:\n";
+            ss << "- Remove the dtype spec for " << in->debugName() << std::endl;
+            ss << "- Disable partial compilation by setting require_full_compilation to True";
+            auto warn_str = ss.str();
+            LOG_WARNING(warn_str);
+            // Overwrite type map with user settings
+            first_use_type_map[in] = {util::TRTDataTypeToScalarType(cfg.convert_info.inputs.find(in)->second.dtype)};
+          }
         }
+      } else {
+        // The user defined the type so no changes are necessary
       }
-    } else {
-      // The user defined the type so no changes are necessary
     }
   }
 }

core/util/jit_util.cpp

@@ -12,19 +12,19 @@ class Device(object):
     Defines a device that can be used to specify target devices for engines
 
     Attributes:
-        device_type (trtorch.DeviceType): Target device type (GPU or DLA). Set implicitly based on if dla_core is specified.
+        device_type (torch_tensorrt.DeviceType): Target device type (GPU or DLA). Set implicitly based on if dla_core is specified.
         gpu_id (int): Device ID for target GPU
         dla_core (int): Core ID for target DLA core
         allow_gpu_fallback (bool): Whether falling back to GPU if DLA cannot support an op should be allowed
     """
 
-    device_type = None  #: (trtorch.DeviceType): Target device type (GPU or DLA). Set implicitly based on if dla_core is specified.
+    device_type = None  #: (torch_tensorrt.DeviceType): Target device type (GPU or DLA). Set implicitly based on if dla_core is specified.
     gpu_id = -1  #: (int) Device ID for target GPU
     dla_core = -1  #: (int) Core ID for target DLA core
     allow_gpu_fallback = False  #: (bool) Whether falling back to GPU if DLA cannot support an op should be allowed
 
     def __init__(self, *args, **kwargs):
-        """ __init__ Method for trtorch.Device
+        """ __init__ Method for torch_tensorrt.Device
 
         Device accepts one of a few construction patterns
 

py/torch_tensorrt/_Device.py

@@ -12,30 +12,30 @@ class Input(object):
     Defines an input to a module in terms of expected shape, data type and tensor format.
 
     Attributes:
-        shape_mode (trtorch.Input._ShapeMode): Is input statically or dynamically shaped
+        shape_mode (torch_tensorrt.Input._ShapeMode): Is input statically or dynamically shaped
         shape (Tuple or Dict): Either a single Tuple or a dict of tuples defining the input shape.
             Static shaped inputs will have a single tuple. Dynamic inputs will have a dict of the form
             ``{
                 "min_shape": Tuple,
                 "opt_shape": Tuple,
                 "max_shape": Tuple
             }``
-        dtype (trtorch.dtype): The expected data type of the input tensor (default: trtorch.dtype.float32)
-        format (trtorch.TensorFormat): The expected format of the input tensor (default: trtorch.TensorFormat.NCHW)
+        dtype (torch_tensorrt.dtype): The expected data type of the input tensor (default: torch_tensorrt.dtype.float32)
+        format (torch_tensorrt.TensorFormat): The expected format of the input tensor (default: torch_tensorrt.TensorFormat.NCHW)
     """
 
     class _ShapeMode(Enum):
         STATIC = 0
         DYNAMIC = 1
 
-    shape_mode = None  #: (trtorch.Input._ShapeMode): Is input statically or dynamically shaped
+    shape_mode = None  #: (torch_tensorrt.Input._ShapeMode): Is input statically or dynamically shaped
     shape = None  #: (Tuple or Dict): Either a single Tuple or a dict of tuples defining the input shape. Static shaped inputs will have a single tuple. Dynamic inputs will have a dict of the form ``{ "min_shape": Tuple, "opt_shape": Tuple, "max_shape": Tuple }``
-    dtype = _enums.dtype.unknown  #: The expected data type of the input tensor (default: trtorch.dtype.float32)
+    dtype = _enums.dtype.unknown  #: The expected data type of the input tensor (default: torch_tensorrt.dtype.float32)
     _explicit_set_dtype = False
-    format = _enums.TensorFormat.contiguous  #: The expected format of the input tensor (default: trtorch.TensorFormat.NCHW)
+    format = _enums.TensorFormat.contiguous  #: The expected format of the input tensor (default: torch_tensorrt.TensorFormat.NCHW)
 
     def __init__(self, *args, **kwargs):
-        """ __init__ Method for trtorch.Input
+        """ __init__ Method for torch_tensorrt.Input
 
         Input accepts one of a few construction patterns
 
@@ -50,13 +50,13 @@ def __init__(self, *args, **kwargs):
                 Note: All three of min_shape, opt_shape, max_shape must be provided, there must be no positional arguments, shape must not be defined and implictly this sets Input's shape_mode to DYNAMIC
             max_shape (Tuple or List, optional): Max size of input tensor's shape range
                 Note: All three of min_shape, opt_shape, max_shape must be provided, there must be no positional arguments, shape must not be defined and implictly this sets Input's shape_mode to DYNAMIC
-            dtype (torch.dtype or trtorch.dtype): Expected data type for input tensor (default: trtorch.dtype.float32)
-            format (torch.memory_format or trtorch.TensorFormat): The expected format of the input tensor (default: trtorch.TensorFormat.NCHW)
+            dtype (torch.dtype or torch_tensorrt.dtype): Expected data type for input tensor (default: torch_tensorrt.dtype.float32)
+            format (torch.memory_format or torch_tensorrt.TensorFormat): The expected format of the input tensor (default: torch_tensorrt.TensorFormat.NCHW)
 
         Examples:
             - Input([1,3,32,32], dtype=torch.float32, format=torch.channel_last)
-            - Input(shape=(1,3,32,32), dtype=trtorch.dtype.int32, format=trtorch.TensorFormat.NCHW)
-            - Input(min_shape=(1,3,32,32), opt_shape=[2,3,32,32], max_shape=(3,3,32,32)) #Implicitly dtype=trtorch.dtype.float32, format=trtorch.TensorFormat.NCHW
+            - Input(shape=(1,3,32,32), dtype=torch_tensorrt.dtype.int32, format=torch_tensorrt.TensorFormat.NCHW)
+            - Input(min_shape=(1,3,32,32), opt_shape=[2,3,32,32], max_shape=(3,3,32,32)) #Implicitly dtype=torch_tensorrt.dtype.float32, format=torch_tensorrt.TensorFormat.NCHW
         """
         if len(args) == 1:
             if not Input._supported_input_size_type(args[0]):
@@ -204,7 +204,7 @@ def _parse_dtype(dtype: Any) -> _enums.dtype:
             return dtype
 
         else:
-            raise TypeError("Input data type needs to be specified with a torch.dtype or a trtorch.dtype, got: " +
+            raise TypeError("Input data type needs to be specified with a torch.dtype or a torch_tensorrt.dtype, got: " +
                             str(type(dtype)))
 
     @staticmethod
@@ -223,7 +223,7 @@ def _parse_format(format: Any) -> _enums.TensorFormat:
 
         else:
             raise TypeError(
-                "Tensor format needs to be specified with either torch.memory_format or trtorch.TensorFormat")
+                "Tensor format needs to be specified with either torch.memory_format or torch_tensorrt.TensorFormat")
 
     @classmethod
     def _from_tensor(cls, t: torch.Tensor):

py/torch_tensorrt/_Input.py

@@ -2,19 +2,19 @@
 from torch_tensorrt import _C
 
 def dump_build_info():
-    """Prints build information about the TRTorch distribution to stdout
+    """Prints build information about the torch_tensorrt distribution to stdout
     """
     print(get_build_info())
 
 
 def get_build_info() -> str:
-    """Returns a string containing the build information of TRTorch distribution
+    """Returns a string containing the build information of torch_tensorrt distribution
 
     Returns:
-        str: String containing the build information for TRTorch distribution
+        str: String containing the build information for torch_tensorrt distribution
     """
     build_info = _C.get_build_info()
-    build_info = "TRTorch Version: " + str(__version__) + '\n' + build_info
+    build_info = "Torch-TensorRT Version: " + str(__version__) + '\n' + build_info
     return build_info
 
 

py/torch_tensorrt/_util.py

@@ -1 +1 @@
-__version__ = "1.0.0a0+808e603f"
+__version__ = "1.0.0a0+483ef591"

py/torch_tensorrt/_version.py

@@ -19,6 +19,8 @@ void RegisterTRTCompileSpec() {
   ADD_FIELD_GET_SET_REGISTRATION(TRTInputRangeTSRegistration, trtorch::pyapi::Input, dtype);
   ADD_FIELD_GET_SET_REGISTRATION(TRTInputRangeTSRegistration, trtorch::pyapi::Input, format);
   ADD_FIELD_GET_SET_REGISTRATION(TRTInputRangeTSRegistration, trtorch::pyapi::Input, input_is_dynamic);
+  ADD_FIELD_GET_SET_REGISTRATION(TRTInputRangeTSRegistration, trtorch::pyapi::Input, explicit_set_dtype);
+
 
   static auto TRTORCH_UNUSED TRTDeviceTSRegistration = torch::class_<trtorch::pyapi::Device>("tensorrt", "_Device")
                                                            .def(torch::init<>())

py/torch_tensorrt/csrc/register_tensorrt_classes.cpp

@@ -12,6 +12,7 @@ namespace backend {
 
 c10::impl::GenericDict TensorRTBackend::compile(c10::IValue mod_val, c10::impl::GenericDict method_compile_spec) {
   auto mod = mod_val.toModule();
+
   auto spec = c10::impl::toTypedDict<std::string, at::IValue>(method_compile_spec);
 
   auto handles = c10::impl::GenericDict(
@@ -22,17 +23,12 @@ c10::impl::GenericDict TensorRTBackend::compile(c10::IValue mod_val, c10::impl::
     const auto& method_name = it->key();
     auto raw_spec = it->value().toCustomClass<trtorch::pyapi::CompileSpec>();
     LOG_DEBUG(raw_spec->stringify());
-    auto cfg = raw_spec->toInternalCompileSpec();
-    auto graph_and_ivals = core::lowering::Lower(mod_, method_name, cfg.lower_info);
-
-    auto g = graph_and_ivals.first;
-    auto params = graph_and_ivals.second;
-    auto named_params = core::ir::get_static_params(g->inputs(), params);
 
+    auto cfg = raw_spec->toInternalCompileSpec();
     auto convert_cfg = std::move(cfg.convert_info);
     auto device_spec = convert_cfg.engine_settings.device;
     auto device = core::runtime::CudaDevice(device_spec.gpu_id, device_spec.device_type);
-    auto serialized_engine = core::conversion::ConvertBlockToEngine(g->block(), convert_cfg, named_params);
+    auto serialized_engine = core::ConvertGraphToTRTEngine(mod_, method_name, cfg);
     auto engine_handle = c10::make_intrusive<core::runtime::TRTEngine>(it->key(), serialized_engine, device);
     handles.insert(method_name, at::IValue(engine_handle));
   }

py/torch_tensorrt/csrc/tensorrt_backend.cpp

@@ -50,7 +50,7 @@ struct Input : torch::CustomClassHolder {
   ADD_FIELD_GET_SET(max, std::vector<int64_t>);
   ADD_FIELD_GET_SET(input_is_dynamic, bool);
   ADD_FIELD_GET_SET(explicit_set_dtype, bool);
-  ADD_ENUM_GET_SET(dtype, DataType, static_cast<int64_t>(DataType::kBool));
+  ADD_ENUM_GET_SET(dtype, DataType, static_cast<int64_t>(DataType::kUnknown));
   ADD_ENUM_GET_SET(format, TensorFormat, static_cast<int64_t>(TensorFormat::kContiguous));
 
   core::ir::Input toInternalInput();