Cache inputinfo structure at compile time, improve test coverage

Author: akhilg-nv

tripy/nvtripy/backend/api/compile.py

@@ -316,9 +316,22 @@ def process_arg(name, arg):
     assert isinstance(func_out, Tensor) or isinstance(
         func_out, Sequence
     ), "This function is only implemented for Tensors or sequences of Tensors"
+
+    # Group leaf input names by top-level argument for efficient runtime extraction
+    leaf_names_by_arg = {}
+    leaf_names = list(input_infos.keys())
+    for arg_name in compiled_arg_names:
+        matching = [
+            leaf
+            for leaf in leaf_names
+            if leaf == arg_name or leaf.startswith(f"{arg_name}.") or leaf.startswith(f"{arg_name}[")
+        ]
+        leaf_names_by_arg[arg_name] = matching
+
     return Executable(
         executable,
         compiled_arg_names,
         return_single_tensor_as_sequence=isinstance(func_out, Sequence),
         input_infos=input_infos,
+        leaf_names_by_arg=leaf_names_by_arg,
     )

tripy/nvtripy/backend/api/executable.py

@@ -46,6 +46,7 @@ def __init__(
         arg_names,
         return_single_tensor_as_sequence: bool,
         input_infos: Dict[str, Union[InputInfo, DimensionInputInfo]],
+        leaf_names_by_arg: Dict[str, Sequence[str]],
     ):
         self._executable = executable
 
