Migrates Tripy to use the tensorrt dialect instead of StableHLO (#607)

Migrates Tripy to use the `tensorrt` dialect instead of StableHLO. The
latter is a lower level representation that the TensorRT network API,
meaning that we first break down high level Tripy operations and then
pattern match them back up. Using the `tensorrt` dialect allows us to go
directly to TensorRT without so many intervening layers of translation,
which reduces complexity and bugs, and improves performance.

Broadly, this change does the following:
- Removes FlatIR
- Make Trace responsible for mapping to MLIR
    - Trace operations represent MLIR operations 1:1
- Makes the frontend API responsible for composing Trace operations

---------

Signed-off-by: yizhuoz004 <[email protected]>
Co-authored-by: Jhalak Patel <[email protected]>
Co-authored-by: Yizhuo Zhang <[email protected]>

Author: 3 people

tripy/.gitignore

@@ -88,3 +88,6 @@ core*
 
 # PyTest profiling results
 *.prof
+
+# PyTorch checkpoints
+*.pt.*

tripy/CONTRIBUTING.md

@@ -70,6 +70,7 @@ pre-commit install
 
 We've written developer guides to help you understand the codebase:
 
+<!-- TODO (pranavm): Update links here -->
 - Start with the
     [architecture](https://nvidia.github.io/TensorRT-Incubator/post0_developer_guides/architecture.html)
     documentation.

tripy/README.md

@@ -45,7 +45,7 @@ guide for details:
         def __init__(self):
             self.conv = tp.Conv(in_channels=1, out_channels=1, kernel_dims=[3, 3])
 
-        def __call__(self, x):
+        def forward(self, x):
             x = self.conv(x)
             x = tp.relu(x)
             return x

tripy/docs/README.md

@@ -101,7 +101,7 @@ Guides are **markdown** files parsed by the Myst parser, therefore:
     - **Example:**
 
         ```md
-        [Fill operation](source:/nvtripy/trace/ops/fill.py)
+        [Broadcast operation](source:/nvtripy/trace/ops/broadcast.py)
         ```
 
     - **Why:** Other links will cause the file to be downloaded instead of linking to the repository.

tripy/docs/conf.py

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# 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");
@@ -27,8 +27,6 @@
 from tests import helper
 
 import nvtripy as tp
-from nvtripy.common.datatype import DATA_TYPES
-from nvtripy.utils.wrappers import TYPE_VERIFICATION
 
 PARAM_PAT = re.compile(":param .*?:")
 
@@ -179,10 +177,10 @@ def process_docstring_impl(app, what, name, obj, options, lines):
                     pname = "*" + pname
 
                 # Type annotations are optional for the `self` parameter unless the API has to be type-verified.
-                if pname != "self" or name in TYPE_VERIFICATION:
+                if pname != "self":
                     assert (
                         pname in documented_args
-                    ), f"Missing documentation for parameter: '{pname}' in: '{obj}'. Please ensure you've included this in the `Args:` section. Note: Documented parameters were: {documented_args} {doc}"
+                    ), f"Missing documentation for parameter: '{pname}' in: '{obj}'. Please ensure you've included this in the `Args:` section. Note: Documented parameters were: {documented_args}"
                     assert (
                         pname in documented_args
                     ), f"Missing documentation for parameter: '{pname}' in: '{obj}'. Please ensure you've included this in the `Args:` section. Note: Documented parameters were: {documented_args}"
@@ -206,69 +204,12 @@ def process_docstring_impl(app, what, name, obj, options, lines):
                     ":returns:" in doc
                 ), f"For: {obj}, return value is not documented. Please ensure you've included a `Returns:` section"
 
