WIP: Add tp.DimensionInputInfo, support shape tensor input

Author: yizhuoz004

tripy/nvtripy/backend/api/compile.py

@@ -20,7 +20,7 @@
 
 from nvtripy import constants, export, utils
 from nvtripy.backend.api.executable import Executable
-from nvtripy.backend.api.input_info import InputInfo
+from nvtripy.backend.api.input_info import InputInfo, DimensionInputInfo
 from nvtripy.backend.mlir import Compiler
 from nvtripy.common.exception import raise_error
 from nvtripy.frontend import Tensor, Trace
@@ -164,6 +164,21 @@ def process_arg(name, arg):
             input_names.add(name)
 
             return tensor
+
+        if isinstance(arg, DimensionInputInfo):
+            from nvtripy.frontend.dimension_size import DimensionSize
+
+            tensor = DimensionSize(arg.value_bounds.opt[0])
+            tensor.name = name
+            tensor.trace_tensor.is_compile_tracer = True
+            assert tensor.trace_tensor.shape == ()
+
+            trace_input_map[name] = tensor
+            shapes.append(arg.value_bounds)
+            input_names.add(name)
+
+            return tensor
+
         return arg
 
     compiled_arg_names = []

tripy/nvtripy/backend/api/executable.py

@@ -190,17 +190,17 @@ def add(a, b):
                 ],
             )
 
-        for tensor in input_tensors:
-            producer = tensor.trace_tensor.producer
-            if not isinstance(producer, Constant) or tensor.device.kind != "gpu":
-                raise_error(
-                    "Inputs to compiled executables must be evaluated tensors on the GPU.",
-                    [
-                        "Got input" + (f" on device '{tensor.device}':" if tensor.device.kind != "gpu" else ":"),
-                        tensor,
-                        "Hint: Try calling `.eval()` on the tensor to ensure it is a GPU constant.",
-                    ],
-                )
+            # for tensor in input_tensors:
+            #     producer = tensor.trace_tensor.producer
+            # if not isinstance(producer, Constant) or tensor.device.kind != "gpu":
+            #     raise_error(
+            #         "Inputs to compiled executables must be evaluated tensors on the GPU.",
+            #         [
+            #             "Got input" + (f" on device '{tensor.device}':" if tensor.device.kind != "gpu" else ":"),
+            #             tensor,
+            #             "Hint: Try calling `.eval()` on the tensor to ensure it is a GPU constant.",
+            #         ],
+            #     )
 
         input_memrefs = [inp.trace_tensor.producer.data for inp in input_tensors]
         try:

tripy/nvtripy/backend/api/input_info.py

@@ -15,7 +15,7 @@
 from typing import Sequence, Tuple, Union
 
 from nvtripy import export
-from nvtripy.common.shape_bounds import ShapeBounds
+from nvtripy.common.shape_bounds import ShapeBounds, ValueBounds
 from nvtripy.frontend.dimension_size import DimensionSize
 from nvtripy.types import IntLike
 from nvtripy.utils import json as json_utils
