Makes various optimizations

Applies various optimizations based on data from NSight traces. For example,
this includes removing all assertions in the critical path and reducing the
amount of Python code involved as much as possible.

Author: pranavm-nvidia

tripy/tests/performance/cases/linear_block.py

@@ -28,9 +28,11 @@
 )
 def linear_block(tripy_dtype, torch_dtype):
 
+    NUM_LAYERS = 15
+
     class LinearBlock(tp.Module):
         def __init__(self):
-            self.layers = [tp.Linear(256, 256, bias=False, dtype=tripy_dtype) for _ in range(10)]
+            self.layers = [tp.Linear(256, 256, bias=False, dtype=tripy_dtype) for _ in range(NUM_LAYERS)]
             for layer in self.layers:
                 # Adjust the weights to prevent FP16 overflows:
                 weight = np.tile(np.array([[-1, 1], [1, -1]], dtype=TRIPY_TO_NUMPY[tripy_dtype]), (128, 128))
@@ -47,7 +49,7 @@ def __init__(self):
             self.layers = torch.nn.ModuleList(
                 [
                     torch.nn.Linear(256, 256, bias=False, dtype=torch_dtype, device=torch.device("cuda"))
-                    for _ in range(10)
+                    for _ in range(NUM_LAYERS)
                 ]
             )
 

tripy/tripy/backend/api/executable.py

@@ -25,6 +25,7 @@
 from tripy.common.exception import raise_error
 from tripy.frontend import Tensor
 from tripy.utils import json as json_utils
+from tripy.utils.stack_info import StackInfo
 
 
 @export.public_api(document_under="compiling_code")
@@ -41,6 +42,7 @@ def __init__(self, executable, arg_names, output_devices):
         self._executable = executable
         self._executor = Executor(self._executable)
         self._arg_names = arg_names
+        self._num_expected_args = len(arg_names)
         self._output_devices = output_devices
         self._executable_signature = self._executable.get_signature("main")
 
