Updates Tripy to work with Python 3.9, moves files into utils

- Updates the `Dockerfile` to use a Python 3.9 container
     (the lowest version we support) so we can be sure we're not relying
     on features that are not available there.

     This unforunately means that we need to drop `lldb` from the container
     since it depends on Python 3.10. However, `gdb` is still available.

- Moves various files into `utils`. We were previously unable to do this
     because the `__init__.py` file of `utils` would import lots of things,
     creating circular dependencies. Now, the top-level `utils` module will
     not import everything, meaning there is no risk of circular dependencies.

- Makes some minor changes in the code to make Python 3.9 work (e.g. using
     `._name` instead of `.__name__`)

Author: pranavm-nvidia

tripy/Dockerfile

@@ -1,10 +1,10 @@
-FROM ubuntu:22.04
+FROM python:3.9
 
 LABEL org.opencontainers.image.description="Tripy development container"
 
 WORKDIR /tripy
 
-SHELL ["/bin/bash", "-c"]
+ENTRYPOINT ["/bin/bash"]
 
 # Setup user account
 ARG uid=1000
@@ -30,21 +30,5 @@ RUN pip install build .[docs,dev,test,build] \
     --extra-index-url https://download.pytorch.org/whl \
     --extra-index-url https://pypi.nvidia.com
 