@@ -78,6 +79,8 @@ def __init__(
         Stores metadata, like shapes and data types, for each input to the executable.
         """
 
+        self._leaf_names_by_arg = leaf_names_by_arg
+
     def __str__(self) -> str:
         params = [
             f"{name}: {str_from_type_annotation(param.annotation)}"
@@ -195,34 +198,36 @@ def add(a, b):
                 ],
             )
 
-        # Recursively build a name->tensor map
-        def extract_inputs(tensors, input_info_names):
-            name_to_tensor = {}
-
-            def extract_recursive(value, name_prefix):
-                if name_prefix in input_info_names:
-                    name_to_tensor[name_prefix] = value
-                    return
-                if isinstance(value, dict):
-                    for key, item in value.items():
-                        nested_name = f"{name_prefix}.{key}"
-                        extract_recursive(item, nested_name)
-                elif isinstance(value, (list, tuple)):
-                    for idx, item in enumerate(value):
-                        nested_name = f"{name_prefix}[{idx}]"
-                        extract_recursive(item, nested_name)
-                else:
-                    print(f"Leaf tensor: {name_prefix}: {value}")
-                    return
-
-            for name_idx, tensor in enumerate(tensors):
-                arg_name = self._arg_names[name_idx]
-                extract_recursive(tensor, arg_name)
-
-            return name_to_tensor
-
+        # Build a name->tensor map using precomputed leaf names to avoid unnecessary recursion
         input_info_names = list(self.input_infos.keys())
-        name_to_tensor = extract_inputs(input_tensors, set(input_info_names))
+        name_to_tensor: Dict[str, Tensor] = {}
+
+        def extract_recursive(value, name_prefix, allowed_names):
+            if name_prefix in allowed_names:
+                name_to_tensor[name_prefix] = value
+                return
+            if isinstance(value, dict):
+                for key, item in value.items():
+                    nested_name = f"{name_prefix}.{key}"
+                    extract_recursive(item, nested_name, allowed_names)
+            elif isinstance(value, (list, tuple)):
+                for idx, item in enumerate(value):
+                    nested_name = f"{name_prefix}[{idx}]"
+                    extract_recursive(item, nested_name, allowed_names)
+            else:
+                return
+
+        for name_idx, tensor in enumerate(input_tensors):
+            arg_name = self._arg_names[name_idx]
+            # Fast path: direct leaf input
+            if arg_name in self.input_infos:
+                name_to_tensor[arg_name] = tensor
+                continue
+            # If this arg has no compiled leaves beneath it, skip any recursion
+            allowed = self._leaf_names_by_arg.get(arg_name)
+            if not allowed:
+                continue
+            extract_recursive(tensor, arg_name, set(allowed))
         try:
             flattened_tensors = [name_to_tensor[name] for name in input_info_names]
         except KeyError as missing:
@@ -267,7 +272,7 @@ def extract_recursive(value, name_prefix):
                 expected_input_dtypes = [
                     info.dtype if isinstance(info, InputInfo) else int32 for info in self.input_infos.values()
                 ]
-                for tensor, dtype, arg_name in zip(input_tensors, expected_input_dtypes, self._arg_names):
+                for tensor, dtype, arg_name in zip(flattened_tensors, expected_input_dtypes, self.input_infos.keys()):
                     if tensor.dtype != dtype:
                         raise_error(
                             f"Unexpected tensor data type.",
@@ -282,16 +287,17 @@ def extract_recursive(value, name_prefix):
                 expected_input_shapes = [
                     info.shape_bounds if isinstance(info, InputInfo) else tuple() for info in self.input_infos.values()
                 ]
-                for tensor, expected_bounds, arg_name in zip(input_tensors, expected_input_shapes, self._arg_names):
+                for tensor, expected_bounds, arg_name in zip(
+                    flattened_tensors, expected_input_shapes, self.input_infos.keys()
+                ):
                     shape = tensor.shape
 
                     if len(shape) != len(expected_bounds.min):
                         raise_error(
                             f"Unexpected tensor rank.",
                             [
                                 f"For tensor: `{arg_name}`, expected a rank of: {len(expected_bounds.min)} but got: {len(shape)}.\n"
-                                f"Note: The provided argument was: ",
-                                tensor,
+                                f"Note: The provided argument was a tensor with shape: {shape}",
                             ],
                         )
 
@@ -302,8 +308,7 @@ def extract_recursive(value, name_prefix):
                                 [
                                     f"For tensor: `{arg_name}`, expected a shape within the bounds: min={expected_bounds.min}, max={expected_bounds.max}, but got: {shape}.\n"
                                     f"Dimension {i} has a shape of {shape[i]}, which is not within the expected bounds of [{expected_bounds.min[i]}, {expected_bounds.max[i]}].\n"
-                                    f"Note: The provided argument was: ",
-                                    tensor,
+                                    f"Note: The provided argument was a tensor with shape: {shape}",
                                 ],
                             )
             raise_error(str(err))

tripy/nvtripy/frontend/tensor.py

@@ -237,6 +237,7 @@ def eval(self) -> "nvtripy.Tensor":
                 name: InputInfo(list(map(int, inp.trace_tensor.shape)), inp.dtype)
                 for name, inp in zip(arg_names, inputs)
             },
+            leaf_names_by_arg={name: [name] for name in arg_names},  # every argument is a direct input
         )
         data = executable(*inputs).trace_tensor.producer.data
 

tripy/tests/backend/api/test_compile.py

@@ -270,7 +270,7 @@ def func(data_dict):
         }
         result = compiled_func(test_dict)
         expected = test_dict["a"]["inner"] + test_dict["b"]["list"][0] + test_dict["b"]["list"][1]
-        assert cp.array_equal(cp.from_dlpack(result), cp.from_dlpack(expected))
+        assert tp.equal(result, expected)
 
     def test_compile_nested_sequence_input_info(self):
         def func(data_list):
@@ -294,11 +294,9 @@ def func(data_list):
         ]
         result = compiled_func(test_list)
         expected = test_list[0] + test_list[1][0] + test_list[1][1]
-        assert cp.array_equal(cp.from_dlpack(result), cp.from_dlpack(expected))
+        assert tp.equal(result, expected)
 
     def test_compile_mixed_containers_and_constants(self):
-        """Test compilation with comprehensive mix: regular InputInfo, dict container, list container, and standalone constant."""
-
         def func(regular_input, data_dict, data_list, const_in_dict, const):
             return (
                 regular_input
@@ -315,19 +313,59 @@ def func(regular_input, data_dict, data_list, const_in_dict, const):
             "x": tp.InputInfo(shape=(2, 3), dtype=tp.float32),
             "y": tp.zeros((2, 3), dtype=tp.float32),
         }
-        list_input = [tp.InputInfo(shape=(2, 3), dtype=tp.float32), tp.ones((2, 3), dtype=tp.float32) * 3]
+        list_input = [tp.ones((2, 3), dtype=tp.float32) * 3, tp.InputInfo(shape=(2, 3), dtype=tp.float32)]
         const_in_dict = {"z": tp.ones((2, 3), dtype=tp.float32) * 5}
         const = tp.ones((2, 3), dtype=tp.float32) * 6
 
         compiled_func = tp.compile(func, args=[regular_input, dict_input, list_input, const_in_dict, const])
 
         # Only InputInfo arguments should be in function signature
         test_regular = tp.ones((2, 3), dtype=tp.float32).eval()
-        test_dict = {"x": (tp.ones((2, 3), dtype=tp.float32) * 2).eval(), "y": tp.zeros((2, 3), dtype=tp.float32)}
-        test_list = [(tp.ones((2, 3), dtype=tp.float32) * 4).eval(), tp.ones((2, 3), dtype=tp.float32) * 3]
+        test_dict = {"x": (tp.ones((2, 3), dtype=tp.float32) * 2).eval()}
+        test_list = [None, (tp.ones((2, 3), dtype=tp.float32) * 4).eval()]
 
         result = compiled_func(test_regular, test_dict, test_list)
         expected = (
-            test_regular + test_dict["x"] + test_dict["y"] + test_list[0] + test_list[1] + const_in_dict["z"] + const
+            test_regular + test_dict["x"] + dict_input["y"] + test_list[1] + list_input[0] + const_in_dict["z"] + const
         )
-        assert cp.array_equal(cp.from_dlpack(result), cp.from_dlpack(expected))
+        assert tp.equal(result, expected)
+
+    def test_compile_missing_nested_input_fails(self):
+        def func(data_dict):
+            return data_dict["a"]["inner"] + data_dict["b"]["list"][1]
+
+        dict_input = {
+            "a": {"inner": tp.InputInfo(shape=(2, 3), dtype=tp.float32)},
+            "b": {"list": [tp.zeros((2, 3), dtype=tp.float32), tp.InputInfo(shape=(2, 3), dtype=tp.float32)]},
+        }
+
+        compiled_func = tp.compile(func, args=[dict_input])
+
+        # Missing b.list[1]
+        bad_dict = {
+            "a": {"inner": tp.ones((2, 3), dtype=tp.float32).eval()},
+            "b": {"list": [tp.ones((2, 3), dtype=tp.float32).eval()]},
+        }
+        with helper.raises(tp.TripyException, match="Missing runtime tensor for input `data_dict\.b\.list\[1\]`."):
+            compiled_func(bad_dict)
+
+        # Wrong shape for b.list[1] should trigger a shape/device validation error
+        wrong_shape = {
+            "a": {"inner": tp.ones((2, 3), dtype=tp.float32).eval()},
+            "b": {"list": [tp.zeros((2, 3), dtype=tp.float32), tp.ones((2, 2), dtype=tp.float32).eval()]},
+        }
+        with helper.raises(tp.TripyException, match="Unexpected tensor shape."):
+            compiled_func(wrong_shape)
+
+    def test_compile_container_mismatch_fails(self):
+        def func(data_list):
+            return data_list[0] + data_list[1][0]
+
+        list_input = [tp.InputInfo(shape=(2, 3), dtype=tp.float32), [tp.InputInfo(shape=(2, 3), dtype=tp.float32)]]
+
+        compiled_func = tp.compile(func, args=[list_input])
+
+        bad_list = [tp.ones((2, 3), dtype=tp.float32).eval(), {"not": tp.ones((2, 3), dtype=tp.float32).eval()}]
+
+        with helper.raises(tp.TripyException, match="Missing runtime tensor for input `data_list\[1\]\[0\]`."):
+            compiled_func(bad_list)