@@ -88,12 +90,12 @@ def add(a, b):
 
             out = compiled_add(a, b)
         """
-        NUM_ARGS = len(args) + len(kwargs)
+        num_positional = len(args)
+        NUM_ARGS = num_positional + len(kwargs)
 
-        input_tensors = []
-        input_tensors.extend(args)
+        input_tensors = list(args)
         # Need to get arguments in the order of self._arg_names, which may be different from kwargs ordering.
-        expected_kwargs = self._arg_names[len(args) :]
+        expected_kwargs = self._arg_names[num_positional:]
         for name in expected_kwargs:
             if name not in kwargs:
                 raise_error(f"Missing argument: {name}", [f"Expected the following arguments: {self._arg_names}"])
@@ -106,16 +108,17 @@ def add(a, b):
                 f"Extra keyword arguments: {list(kwargs.keys())}",
                 [
                     f"Expected the following arguments: {self._arg_names}.\n"
-                    f"Note: The following arguments were already provided as positional arguments: {self._arg_names[:len(args)]}"
+                    f"Note: The following arguments were already provided as positional arguments: {self._arg_names[:num_positional]}"
                 ],
             )
 
         # We do this after kwarg checks since those will be more informative (we can explain which arguments are missing/extra).
-        if NUM_ARGS != len(self._arg_names):
+
+        if NUM_ARGS != self._num_expected_args:
             raise_error(
                 "Incorrect number of arguments.",
                 [
-                    f"Expected {len(self._arg_names)} arguments but got {NUM_ARGS}.\n"
+                    f"Expected {self._num_expected_args} arguments but got {NUM_ARGS}.\n"
                     f"Note: Expected arguments were: {self._arg_names}",
                 ],
             )
@@ -158,8 +161,6 @@ def add(a, b):
                             )
             raise
 
-        from tripy.utils.stack_info import StackInfo
-
         output_tensors = [Tensor(output, stack_info=StackInfo([])) for output in executor_outputs]
         if len(output_tensors) == 1:
             output_tensors = output_tensors[0]

tripy/tripy/backend/mlir/executor.py

@@ -17,31 +17,34 @@
 
 from typing import List
 
-import mlir_tensorrt.compiler.api as compiler
 import mlir_tensorrt.runtime.api as runtime
 
+from tripy.backend.api.stream import default_stream
 from tripy.backend.mlir.memref import create_empty_memref
+from tripy.backend.mlir.utils import MLIRRuntimeClient, convert_runtime_dtype_to_tripy_dtype
 from tripy.backend.utils import TensorInfo
 from tripy.common import datatype, device
 from tripy.common.exception import raise_error
+from tripy.common.utils import convert_list_to_array
 from tripy.utils import make_tuple
 
 
 class Executor:
     def __init__(self, executable: runtime.Executable) -> None:
-        from tripy.backend.api.stream import default_stream
-        from tripy.backend.mlir.utils import MLIRRuntimeClient
-
         self.runtime_client = MLIRRuntimeClient()
         session_options = runtime.RuntimeSessionOptions(num_devices=1, device_id=0)
         self.session = runtime.RuntimeSession(session_options, executable)
         self.device = self.runtime_client.get_devices()[0]  # Assume a single device is available.
         self.signature = executable.get_signature("main")
         self.stream = default_stream()
+        self.num_input_args = self.signature.get_num_input_args()
+        self.num_output_args = self.signature.get_num_output_args()
+        self.output_args = [
+            self.signature.get_arg(index + self.num_input_args) for index in range(self.num_output_args)
+        ]
+        self.output_memrefs = [runtime.MemRefType(out) for out in self.output_args]
 
     def _create_shape_memref(self, shape):
-        from tripy.common.utils import convert_list_to_array
-
         shape = make_tuple(shape)
         if len(shape) == 0:
             # create an empty memref
@@ -55,34 +58,21 @@ def _create_shape_memref(self, shape):
             stream=self.stream._active_cuda_stream,
         )
 
-    def _get_inputs_runtime_shape(self, inputs):
-        inputs_shape = []
-        for input in inputs:
-            inputs_shape.append(input.trace_tensor.producer.data.shape)
-        return inputs_shape
-
     def _get_outputs_shape(self):
-        offset = self.signature.get_num_input_args()
         outputs_shape = []
         all_outputs_known = True
-        for output_index in range(self.signature.get_num_output_args()):
-            arg_index = output_index + offset
-            arg = self.signature.get_arg(arg_index)
-            assert compiler.MemRefType.isinstance(arg)
-            memref = runtime.MemRefType(arg)
-            rank = len(memref.shape)
-
+        for memref in self.output_memrefs:
             outputs_shape.append(memref.shape)
-            if rank > 0:
-                all_outputs_known &= all(dim >= 0 for dim in memref.shape)
+            all_outputs_known &= all(dim >= 0 for dim in memref.shape)
         return outputs_shape, all_outputs_known
 
-    def _execute_shape_inference(self, inputs_shape, outputs_shape):
-        # Only execute shape inference if shape function name is valid.
-        assert (
-            self.signature.get_shape_func_name()
-        ), f"Shape inference function is missing while output shapes are not known."
+    def _get_inputs_runtime_shape(self, inputs):
+        inputs_shape = []
+        for input in inputs:
+            inputs_shape.append(input.trace_tensor.producer.data.shape)
+        return inputs_shape
 
+    def _execute_shape_inference(self, inputs_shape, outputs_shape):
         inputs_shape_memref = [self._create_shape_memref(inp_shape) for inp_shape in inputs_shape]
         outputs_shape_memref = [self._create_shape_memref(out_shape) for out_shape in outputs_shape]
         self.session.execute_function(
@@ -93,41 +83,24 @@ def _execute_shape_inference(self, inputs_shape, outputs_shape):
         return outputs_runtime_shape
 
     def _get_output_tensor_info(self, outputs_runtime_shape, output_devices):
-        from tripy.backend.mlir.utils import convert_runtime_dtype_to_tripy_dtype
-
-        offset = self.signature.get_num_input_args()
         outputs_tensor_info = []
-        for output_index in range(self.signature.get_num_output_args()):
-            arg_index = output_index + offset
-            arg = self.signature.get_arg(arg_index)
-            assert compiler.MemRefType.isinstance(arg) or compiler.ScalarType.isinstance(
-                arg
-            ), "Argument must be either MemRefType or ScalarType"
-            assert compiler.MemRefType.isinstance(
-                arg
-            ), "ScalarType argument are not yet supported"  # 158: Add scalar type output argument support.
-            memref = compiler.MemRefType(arg)
+        for index in range(self.num_output_args):
+            memref = self.output_memrefs[index]
             dtype = convert_runtime_dtype_to_tripy_dtype(memref.dtype)
-            device_type = "gpu" if memref.address_space == runtime.PointerType.device else "cpu"
-            if output_devices[output_index]:
-                device_type = output_devices[output_index].kind
-            is_static_shape = all(dim >= 0 for dim in memref.shape)
-            if is_static_shape:
-                outputs_tensor_info.append(
-                    TensorInfo(len(memref.shape), tuple(memref.shape), dtype, device(device_type))
-                )
-            else:
-                runtime_shape = [
-                    rs if dim < 0 else dim for dim, rs in zip(memref.shape, outputs_runtime_shape[output_index])
-                ]
-                outputs_tensor_info.append(
-                    TensorInfo(
-                        len(runtime_shape),
-                        tuple(runtime_shape),
-                        dtype,
-                        device(device_type),
-                    )
+
+            output_device = output_devices[index]
+            if not output_device:
+                output_device = device(("gpu" if memref.address_space == runtime.PointerType.device else "cpu", 0))
+
+            runtime_shape = [rs if dim < 0 else dim for dim, rs in zip(memref.shape, outputs_runtime_shape[index])]
+            outputs_tensor_info.append(
+                TensorInfo(
+                    len(runtime_shape),
+                    tuple(runtime_shape),
+                    dtype,
+                    output_device,
                 )
+            )
         return outputs_tensor_info
 
     def get_output_tensor_runtime_info(self, inputs, output_devices=List[device]):

tripy/tripy/backend/mlir/memref.py

@@ -41,7 +41,7 @@ def _cached_create_empty_memref(shape: Sequence[int], dtype: str, device_kind: s
 def create_empty_memref(
     shape: Sequence[int],
     dtype: str,
-    device: tp_device = tp_device("gpu"),
+    device: tp_device = tp_device(("gpu", 0)),
     stream=None,
     use_cache: bool = True,
 ):

tripy/tripy/backend/mlir/utils.py

@@ -288,15 +288,17 @@ def redirect_stderr() -> BinaryIO:
 
 
 def convert_tripy_dtype_to_runtime_dtype(dtype: datatype.dtype) -> runtime.ScalarTypeCode:
-    if dtype not in TRIPY_DTYPE_TO_MLIR_TRT:
+    try:
+        return TRIPY_DTYPE_TO_MLIR_TRT[dtype]
+    except KeyError:
         raise_error(f"Data type: '{dtype}' does not have a corresponding runtime data type")
-    return TRIPY_DTYPE_TO_MLIR_TRT.get(dtype)
 
 
 def convert_runtime_dtype_to_tripy_dtype(dtype: runtime.ScalarTypeCode) -> datatype.dtype:
-    if dtype not in MLIR_TRT_TO_TRIPY_DTYPE:
+    try:
+        return MLIR_TRT_TO_TRIPY_DTYPE[dtype]
+    except KeyError:
         raise_error(f"Data type: '{dtype}' does not have a corresponding tripy data type")
-    return MLIR_TRT_TO_TRIPY_DTYPE.get(dtype)
 
 
 def is_any_dim_dynamic(mlir_tensor):

tripy/tripy/common/device.py

@@ -22,6 +22,8 @@
 from tripy.common.exception import TripyException
 from tripy.utils.json import Decoder, Encoder
 
+_VALID_KINDS = {"cpu", "gpu"}
+
 
 @export.public_api()
 @dataclass
@@ -60,27 +62,30 @@ def __init__(self, device: str) -> None:
             assert gpu_1.kind == "gpu"
             assert gpu_1.index == 1
         """