-    if name in TYPE_VERIFICATION:
-        add_text_index = -1
-        for index, block in enumerate(blocks):
-
-            def insert_block(text):
-                nonlocal index
-
-                blocks.insert(index, text)
-                index += 1
-
-            if re.search(r".. code-block::", block):
-                type_dict = TYPE_VERIFICATION[name].dtypes
-                insert_block("TYPE CONSTRAINTS:")
-                # Add the dtype constraint name and the dtypes that correlate.
-                for type_name, dt in type_dict.items():
-                    insert_block(
-                        f"    - **{type_name}**: :class:`"
-                        + "`, :class:`".join(
-                            sorted(
-                                set(dt),
-                                key=lambda dtype: (
-                                    tuple(typ.__name__ for typ in DATA_TYPES[dtype].__bases__),
-                                    DATA_TYPES[dtype].itemsize,
-                                ),
-                            )
-                        )
-                        + "`",
-                    )
-                insert_block("\n")
-
-                if TYPE_VERIFICATION[name].exceptions:
-                    # Add the dtype exceptions.
-                    insert_block("UNSUPPORTED TYPE COMBINATIONS:")
-                    for exception_dict in TYPE_VERIFICATION[name].exceptions:
-                        insert_block(
-                            "    - "
-                            + ", ".join([f"**{key}**\ =\ :class:`{val}`" for key, val in exception_dict.items()]),
-                        )
-                    insert_block("\n")
-                break
-
-            if re.search(r":param \w+: ", block):
-                param_name = re.match(r":param (\w+): ", block).group(1)
-                # Add dtype constraint to start of each parameter description.
-                if TYPE_VERIFICATION[name].constraints.get(param_name, None):
-                    add_text_index = re.search(r":param \w+: ", block).span()[1]
-                    blocks[index] = (
-                        f"{block[0:add_text_index]}[*dtype=*\ **{TYPE_VERIFICATION[name].constraints[param_name]}**\ ] {block[add_text_index:]}"
-                    )
-
-            if TYPE_VERIFICATION[name].return_dtype is not None and re.search(r":returns:", block):
-                add_text_index = re.search(r":returns:", block).span()[1] + 1
-                # Add dtype constraint to start of returns description.
-                blocks[index] = (
-                    f"{block[0:add_text_index]}[*dtype=*\ **{TYPE_VERIFICATION[name].return_dtype}**\ ] {block[add_text_index:]}"
-                )
-
     seen_classes.add(name)
 
     def allow_no_example():
-        # `tp.Module`s include examples in their constructors, so their __call__ methods don't require examples.
+        # `tp.Module`s include examples in their constructors, so their forward methods don't require examples.
         is_tripy_module_call_method = False
-        if what == "method" and obj.__name__ == "__call__":
+        if what == "method" and obj.__name__ == "forward":
             class_name = "nvtripy." + name.rpartition(".")[0]
             # Class names are prefixed with nvtripy.<...>, so we need to import it here to make eval() work.
             import nvtripy
@@ -336,7 +277,8 @@ def process_docstring(app, what, name, obj, options, lines):
     try:
         process_docstring_impl(app, what, name, obj, options, lines)
     except:
-        print(f"Error while processing {what}: {name} ({obj})")
+        sep = "\n"
+        print(f"Error while processing {what}: {name} ({obj}).\nNote: Docstring was: {sep.join(lines)}")
         raise
 
 

tripy/docs/generate_rsts.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# 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");
@@ -19,16 +19,15 @@
 import inspect
 import os
 import re
-import shutil
 import subprocess as sp
 from collections import defaultdict
 from dataclasses import dataclass
 from textwrap import dedent, indent
 from typing import Dict, List, Set
 
 import nvtripy as tp
-from tests import helper
 from nvtripy.export import PUBLIC_APIS
+from tests import helper
 
 
 @dataclass
@@ -93,7 +92,7 @@ def build_api_doc(api, include_heading=True):
     )
 
 
