Enable direct fast access of inputs inside containers

Author: akhilg-nv

tripy/nvtripy/backend/api/compile.py

@@ -157,6 +157,7 @@ def add(a, b):
     input_names = set()
     input_infos = {}
     trace_inputs = []  # flattened list of trace input tensors in argument order
+    access_plan_by_name: Dict[str, tuple] = {}
 
     # Set up names for the weights in the module to make the trace easier to read.
     if isinstance(func, Module):
@@ -207,18 +208,22 @@ def process_arg_input_info(name, arg):
 
         return arg
 
-    def process_arg_and_flag(name, arg):
+    def process_arg_and_flag(top_arg_name, name, arg, steps):
         # Handle individual InputInfo or DimensionInputInfo objects
         if isinstance(arg, (InputInfo, DimensionInputInfo)):
-            return process_arg_input_info(name, arg), True
+            tensor_or_dim = process_arg_input_info(name, arg)
+            access_plan_by_name[name] = (top_arg_name, tuple(steps))
+            return tensor_or_dim, True
 
         # Handle containers of InputInfo objects
         if isinstance(arg, dict):
             result = {}
             has_input = False
             for key, value in arg.items():
                 nested_name = f"{name}.{key}"
-                processed_child, child_has_input = process_arg_and_flag(nested_name, value)
+                processed_child, child_has_input = process_arg_and_flag(
+                    top_arg_name, nested_name, value, (*steps, str(key))
+                )
                 result[key] = processed_child
                 has_input = has_input or child_has_input
             return result, has_input
@@ -227,15 +232,15 @@ def process_arg_and_flag(name, arg):
             has_input = False
             for idx, value in enumerate(arg):
                 nested_name = f"{name}[{idx}]"
-                processed_child, child_has_input = process_arg_and_flag(nested_name, value)
+                processed_child, child_has_input = process_arg_and_flag(top_arg_name, nested_name, value, (*steps, idx))
                 result_list.append(processed_child)
                 has_input = has_input or child_has_input
             return type(arg)(result_list), has_input  # preserve sequence type
 
         return arg, False
 
     def process_arg(name, arg):
-        processed, has_input = process_arg_and_flag(name, arg)
+        processed, has_input = process_arg_and_flag(name, name, arg, ())
         if has_input:
             input_names.add(name)
         return processed
@@ -317,21 +322,10 @@ def process_arg(name, arg):
         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,
+        access_plan_by_name=access_plan_by_name,
     )

tripy/nvtripy/backend/api/executable.py

@@ -46,7 +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]],
+        access_plan_by_name: Dict[str, Tuple[str, Tuple[Union[str, int], ...]]],
     ):
         self._executable = executable
 
@@ -79,7 +79,23 @@ def __init__(
         Stores metadata, like shapes and data types, for each input to the executable.
         """
 
-        self._leaf_names_by_arg = leaf_names_by_arg
+        # Build accessor map from compile-time access plans
+        self._accessor_map: Dict[str, callable] = {}
+        name_to_index = {name: idx for idx, name in enumerate(self._arg_names)}
+
+        def make_accessor(arg_index: int, steps: Tuple[Union[str, int], ...]):
+            def accessor(inputs, idx=arg_index, stps=steps):
+                v = inputs[idx]
+                for s in stps:
+                    v = v[s]
+                return v
+
+            return accessor
+
+        self._access_plan_by_name = access_plan_by_name
+        for leaf_name, (arg_name, steps) in self._access_plan_by_name.items():
+            idx = name_to_index[arg_name]
+            self._accessor_map[leaf_name] = make_accessor(idx, steps)
 
     def __str__(self) -> str:
         params = [
@@ -198,46 +214,20 @@ def add(a, b):
                 ],
             )
 
-        # Build a name->tensor map using precomputed leaf names to avoid unnecessary recursion
+        # Fetch flattened tensors directly via accessors
         input_info_names = list(self.input_infos.keys())
-        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:
-            raise_error(
-                f"Missing runtime tensor for input `{missing.args[0]}`.",
-                [
-                    "Ensure your provided containers include tensors for all compiled inputs.",
-                    f"Expected inputs: {input_info_names}",
-                ],
-            )
+        flattened_tensors = []
+        for name in input_info_names:
+            try:
+                flattened_tensors.append(self._accessor_map[name](input_tensors))
+            except Exception:
+                raise_error(
+                    f"Missing runtime tensor for input `{name}`.",
+                    [
+                        "Ensure your provided collections include tensors for all compiled inputs.",
+                        f"Expected inputs: {input_info_names}",
+                    ],
+                )
         expected_devices = ["gpu" if isinstance(info, InputInfo) else "cpu" for info in self.input_infos.values()]
 
         # Validate flattened tensors against input_infos
@@ -398,7 +388,7 @@ def encode_executable(executable):
         "executable": base64.b64encode(executable._executable.serialize()).decode(),
         "_return_single_tensor_as_sequence": executable._return_single_tensor_as_sequence,
         "input_infos": executable.input_infos,
-        "leaf_names_by_arg": executable._leaf_names_by_arg,
+        "access_plan_by_name": executable._access_plan_by_name,
     }
 
 
@@ -410,5 +400,5 @@ def decode_executable(executable_dict):
         executable_dict["arg_names"],
         return_single_tensor_as_sequence=executable_dict["_return_single_tensor_as_sequence"],
         input_infos=executable_dict["input_infos"],
-        leaf_names_by_arg=executable_dict.get("leaf_names_by_arg"),
+        access_plan_by_name=executable_dict["access_plan_by_name"],
     )

tripy/nvtripy/frontend/tensor.py

@@ -237,7 +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
+            access_plan_by_name={name: (name, tuple()) for name in arg_names},
         )
         data = executable(*inputs).trace_tensor.producer.data