-
-        kind, _, index = device.partition(":")
-        kind = kind.lower()
-
-        if index:
-            try:
-                index = int(index)
-            except ValueError:
-                raise TripyException(f"Could not interpret: {index} as an integer")
-        else:
-            index = 0
-
-        if index < 0:
-            raise TripyException(f"Device index must be a non-negative integer, but was: {index}")
-
-        VALID_KINDS = {"cpu", "gpu"}
-        if kind not in VALID_KINDS:
-            raise TripyException(f"Unrecognized device kind: {kind}. Choose from: {list(VALID_KINDS)}")
-
-        self.kind = kind
-        self.index = index
+        try:
+            # Fast constructor for the critical path. If a Tuple[str, int] is provided, then
+            # we bypass all the logic to parse the information from a string.
+            self.kind, self.index = device
+        except ValueError:
+            kind, _, index = device.partition(":")
+            kind = kind.lower()
+
+            if index:
+                try:
+                    index = int(index)
+                except ValueError:
+                    raise TripyException(f"Could not interpret: {index} as an integer")
+            else:
+                index = 0
+
+            if index < 0:
+                raise TripyException(f"Device index must be a non-negative integer, but was: {index}")
+
+            if kind not in _VALID_KINDS:
+                raise TripyException(f"Unrecognized device kind: {kind}. Choose from: {list(_VALID_KINDS)}")
+
+            self.kind = kind
+            self.index = index
 
     def __str__(self) -> str:
         return f"{self.kind}:{self.index}"

