Address comments

Author: yizhuoz004

tripy/nvtripy/backend/api/shape_bounds.py renamed to tripy/nvtripy/backend/api/bounds.py

@@ -25,62 +25,37 @@
 
 @export.public_api(document_under="compiling_code/input_info", document_init_sig=False)
 @dataclass
-class ShapeBounds:
+class Bounds:
     min: Tuple[IntLike]
     """
-    The minimum shape.
+    The minimum value.
     """
     opt: Tuple[IntLike]
     """
-    The shape to optimize for.
+    The value to optimize for.
     """
     max: Tuple[IntLike]
     """
-    The maximum shape.
+    The maximum value.
     """
 
     def is_static(self):
         return self.min == self.opt == self.max
 
 
-@json_utils.Encoder.register(ShapeBounds)
-def encode_shape_bounds(shape_bounds):
+@json_utils.Encoder.register(Bounds)
+def encode_bounds(bounds):
     return {
-        "min": shape_bounds.min,
-        "opt": shape_bounds.opt,
-        "max": shape_bounds.max,
+        "min": bounds.min,
+        "opt": bounds.opt,
+        "max": bounds.max,
     }
 
 
-@json_utils.Decoder.register(ShapeBounds)
-def decode_shape_bounds(shape_bounds_dict):
-    return ShapeBounds(
-        min=tuple(shape_bounds_dict["min"]),
-        opt=tuple(shape_bounds_dict["opt"]),
-        max=tuple(shape_bounds_dict["max"]),
-    )
-
-
-@dataclass
-class ValueBounds:
-    min: Tuple[IntLike]
-    opt: Tuple[IntLike]
-    max: Tuple[IntLike]
-
-
-@json_utils.Encoder.register(ValueBounds)
-def encode_value_bounds(value_bounds):
-    return {
-        "min": tuple(value_bounds.min),
-        "opt": tuple(value_bounds.opt),
-        "max": tuple(value_bounds.max),
-    }
-
-
-@json_utils.Decoder.register(ValueBounds)
-def decode_value_bounds(value_bounds_dict):
-    return ValueBounds(
-        min=tuple(value_bounds_dict["min"]),
-        opt=tuple(value_bounds_dict["opt"]),
-        max=tuple(value_bounds_dict["max"]),
+@json_utils.Decoder.register(Bounds)
+def decode_bounds(bounds_dict):
+    return Bounds(
+        min=tuple(bounds_dict["min"]),
+        opt=tuple(bounds_dict["opt"]),
+        max=tuple(bounds_dict["max"]),
     )

tripy/nvtripy/backend/api/input_info.py

@@ -16,7 +16,7 @@
 
 from nvtripy import export
 from nvtripy.backend.api.named_dimension import NamedDimension
-from nvtripy.backend.api.shape_bounds import ShapeBounds, ValueBounds
+from nvtripy.backend.api.bounds import Bounds
 from nvtripy.frontend.dimension_size import DimensionSize
 from nvtripy.types import IntLike
 from nvtripy.utils import json as json_utils
@@ -97,7 +97,7 @@ def __init__(
         A mapping of dimension indices to their names, if set.
         """
 
-        self.shape_bounds: ShapeBounds = ShapeBounds(tuple(min_shape), tuple(opt_shape), tuple(max_shape))
+        self.shape_bounds: Bounds = Bounds(tuple(min_shape), tuple(opt_shape), tuple(max_shape))
         """
         The shape bounds of the input.
         """
@@ -152,7 +152,9 @@ def __init__(self, value_bounds: Tuple[IntLike, IntLike, IntLike]) -> None:
             assert dim_inp.opt == 2
             assert dim_inp.max == 3
         """
-        self.value_bounds = ValueBounds(
+        # Evaluate `DimensionSize` early to avoid duplicate evaluation
+        value_bounds = tuple(map(int, value_bounds))
+        self.value_bounds = Bounds(
             min=tuple([value_bounds[0]]), opt=tuple([value_bounds[1]]), max=tuple([value_bounds[2]])
         )
 
@@ -171,5 +173,6 @@ def encode_dim_input_info(dim_input_info):
 
 @json_utils.Decoder.register(DimensionInputInfo)
 def decode_dim_input_info(dim_input_info_dict):
-    dim_input_info_dict.value_bounds = dim_input_info_dict["value_bounds"]
-    return dim_input_info_dict
+    dim_input_info = DimensionInputInfo((-1, -1, -1))
+    dim_input_info.value_bounds = dim_input_info_dict["value_bounds"]
+    return dim_input_info

tripy/nvtripy/frontend/dimension_size.py

@@ -48,7 +48,9 @@ def __str__(self) -> str:
         return str(val)
 
     def eval(self) -> "nvtripy.Tensor":
+        from nvtripy.common import device
         from nvtripy.trace.ops.shape import GetDimensionSize, Shape
+        from nvtripy.frontend.ops.copy import copy
 
         # TODO (#593): Generalize this to any branchy graph:
         # If we find a pattern like Shape -> GetDimensionSize, we want to eval the Shape operation
@@ -62,4 +64,4 @@ def eval(self) -> "nvtripy.Tensor":
             dim_size.outputs[0].is_compile_tracer = self.trace_tensor.is_compile_tracer
             self.trace_tensor = dim_size.outputs[0]
 
-        return super().eval()
+        return copy(super().eval(), device("cpu"))

tripy/nvtripy/frontend/tensor.py

@@ -181,8 +181,6 @@ def eval(self) -> "nvtripy.Tensor":
         """
         Immediately evaluates this tensor. By default, tensors are evaluated lazily.
 
-        .. note:: The evaluated tensor will always be in **GPU memory**.
-
         Returns:
             The evaluated tensor.
 

tripy/nvtripy/trace/trace.py

@@ -227,7 +227,8 @@ def num_known_dims(ranked_tensor_type):
                                         for idx, name in input_info.dimension_names.items()
                                     }
                                 )
-                            elif isinstance(input_info, DimensionInputInfo):
+                            else:
+                                assert isinstance(input_info, DimensionInputInfo)
                                 value_bounds = input_info.value_bounds
                                 attr["tensorrt.value_bounds"] = ir.Attribute.parse(
                                     f"#tensorrt.shape_profile<min={list(value_bounds.min)}, opt={list(value_bounds.opt)}, max={list(value_bounds.max)}>"

tripy/tests/backend/api/conftest.py

@@ -12,6 +12,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import nvtripy as tp
+
+
 def add(a, b):
     return a + b
 
@@ -38,3 +41,8 @@ def variadic_positional(*args):
 
 def variadic_keyword(**kwargs):
     return sum(kwargs.values())
+
+
+def dynamic_reshape(a, b):
+    a = a + a
+    return tp.reshape(a, (-1, b))

tripy/tests/backend/api/test_compile.py

@@ -224,3 +224,20 @@ def func(a):
             tp.TripyException, match="Tensors that are not inputs to compiled functions must reside in CPU memory."
         ):
             tp.compile(func, args=[tp.InputInfo((2, 3), dtype=tp.float32)])
+
+    def test_dimension_input(self):
+        dummy = tp.ones((3, 4))
+        compiled = tp.compile(
+            dynamic_reshape,
+            args=[
+                tp.InputInfo(shape=((2, 4, 6), 4), dtype=tp.float32),
+                tp.DimensionInputInfo(value_bounds=(2, dummy.shape[1], 6)),
+            ],
+        )
+        for reshape_dim in [2, 4, 6]:
+            inp_cp = cp.arange(12, dtype=cp.float32).reshape((3, 4))
+            inp = tp.Tensor(inp_cp)
+            dim_inp = tp.DimensionSize(reshape_dim)
+            out = compiled(inp, dim_inp)
+            expected = (inp_cp + inp_cp).reshape((-1, reshape_dim))
+            assert cp.array_equal(cp.from_dlpack(out), expected)

tripy/tests/backend/api/test_input_info.py

@@ -70,10 +70,17 @@ def test_dimension_names(self):
         assert inp.shape_bounds.max == (3,)
         assert inp.dimension_names == {0: "batch"}
 
-    def test_serialize(self):
+    def test_serialize_input_info(self):
         batch = tp.NamedDimension("batch", 1, 2, 3)
         inp_info = tp.InputInfo(shape=[batch, 3, 28, 28], dtype=tp.float32)
 
         deserialized = json_utils.from_json(json_utils.to_json(inp_info))
 
         assert inp_info.__dict__ == deserialized.__dict__
+
+    def test_serialize_dim_input_info(self):
+        dim_inp_info = tp.DimensionInputInfo(value_bounds=(2, 4, 6))
+
+        deserialized = json_utils.from_json(json_utils.to_json(dim_inp_info))
+
+        assert dim_inp_info.__dict__ == deserialized.__dict__

tripy/tests/integration/test_dimension_input.py

@@ -229,8 +229,8 @@ def test_str_for_dynamic_shapes(self):
             == dedent(
                 r"""
                 def main(
-                    a : tensor<?xi32:gpu:0> : InputInfo<ShapeBounds(min=(2,), opt=(3,), max=(4,)), dimension names: {0: 'dim'}, dtype: int32>,
-                    b : tensor<?xi32:gpu:0> : InputInfo<ShapeBounds(min=(2,), opt=(3,), max=(4,)), dimension names: {}, dtype: int32>
+                    a : tensor<?xi32:gpu:0> : InputInfo<Bounds(min=(2,), opt=(3,), max=(4,)), dimension names: {0: 'dim'}, dtype: int32>,
+                    b : tensor<?xi32:gpu:0> : InputInfo<Bounds(min=(2,), opt=(3,), max=(4,)), dimension names: {}, dtype: int32>
                 ) -> (
                     c : tensor<?xi32:gpu:0>
                 ):