docs(//py): Adding documentation on the limitation of the backend API

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

Author: narendasan

py/trtorch/_compile_spec.py

@@ -214,6 +214,10 @@ def TensorRTCompileSpec(compile_spec: Dict[str, Any]) -> torch.classes.tensorrt.
             One key is required which is ``input_shapes``, describing the input sizes or ranges for inputs
             to the graph. All other keys are optional. Entries for each method to be compiled.
 
+            Note: Partial compilation of TorchScript modules is not supported through the PyTorch TensorRT backend
+            If you need this feature, use trtorch.compile to compile your module. Usage of the resulting module is
+            as if you were using the TensorRT integration.
+
             .. code-block:: py
 
                 CompileSpec = {
@@ -272,7 +276,9 @@ def TensorRTCompileSpec(compile_spec: Dict[str, Any]) -> torch.classes.tensorrt.
     d._set_allow_gpu_fallback(parsed_spec.device.allow_gpu_fallback)
 
     if parsed_spec.torch_fallback.enabled:
-        raise RuntimeError("Partial module compilation is not currently supported via the PyTorch to_backend API integration. If you need partial compilation, use trtorch.compile")
+        raise RuntimeError(
+            "Partial module compilation is not currently supported via the PyTorch TensorRT backend. If you need partial compilation, use trtorch.compile"
+        )
 
     torch_fallback = torch.classes.tensorrt._TorchFallback()
     torch_fallback._set_enabled(parsed_spec.torch_fallback.enabled)

py/trtorch/_compiler.py

@@ -50,11 +50,11 @@ def compile(module: torch.jit.ScriptModule, compile_spec: Any) -> torch.jit.Scri
                     "workspace_size": 0, # Maximum size of workspace given to TensorRT
                     "max_batch_size": 0, # Maximum batch size (must be >= 1 to be set, 0 means not set)
                     "torch_fallback": {
-                        "enabled": True,
+                        "enabled": True, # Turn on or turn off falling back to PyTorch if operations are not supported in TensorRT
                         "force_fallback_ops": [
-                            "aten::max_pool2d"
+                            "aten::max_pool2d" # List of specific ops to require running in PyTorch
                         ],
-                        "min_block_size": 1
+                        "min_block_size": 3 # Minimum number of ops an engine must incapsulate to be run in TensorRT
                     }
                 }
 

py/trtorch/csrc/register_tensorrt_classes.cpp

@@ -5,33 +5,32 @@ namespace backend {
 namespace {
 
 #define ADD_FIELD_GET_SET_REGISTRATION(registry, class_name, field_name) \
-  (registry).def("_set_" #field_name, &class_name::set_##field_name);     \
+  (registry).def("_set_" #field_name, &class_name::set_##field_name);    \
   (registry).def("_get_" #field_name, &class_name::get_##field_name);
 
 void RegisterTRTCompileSpec() {
   static auto TRTORCH_UNUSED TRTInputRangeTSRegistration =
-    torch::class_<trtorch::pyapi::InputRange>("tensorrt", "_InputRange")
-      .def(torch::init<>())
-      .def("__str__", &trtorch::pyapi::InputRange::to_str);
+      torch::class_<trtorch::pyapi::InputRange>("tensorrt", "_InputRange")
+          .def(torch::init<>())
+          .def("__str__", &trtorch::pyapi::InputRange::to_str);
 
   ADD_FIELD_GET_SET_REGISTRATION(TRTInputRangeTSRegistration, trtorch::pyapi::InputRange, min);
   ADD_FIELD_GET_SET_REGISTRATION(TRTInputRangeTSRegistration, trtorch::pyapi::InputRange, opt);
   ADD_FIELD_GET_SET_REGISTRATION(TRTInputRangeTSRegistration, trtorch::pyapi::InputRange, max);
 
-  static auto TRTORCH_UNUSED TRTDeviceTSRegistration =
-    torch::class_<trtorch::pyapi::Device>("tensorrt", "_Device")
-      .def(torch::init<>())
-      .def("__str__", &trtorch::pyapi::Device::to_str);
+  static auto TRTORCH_UNUSED TRTDeviceTSRegistration = torch::class_<trtorch::pyapi::Device>("tensorrt", "_Device")
+                                                           .def(torch::init<>())
+                                                           .def("__str__", &trtorch::pyapi::Device::to_str);
 
   ADD_FIELD_GET_SET_REGISTRATION(TRTDeviceTSRegistration, trtorch::pyapi::Device, device_type);
   ADD_FIELD_GET_SET_REGISTRATION(TRTDeviceTSRegistration, trtorch::pyapi::Device, gpu_id);
   ADD_FIELD_GET_SET_REGISTRATION(TRTDeviceTSRegistration, trtorch::pyapi::Device, dla_core);
   ADD_FIELD_GET_SET_REGISTRATION(TRTDeviceTSRegistration, trtorch::pyapi::Device, allow_gpu_fallback);
 
   static auto TRTORCH_UNUSED TRTFallbackTSRegistration =
-    torch::class_<trtorch::pyapi::TorchFallback>("tensorrt", "_TorchFallback")
-      .def(torch::init<>())
-      .def("__str__", &trtorch::pyapi::TorchFallback::to_str);
+      torch::class_<trtorch::pyapi::TorchFallback>("tensorrt", "_TorchFallback")
+          .def(torch::init<>())
+          .def("__str__", &trtorch::pyapi::TorchFallback::to_str);
 
   ADD_FIELD_GET_SET_REGISTRATION(TRTFallbackTSRegistration, trtorch::pyapi::TorchFallback, enabled);
   ADD_FIELD_GET_SET_REGISTRATION(TRTFallbackTSRegistration, trtorch::pyapi::TorchFallback, min_block_size);

py/trtorch/csrc/tensorrt_classes.h

@@ -93,7 +93,6 @@ struct TorchFallback : torch::CustomClassHolder {
   std::string to_str();
 };
 
-
 enum class EngineCapability : int8_t {
   kDEFAULT,
   kSAFE_GPU,