tripy/tripy/frontend/tensor.py

@@ -18,6 +18,8 @@
 from textwrap import indent
 from typing import Any, List, Optional, Sequence, Union
 
+import mlir_tensorrt.runtime.api as runtime
+
 # Import ops to populate the registry before we define our Tensor class
 import tripy.frontend.ops
 import tripy.frontend.trace.ops
@@ -27,8 +29,7 @@
 from tripy.common.exception import raise_error
 from tripy.frontend.ops.registry import TENSOR_METHOD_REGISTRY
 from tripy.frontend.trace.ops import Storage
-
-import mlir_tensorrt.runtime.api as runtime
+from tripy.frontend.trace.tensor import TraceTensor
 
 
 class TensorMeta(type):
@@ -93,7 +94,6 @@ def __init__(
 
             tensor = tp.Tensor([1.0, 2.0, 3.0], dtype=tp.float32)
         """
-        from tripy.frontend.trace.tensor import TraceTensor
 
         # We include code for everything above the `BaseTraceOp.build` function, which is called at most
         # this many stack frames above the constructor.
@@ -120,17 +120,6 @@ def __init__(
         else:
             Storage.build_internal([], [self.trace_tensor], data, dtype, device)
 
-        # Storage should populate attrs of trace_tensor
-        assert all(
-            attr is not None
-            for attr in [
-                self.trace_tensor.shape,
-                self.trace_tensor.dtype,
-                self.trace_tensor.device,
-                self.trace_tensor.producer,
-            ]
-        )
-
         # Explicit cast if necessary
         # TODO(#155): Add copy as well when host allocation is fixed
         #             Also make device as a property, similar to dtype
@@ -172,13 +161,13 @@ def rank(self):
         return self.trace_tensor.rank
 
     def eval(self) -> runtime.MemRefValue:
+        if isinstance(self.trace_tensor.producer, Storage) and self.trace_tensor.producer.has_memref:
+            return self.trace_tensor.producer.data
+
         from tripy.backend.mlir.compiler import Compiler
         from tripy.backend.mlir.executor import Executor
         from tripy.frontend.trace import Trace
 
-        if isinstance(self.trace_tensor.producer, Storage) and self.trace_tensor.producer.has_memref:
-            return self.trace_tensor.producer.data
-
         trace = Trace([self])
         flat_ir = trace.to_flat_ir()
         mlir = flat_ir.to_mlir()