@@ -57,7 +57,6 @@ def __init__(
         """
         is_int_like = lambda arg: any(isinstance(arg, typ) for typ in {int, DimensionSize})
 
-        # TODO (#252): Allow `shape` to be a shape tensor
         min_shape = []
         opt_shape = []
         max_shape = []
@@ -100,3 +99,44 @@ def decode_input_info(input_info_dict):
     input_info = InputInfo(shape=[], dtype=input_info_dict["dtype"])
     input_info.shape_bounds = input_info_dict["shape_bounds"]
     return input_info
+
+@export.public_api(document_under="compiling_code")
+class DimensionInputInfo:
+    """
+    Captures information about a dimension size input to a compiled function.
+    """
+
+    def __init__(self, value_bounds: Tuple[IntLike, IntLike, IntLike]) -> None:
+        """
+        Args:
+            value_bounds: The value bound of the dimension size input, consisting of minimum, optimum, and maximum values.
+
+        .. code-block:: python
+            :linenos:
+            :caption: Dynamic Dimensions
+
+            # The dimension size will support values in the range [1, 3],
+            # optimizing for a size of 2.
+            dim_inp = tp.DimensionInputInfo((1, 2, 3))
+            assert dim_inp.min == 1
+            assert dim_inp.opt == 2
+            assert dim_inp.max == 3
+        """
+        self.value_bounds = ValueBounds(min=tuple(value_bounds[0]), opt=tuple(value_bounds[1]), max=tuple(value_bounds[2]))
+
+    def __str__(self) -> str:
+        return (
+            f"DimensionInputInfo(min={self.value_bounds.min}, opt={self.value_bounds.opt}, max={self.value_bounds.max})"
+        )
+
+@json_utils.Encoder.register(DimensionInputInfo)
+def encode_input_info(dim_input_info):
+    return {
+        "value_bounds": dim_input_info.value_bounds,
+    }
+
+
+@json_utils.Decoder.register(DimensionInputInfo)
+def decode_input_info(dim_input_info_dict):
+    dim_input_info_dict.value_bounds = dim_input_info_dict["value_bounds"]
+    return dim_input_info_dict

tripy/nvtripy/common/shape_bounds.py

@@ -59,3 +59,26 @@ def decode_shape_bounds(shape_bounds_dict):
         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_shape_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_shape_bounds(shape_bounds_dict):
+    return ValueBounds(
+        min=tuple(shape_bounds_dict["min"]),
+        opt=tuple(shape_bounds_dict["opt"]),
+        max=tuple(shape_bounds_dict["max"]),
+    )

tripy/nvtripy/trace/trace.py

@@ -16,7 +16,7 @@
 #
 
 from textwrap import indent
-from typing import Dict, List, Optional, Sequence, Set
+from typing import Dict, List, Optional, Sequence, Set, Union
 
 from mlir_tensorrt.compiler import ir
 from mlir_tensorrt.compiler.dialects import func as func_dialect
@@ -29,7 +29,7 @@
     redirect_stderr,
 )
 from nvtripy.common.exception import raise_error
-from nvtripy.common.shape_bounds import ShapeBounds
+from nvtripy.common.shape_bounds import ShapeBounds, ValueBounds
 from nvtripy.logging import logger
 from nvtripy.trace.tensor import TraceTensor
 from nvtripy.trace.utils import topological_sort
@@ -44,7 +44,7 @@ def __init__(
         self,
         outputs: Sequence[TraceTensor],
         inputs: Sequence[TraceTensor] = [],
-        shapes: Optional[Sequence[ShapeBounds]] = None,
+        shapes: Optional[Sequence[Union[ShapeBounds, ValueBounds]]] = None,
         name: str = "main",
     ) -> None:
         """
@@ -199,9 +199,16 @@ def num_known_dims(ranked_tensor_type):
                     for idx in range(len(self.inputs)):
                         attr = {}
                         if self.shapes:
-                            attr["tensorrt.shape_profile"] = ir.Attribute.parse(
+                            attr_name = (
+                                "tensorrt.value_bounds"
+                                if isinstance(self.shapes[idx], ValueBounds)
+                                else "tensorrt.shape_profile"
+                            )
+                            attr[attr_name] = ir.Attribute.parse(
                                 f"#tensorrt.shape_profile<min={list(self.shapes[idx].min)}, opt={list(self.shapes[idx].opt)}, max={list(self.shapes[idx].max)}>"
                             )
+                        if idx == 1:
+                            attr["tensorrt.host_tensor"] = ir.UnitAttr.get()
 
                         arg_attrs.append(ir.DictAttr.get(attr))
 

tripy/test_out_of_bound.py

@@ -0,0 +1,45 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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
+import cupy as cp
+
+from nvtripy.logging import logger
+
+logger.verbosity = "ir"
+import mlir_tensorrt.runtime.api as runtime
+
+
+def func(x):
+    x = x + x
+    return x
+
+
+compiled_func = tp.compile(func, args=[tp.InputInfo(shape=((2, 4, 6), 4), dtype=tp.float32)])
+
+sig = compiled_func._executable_signature
+
+for idx in range(2):
+
+    arg = sig.get_arg(idx)
+    memref = runtime.MemRefType(arg)
+    print(f"Arg {idx}: ", memref.address_space)
+
+    print("Shape: ", memref.shape)
+    bound = sig.get_arg_bound(idx)
+    print(f"Bound: {bound.min()}, {bound.max()}")
+
+# inp = cp.ones((8, 4), dtype=cp.float32)
+# inp = tp.Tensor(inp)
+# out = compiled_func(inp)

tripy/test_shape_input.py

@@ -0,0 +1,60 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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
+import cupy as cp
+
+from nvtripy.logging import logger
+
+logger.verbosity = "ir"
+import mlir_tensorrt.runtime.api as runtime
+
+
+def func(x, y):
+    x = x + x
+    x = tp.reshape(x, (-1, y))
+    return x
+
+
+compiled_func = tp.compile(
+    func, args=[tp.InputInfo(shape=((2, 4, 6), 4), dtype=tp.float32), tp.DimensionInputInfo(value_bounds=(1, 2, 3))]
+)
+
+print("compilation complete.")
+
+sig = compiled_func._executable_signature
+
+input_info = compiled_func._get_input_info()
+print(input_info[0])
+print(input_info[1])
+for idx in range(2):
+
+    arg = sig.get_arg(idx)
+    memref = runtime.MemRefType(arg)
+    print(f"Arg {idx}: ", memref.address_space)
+
+    print("Shape: ", memref.shape)
+    bound = sig.get_arg_bound(idx)
+    print(f"Bound: {bound.min()}, {bound.max()}")
+
+
+# import pdb
+# pdb.set_trace()
+
+
+inp = cp.ones((4, 4), dtype=cp.float32)
+inp = tp.Tensor(inp)
+dim_inp = tp.DimensionSize(2)
+out = compiled_func(inp, dim_inp)
+print(out)