-# Installl lldb for debugging purposes in Tripy container.
-# The LLVM version should correspond on LLVM_VERSION specified in https://github.com/NVIDIA/TensorRT-Incubator/blob/main/mlir-tensorrt/build_tools/docker/Dockerfile#L30.
-ARG LLVM_VERSION=17
-ENV LLVM_VERSION=$LLVM_VERSION
-ENV LLVM_PACKAGES="lldb-${LLVM_VERSION}"
-RUN echo "deb http://apt.llvm.org/jammy/ llvm-toolchain-jammy-$LLVM_VERSION main" > /etc/apt/sources.list.d/llvm.list && \
-    echo "deb-src http://apt.llvm.org/jammy/ llvm-toolchain-jammy-$LLVM_VERSION main" >> /etc/apt/sources.list.d/llvm.list && \
-    wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key 2>/dev/null > /etc/apt/trusted.gpg.d/apt.llvm.org.asc && \
-    wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null | gpg --dearmor - | tee /usr/share/keyrings/kitware-archive-keyring.gpg >/dev/null && \
-    echo 'deb [signed-by=/usr/share/keyrings/kitware-archive-keyring.gpg] https://apt.kitware.com/ubuntu/ jammy main' | tee /etc/apt/sources.list.d/kitware.list >/dev/null && \
-    apt-get update && \
-    apt-get install -y ${LLVM_PACKAGES} && \
-    apt-get clean -y && \
-    rm -rf /var/lib/apt/lists/* && \
-    ln -s /usr/bin/lldb-17 /usr/bin/lldb
-
 # Export tripy into the PYTHONPATH so it doesn't need to be installed after making changes
 ENV PYTHONPATH=/tripy

tripy/docs/README.md

@@ -45,7 +45,7 @@ which specifies doc metadata for each API (e.g. location).
 - Docstring must include *at least* **one [code example](#code-examples)**.
 
 - If the function accepts `tp.Tensor`s, must indicate **data type constraints**
-    with the [`wrappers.interface`](../nvtripy/wrappers.py) decorator.
+    with the [`wrappers.interface`](../nvtripy/utils/wrappers.py) decorator.
 
 **Example:**
 

tripy/docs/conf.py

@@ -28,7 +28,7 @@
 
 import nvtripy as tp
 from nvtripy.common.datatype import DATA_TYPES
-from nvtripy.wrappers import TYPE_VERIFICATION
+from nvtripy.utils.wrappers import TYPE_VERIFICATION
 
 PARAM_PAT = re.compile(":param .*?:")
 

tripy/docs/post0_developer_guides/how-to-add-new-ops.md

@@ -176,7 +176,8 @@ it as a `nvtripy.Module` under [`frontend/module`](source:/nvtripy/frontend/modu
 
 ```py
 # doc: no-eval
-from nvtripy import export, wrappers
+from nvtripy import export
+from nvtripy.utils import wrappers
 from nvtripy.types import ShapeLike
 
 # We can use the `export.public_api()` decorator to automatically export this

tripy/nvtripy/backend/api/compile.py

@@ -153,7 +153,7 @@ def process_arg(name, arg):
         return arg
 
     new_args = []
-    positional_arg_info, _ = utils.get_positional_arg_names(func, *args)
+    positional_arg_info, _ = utils.utils.get_positional_arg_names(func, *args)
     for name, arg in positional_arg_info:
         new_args.append(process_arg(name, arg))
 
@@ -165,7 +165,7 @@ def process_arg(name, arg):
     # as `InputInfo`s, but the order needs to match the signature of the original function.
     compiled_arg_names = [name for name in signature.parameters.keys() if name in input_names]
 
-    trace_outputs = utils.make_list(func(*new_args, **new_kwargs))
+    trace_outputs = utils.utils.make_list(func(*new_args, **new_kwargs))
 
     if not trace_outputs:
         raise_error(

tripy/nvtripy/backend/api/executable.py

@@ -14,18 +14,17 @@
 # limitations under the License.
 import base64
 import inspect
-from typing import Sequence, Union, Tuple
+from dataclasses import dataclass
+from typing import Sequence, Tuple, Union
 
 import mlir_tensorrt.runtime.api as runtime
-
 from nvtripy import export
 from nvtripy.backend.mlir import Executor
 from nvtripy.backend.mlir import utils as mlir_utils
 from nvtripy.common.exception import raise_error
 from nvtripy.frontend import Tensor
-from nvtripy.function_registry import str_from_type_annotation
 from nvtripy.utils import json as json_utils
-from dataclasses import dataclass
+from nvtripy.utils.types import str_from_type_annotation
 
 
 # TODO(MLIR-TRT #923): Can generalize `InputInfo` and drop this class.

tripy/nvtripy/backend/mlir/compiler.py

@@ -76,7 +76,7 @@ def compile_stabehlo_program(self, code: str) -> compiler.Executable:
             return compiler.compiler_stablehlo_to_executable(self.compiler_client, module.operation, opts)
 
     # The optional flat_ir parameter is used to generate nicer error messages.
-    @utils.log_time
+    @utils.utils.log_time
     def compile(self, mlir_module: ir.Module, flat_ir: Optional["FlatIR"] = None) -> compiler.Executable:
         logger.mlir(lambda: f"{mlir_module.operation.get_asm(large_elements_limit=32)}\n")
         opts = self._make_mlir_opts(self.trt_builder_opt_level)

tripy/nvtripy/backend/mlir/executor.py

@@ -18,15 +18,14 @@
 from typing import List
 
 import mlir_tensorrt.runtime.api as runtime
-
 from nvtripy.backend.api.stream import default_stream
 from nvtripy.backend.mlir.memref import create_memref
 from nvtripy.backend.mlir.utils import MLIRRuntimeClient, convert_runtime_dtype_to_tripy_dtype
 from nvtripy.backend.utils import TensorInfo
 from nvtripy.common import datatype, device
 from nvtripy.common.exception import raise_error
 from nvtripy.common.utils import convert_list_to_array
-from nvtripy.utils import make_tuple
+from nvtripy.utils.utils import make_tuple
 
 
 class Executor:

tripy/nvtripy/backend/mlir/memref.py

@@ -15,14 +15,12 @@
 # limitations under the License.
 #
 
-import math
 import re
 
 import mlir_tensorrt.runtime.api as runtime
-
 from nvtripy.backend.mlir import utils as mlir_utils
 from nvtripy.common import device as tp_device
-from nvtripy.utils import raise_error
+from nvtripy.common.exception import raise_error
 
 EMPTY_MEMREF_CACHE = {}
 

tripy/nvtripy/backend/mlir/utils.py

@@ -368,7 +368,7 @@ def get_flat_ir_operation(output_names):
         return (
             [
                 "This error occured while trying to compile the following FlatIR expression:",
-                utils.code_pretty_str(str(op)),
+                utils.utils.code_pretty_str(str(op)),
                 "\n",
             ]
             + (