-def build_markdown_doc(source_path):
+def build_rst_markdown_include(source_path):
     return dedent(
         f"""
                 .. include:: {source_path}
@@ -181,7 +180,24 @@ def build_index_file(name, constituents, include_heading=True, caption=None):
     )
 
 
+def write_if_different(filepath, content):
+    if os.path.exists(filepath):
+        with open(filepath, "r") as f:
+            existing_content = f.read()
+        if existing_content == content:
+            print(f"No changes for {filepath}")
+            # Update the file timestamp so we know its up-to-date in timestamp checks.
+            os.utime(filepath, None)
+            return
+    with open(filepath, "w") as f:
+        f.write(content)
+    print(f"Updated {filepath}")
+
+
 def process_guide(guide_path: str, processed_guide_path: str):
+    if os.path.exists(processed_guide_path) and os.path.getmtime(guide_path) <= os.path.getmtime(processed_guide_path):
+        print(f"Skipping {guide_path} as it is not newer than {processed_guide_path}")
+        return
 
     os.makedirs(os.path.dirname(processed_guide_path), exist_ok=True)
 
@@ -194,7 +210,7 @@ def process_guide(guide_path: str, processed_guide_path: str):
     for index, block in enumerate(blocks):
         print(f"Processing block {index} (lang={block.lang}) in: {guide_path}: ", end="")
 
-        should_eval = not block.has_marker("doc: no_eval")
+        should_eval = not block.has_marker("doc: no_eval_or_format")
         if should_eval and block.lang.startswith("py"):
             print("Evaluating Python block")
 
@@ -229,6 +245,7 @@ def add_block(kind, contents, lang, code_indentation):
                     format_contents=add_block,
                     err_msg=f"Error while executing code block {index} (line {block.line_number}) from {guide_path}. ",
                     local_vars=code_locals,
+                    force_no_print_locals=block.has_marker("doc: no_print_locals"),
                 )
             )
 
@@ -254,8 +271,7 @@ def add_block(kind, contents, lang, code_indentation):
             else:
                 print("Omitting block")
 
-    with open(processed_guide_path, "w") as fout:
-        fout.write("\n".join(new_blocks))
+    write_if_different(processed_guide_path, "\n".join(new_blocks))
 
 
 def main():
@@ -272,8 +288,6 @@ def main():
 
     args = parser.parse_args()
 
-    shutil.rmtree(args.output, ignore_errors=True)
-
     def make_output_path(*components):
         return os.path.join(args.output, *components)
 
@@ -282,6 +296,7 @@ def make_output_path(*components):
     def is_file(document_under):
         return bool(os.path.splitext(document_under)[1])
 
+    rst_files = {}
     seen_apis = set()
     for api in PUBLIC_APIS:
         name = get_name(api)
@@ -314,9 +329,12 @@ def is_file(document_under):
 
         os.makedirs(os.path.dirname(rst_path), exist_ok=True)
 
-        pre_existing_file = os.path.exists(rst_path)
-        with open(rst_path, "a") as f:
-            f.write(build_api_doc(api, include_heading=not pre_existing_file))
+        # Only include a heading if we didn't already create the file:
+        include_heading = False
+        if rst_path not in rst_files:
+            include_heading = True
+            rst_files[rst_path] = ""
+        rst_files[rst_path] += build_api_doc(api, include_heading=include_heading)
 
     def str_from_hierarchy(obj):
         if isinstance(obj, dict):
@@ -374,13 +392,11 @@ def str_from_hierarchy(obj):
 
             # Do not create RSTs for top-level README
             if not is_top_level_dir:
-                with open(guide_out_path, "w") as f:
-                    # Grab the path of the .md file relative to the directory containing the .rst file so we can include it correctly.
-                    f.write(
-                        build_markdown_doc(
-                            source_path=os.path.relpath(processed_guide, os.path.dirname(guide_out_path))
-                        )
-                    )
+                content = build_rst_markdown_include(
+                    source_path=os.path.relpath(processed_guide, os.path.dirname(guide_out_path))
+                )
+                rst_files[guide_out_path] = content
+
             guides.append(os.path.join(basename, os.path.splitext(guide_filename)[0]))
 
         # We have special treatment for the files in the top-level directory.
@@ -391,24 +407,33 @@ def str_from_hierarchy(obj):
         is_root = path == ""
 
         index_path = make_output_path(path, "index.rst")
-        pre_existing_file = os.path.exists(index_path)
+        pre_existing_file = index_path in rst_files
 
         if is_root:
             assert (
                 not pre_existing_file
             ), f"APIs should *not* target the root index file directly! Please remove any `document_under='index.rst'` arguments!"
-
-        with open(index_path, "a") as f:
-            f.write(
-                build_root_index_file(constituents, guide_sets, processed_markdown_dirname)
-                if is_root
-                else build_index_file(
-                    name=to_title(os.path.basename(os.path.dirname(index_path))),
-                    constituents=constituents,
-                    include_heading=not pre_existing_file,
-                    caption="See also:" if pre_existing_file else None,
-                )
+            content = build_root_index_file(constituents, guide_sets, processed_markdown_dirname)
+        else:
+            content = build_index_file(
+                name=to_title(os.path.basename(os.path.dirname(index_path))),
+                constituents=constituents,
+                include_heading=not pre_existing_file,
+                caption="See also:" if pre_existing_file else None,
             )
+        if not pre_existing_file:
+            rst_files[index_path] = ""
+        rst_files[index_path] += content
+
+    # Delete any existing RST files that are not in rst_files
+    for existing_file in glob.glob(os.path.join(args.output, "**", "*.rst"), recursive=True):
+        if existing_file not in rst_files:
+            print(f"Removing stale file: {existing_file}")
+            os.remove(existing_file)
+
+    # Write the new/updated RST files
+    for path, content in rst_files.items():
+        write_if_different(path, content)
 
 
 if __name__ == "__main__":