feat: Add optional tensor domain argument to Input class (#1537)

Author: gs-olive

core/ir/Input.cpp

@@ -69,11 +69,16 @@ bool valid_input_dtype(nvinfer1::DataType dtype) {
   }
 }
 
+bool valid_input_domain(std::vector<double> domain) {
+  return (domain.size() == 2) && (domain[0] < domain[1]);
+}
+
 Input::Input(
     std::vector<int64_t> shape,
     at::ScalarType dtype,
     nvinfer1::TensorFormat format,
-    bool dtype_is_user_defined) {
+    bool dtype_is_user_defined,
+    std::vector<double> tensor_domain) {
   if (shape.size() > 5) {
     LOG_WARNING("Verify that this dim size is accepted");
   }
@@ -93,6 +98,11 @@ Input::Input(
           << "), Torch-TensorRT only supports contiguous format (NCHW) except with input type Float32 where channel last (NHWC) is also supported");
   this->format = format;
   this->dtype_is_user_defined = dtype_is_user_defined;
+
+  TORCHTRT_CHECK(
+      valid_input_domain(tensor_domain),
+      "Unsupported tensor domain: [" << tensor_domain[0] << ", " << tensor_domain[1] << ")");
+  this->tensor_domain = tensor_domain;
 }
 
 Input::Input(
@@ -101,7 +111,8 @@ Input::Input(
     std::vector<int64_t> max_shape,
     at::ScalarType dtype,
     nvinfer1::TensorFormat format,
-    bool dtype_is_user_defined) {
+    bool dtype_is_user_defined,
+    std::vector<double> tensor_domain) {
   if (min_shape.size() > 5 || opt_shape.size() > 5 || max_shape.size() > 5) {
     LOG_WARNING("Verify that this dim size is accepted");
   }
@@ -146,6 +157,10 @@ Input::Input(
           << "), Torch-TensorRT only supports contiguous format (NCHW) except with input type Float32 where channel last (NHWC) is also supported");
   this->format = format;
   this->dtype_is_user_defined = dtype_is_user_defined;
+  TORCHTRT_CHECK(
+      valid_input_domain(tensor_domain),
+      "Unsupported tensor domain: [" << tensor_domain[0] << ", " << tensor_domain[1] << ")");
+  this->tensor_domain = tensor_domain;
 }
 
 std::ostream& operator<<(std::ostream& os, const Input& input) {

core/ir/ir.h

@@ -31,19 +31,22 @@ struct Input : torch::CustomClassHolder {
       std::vector<int64_t> shape,
       at::ScalarType dtype = at::kFloat,
       nvinfer1::TensorFormat format = nvinfer1::TensorFormat::kLINEAR,
-      bool dtype_is_user_defined = false);
+      bool dtype_is_user_defined = false,
+      std::vector<double> tensor_domain = std::vector<double>{0, 2});
   Input(
       std::vector<int64_t> min_shape,
       std::vector<int64_t> opt_shape,
       std::vector<int64_t> max_shape,
       at::ScalarType dtype = at::kFloat,
       nvinfer1::TensorFormat format = nvinfer1::TensorFormat::kLINEAR,
-      bool dtype_is_used_defined = false);
+      bool dtype_is_user_defined = false,
+      std::vector<double> tensor_domain = std::vector<double>{0, 2});
 
   friend std::ostream& operator<<(std::ostream& os, const Input& input);
 
   bool input_is_dynamic = false;
   bool dtype_is_user_defined = false;
+  std::vector<double> tensor_domain;
   nvinfer1::Dims input_shape;
   nvinfer1::Dims min;
   nvinfer1::Dims max;

core/partitioning/shape_analysis.cpp

@@ -26,8 +26,8 @@ at::Tensor generateSingleInput(
   }
 
   // Initialize min and max ranges for random number selection
-  int LoValIncl = 0;
-  int HiValExcl = 2;
+  double LoValIncl = input.tensor_domain[0];
+  double HiValExcl = input.tensor_domain[1];
 
   auto type = at::kFloat;
   if (type_opt) {
@@ -36,6 +36,10 @@ at::Tensor generateSingleInput(
     LOG_WARNING("Input type for doing shape analysis could not be determined, defaulting to F32");
   }
 
+  LOG_DEBUG(
+      "Using the Range: [" << LoValIncl << ", " << HiValExcl
+                           << ") as a random range for shape analysis on input with data type " << type);
+
   // Make the value range for input tensor a uniform (float) distribution
   // over [LoValIncl, HiValExcl), then cast to the desired dtype
   auto in = ((HiValExcl - LoValIncl) * at::rand(util::toVec(input_shape), {at::kCUDA}) + LoValIncl).to(type);

cpp/include/torch_tensorrt/torch_tensorrt.h

@@ -381,6 +381,8 @@ struct Input : torch::CustomClassHolder {
   DataType dtype;
   /// Expected tensor format for the input
   TensorFormat format;
+  /// Expected allowed domain for tensor input
+  std::vector<double> tensor_domain;
 
   Input() {}
   /**
@@ -394,6 +396,22 @@ struct Input : torch::CustomClassHolder {
    */
   TORCHTRT_API Input(std::vector<int64_t> shape, TensorFormat format = TensorFormat::kContiguous);
 
+  /**
+   * @brief Construct a new Input spec object for static input size from
+   * c10::ArrayRef (the type produced by tensor.sizes()), vector, optional arguments
+   * allow the user to configure expected input shape tensor format
+   * dtype (Expected data type for the input) defaults to PyTorch
+   * / traditional TRT convection (FP32 for FP32 only, FP16 for FP32 and FP16, FP32 for Int8)
+   *
+   * @param shape Input tensor shape
+   * @param tensor_domain Allowed range for tensor inputs [low, high)
+   * @param format Expected tensor format for the input (Defaults to contiguous)
+   */
+  TORCHTRT_API Input(
+      std::vector<int64_t> shape,
+      std::vector<double> tensor_domain,
+      TensorFormat format = TensorFormat::kContiguous);
+
   /**
    * @brief Construct a new Input spec object for static input size from
    * vector, optional arguments allow the user to configure expected input shape
@@ -406,6 +424,23 @@ struct Input : torch::CustomClassHolder {
    */
   TORCHTRT_API Input(std::vector<int64_t> shape, DataType dtype, TensorFormat format = TensorFormat::kContiguous);
 
+  /**
+   * @brief Construct a new Input spec object for static input size from
+   * vector, optional arguments allow the user to configure expected input shape
+   * tensor format
+   *
+   * @param shape Input tensor shape
+   * @param dtype Expected data type for the input (Defaults to the type of the weights in the first tensor
+   * calculation if detectable else Float32)
+   * @param tensor_domain Allowed range for tensor inputs [low, high)
+   * @param format Expected tensor format for the input (Defaults to contiguous)
+   */
+  TORCHTRT_API Input(
+      std::vector<int64_t> shape,
+      DataType dtype,
+      std::vector<double> tensor_domain,
+      TensorFormat format = TensorFormat::kContiguous);
+
   /**
    * @brief Construct a new Input spec object for static input size from
    * c10::ArrayRef (the type produced by tensor.sizes()), vector, optional arguments
@@ -418,6 +453,22 @@ struct Input : torch::CustomClassHolder {
    */
   TORCHTRT_API Input(c10::ArrayRef<int64_t> shape, TensorFormat format = TensorFormat::kContiguous);
 
+  /**
+   * @brief Construct a new Input spec object for static input size from
+   * c10::ArrayRef (the type produced by tensor.sizes()), vector, optional arguments
+   * allow the user to configure expected input shape tensor format
+   * dtype (Expected data type for the input) defaults to PyTorch
+   * / traditional TRT convection (FP32 for FP32 only, FP16 for FP32 and FP16, FP32 for Int8)
+   *
+   * @param shape Input tensor shape
+   * @param tensor_domain Allowed range for tensor inputs [low, high)
+   * @param format Expected tensor format for the input (Defaults to contiguous)
+   */
+  TORCHTRT_API Input(
+      c10::ArrayRef<int64_t> shape,
+      std::vector<double> tensor_domain,
+      TensorFormat format = TensorFormat::kContiguous);
+
   /**
    * @brief Construct a new Input spec object for static input size from
    * c10::ArrayRef (the type produced by tensor.sizes()), vector, optional arguments
@@ -430,6 +481,23 @@ struct Input : torch::CustomClassHolder {
    */
   TORCHTRT_API Input(c10::ArrayRef<int64_t> shape, DataType dtype, TensorFormat format = TensorFormat::kContiguous);
 
+  /**
+   * @brief Construct a new Input spec object for static input size from
+   * c10::ArrayRef (the type produced by tensor.sizes()), vector, optional arguments
+   * allow the user to configure expected input shape tensor format
+   *
+   * @param shape Input tensor shape
+   * @param dtype Expected data type for the input (Defaults to the type of the weights in the first tensor
+   * calculation if detectable else Float32)
+   * @param tensor_domain Allowed range for tensor inputs [low, high)
+   * @param format Expected tensor format for the input (Defaults to contiguous)
+   */
+  TORCHTRT_API Input(
+      c10::ArrayRef<int64_t> shape,
+      DataType dtype,
+      std::vector<double> tensor_domain,
+      TensorFormat format = TensorFormat::kContiguous);
+
   /**
    * @brief Construct a new Input spec object dynamic input size from
    * c10::ArrayRef (the type produced by tensor.sizes()) for min, opt, and max
@@ -446,6 +514,24 @@ struct Input : torch::CustomClassHolder {
       std::vector<int64_t> opt_shape,
       std::vector<int64_t> max_shape,
       TensorFormat format = TensorFormat::kContiguous);
+  /**
+   * @brief Construct a new Input spec object dynamic input size from
+   * c10::ArrayRef (the type produced by tensor.sizes()) for min, opt, and max
+   * supported sizes. dtype (Expected data type for the input) defaults to PyTorch
+   * / traditional TRT convection (FP32 for FP32 only, FP16 for FP32 and FP16, FP32 for Int8)
+   *
+   * @param min_shape Minimum shape for input tensor
+   * @param opt_shape Target optimization shape for input tensor
+   * @param max_shape Maximum acceptible shape for input tensor
+   * @param tensor_domain Allowed range for tensor inputs [low, high)
+   * @param format Expected tensor format for the input (Defaults to contiguous)
+   */
+  TORCHTRT_API Input(
+      std::vector<int64_t> min_shape,
+      std::vector<int64_t> opt_shape,
+      std::vector<int64_t> max_shape,
+      std::vector<double> tensor_domain,
+      TensorFormat format = TensorFormat::kContiguous);
 
   /**
    * @brief Construct a new Input spec object for a dynamic input size from vectors
@@ -466,6 +552,44 @@ struct Input : torch::CustomClassHolder {
       DataType dtype,
       TensorFormat format = TensorFormat::kContiguous);
 
+  /**
+   * @brief Construct a new Input spec object for a dynamic input size from vectors
+   * for minimum shape, optimal shape, and max shape supported sizes optional arguments
+   * allow the user to configure expected input shape tensor format
+   *
+   * @param min_shape Minimum shape for input tensor
+   * @param opt_shape Target optimization shape for input tensor
+   * @param max_shape Maximum acceptible shape for input tensor
+   * @param dtype Expected data type for the input (Defaults to the type of the weights in the first tensor
+   * calculation if detectable else Float32)
+   * @param tensor_domain Allowed range for tensor inputs [low, high)
+   * @param format Expected tensor format for the input (Defaults to contiguous)
+   */
+  TORCHTRT_API Input(
+      std::vector<int64_t> min_shape,
+      std::vector<int64_t> opt_shape,
+      std::vector<int64_t> max_shape,
+      DataType dtype,
+      std::vector<double> tensor_domain,
+      TensorFormat format = TensorFormat::kContiguous);
+
+  /**
+   * @brief Construct a new Input spec object dynamic input size from
+   * c10::ArrayRef (the type produced by tensor.sizes()) for min, opt, and max
+   * supported sizes. dtype (Expected data type for the input) defaults to PyTorch
+   * / traditional TRT convection (FP32 for FP32 only, FP16 for FP32 and FP16, FP32 for Int8)
+   *
+   * @param min_shape Minimum shape for input tensor
+   * @param opt_shape Target optimization shape for input tensor
+   * @param max_shape Maximum acceptible shape for input tensor
+   * @param format Expected tensor format for the input (Defaults to contiguous)
+   */
+  TORCHTRT_API Input(
+      c10::ArrayRef<int64_t> min_shape,
+      c10::ArrayRef<int64_t> opt_shape,
+      c10::ArrayRef<int64_t> max_shape,
+      TensorFormat format = TensorFormat::kContiguous);
+
   /**
    * @brief Construct a new Input spec object dynamic input size from
    * c10::ArrayRef (the type produced by tensor.sizes()) for min, opt, and max
@@ -475,12 +599,33 @@ struct Input : torch::CustomClassHolder {
    * @param min_shape Minimum shape for input tensor
    * @param opt_shape Target optimization shape for input tensor
    * @param max_shape Maximum acceptible shape for input tensor
+   * @param tensor_domain Allowed range for tensor inputs [low, high)
    * @param format Expected tensor format for the input (Defaults to contiguous)
    */
   TORCHTRT_API Input(
       c10::ArrayRef<int64_t> min_shape,
       c10::ArrayRef<int64_t> opt_shape,
       c10::ArrayRef<int64_t> max_shape,
+      std::vector<double> tensor_domain,
+      TensorFormat format = TensorFormat::kContiguous);
+
+  /**
+   * @brief Construct a new Input spec object dynamic input size from
+   * c10::ArrayRef (the type produced by tensor.sizes()) for min, opt, and max
+   * supported sizes
+   *
+   * @param min_shape Minimum shape for input tensor
+   * @param opt_shape Target optimization shape for input tensor
+   * @param max_shape Maximum acceptible shape for input tensor
+   * @param dtype Expected data type for the input (Defaults to the type of the weights in the first tensor
+   * calculation if detectable else Float32)
+   * @param format Expected tensor format for the input (Defaults to contiguous)
+   */
+  TORCHTRT_API Input(
+      c10::ArrayRef<int64_t> min_shape,
+      c10::ArrayRef<int64_t> opt_shape,
+      c10::ArrayRef<int64_t> max_shape,
+      DataType dtype,
       TensorFormat format = TensorFormat::kContiguous);
 
   /**
@@ -493,13 +638,15 @@ struct Input : torch::CustomClassHolder {
    * @param max_shape Maximum acceptible shape for input tensor
    * @param dtype Expected data type for the input (Defaults to the type of the weights in the first tensor
    * calculation if detectable else Float32)
+   * @param tensor_domain Allowed range for tensor inputs [low, high)
    * @param format Expected tensor format for the input (Defaults to contiguous)
    */
   TORCHTRT_API Input(
       c10::ArrayRef<int64_t> min_shape,
       c10::ArrayRef<int64_t> opt_shape,
       c10::ArrayRef<int64_t> max_shape,
       DataType dtype,
+      std::vector<double> tensor_domain,
       TensorFormat format = TensorFormat::kContiguous);
 
   /**