Address review comments

jhalakpatel · jhalakpatel · commit aa378072885b · 2024-10-02T16:35:37.000-07:00
diff --git a/tripy/tests/flat_ir/ops/test_broadcast.py b/tripy/tests/flat_ir/ops/test_broadcast.py
@@ -33,6 +33,6 @@ def test_str(self):
         broadcast = func_broadcast.ops[-1]
         assert isinstance(broadcast, DynamicBroadcastOp)
         assert re.match(
-            r"t_inter3: \[rank=\(2\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(t_inter4, t_inter2, broadcast_dim=\[\]\)",
+            r"t_inter[0-9]+: \[rank=\(2\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(t_inter[0-9]+, t_inter[0-9]+, broadcast_dim=\[\]\)",
             str(broadcast),
         )
diff --git a/tripy/tests/flat_ir/ops/test_gather.py b/tripy/tests/flat_ir/ops/test_gather.py
@@ -19,7 +19,7 @@
 import tripy as tp
 
 from tripy.flat_ir.ops import DynamicGatherOp
-from tripy.flat_ir.ops.base import FlatIRFunction
+from tripy.flat_ir.function import FlatIRFunction
 from tripy.frontend.trace import Trace
 import re
 
diff --git a/tripy/tests/flat_ir/ops/test_maximum.py b/tripy/tests/flat_ir/ops/test_maximum.py
@@ -18,8 +18,8 @@
 import re
 import tripy as tp
 from tripy.frontend.trace import Trace
+from tripy.flat_ir.function import FlatIRFunction
 from tripy.flat_ir.ops import MaxOp
-from tripy.flat_ir.ops.base import FlatIRFunction
 
 
 class TestMaxOp:
diff --git a/tripy/tests/flat_ir/ops/test_reduce.py b/tripy/tests/flat_ir/ops/test_reduce.py
@@ -17,8 +17,8 @@
 
 import tripy as tp
 from tripy.frontend.trace import Trace
+from tripy.flat_ir.function import FlatIRFunction
 from tripy.flat_ir.ops import ArgMinMaxOp, ConvertOp, DivideOp, DynamicBroadcastOp, MulOp, ReduceOp
-from tripy.flat_ir.ops.base import FlatIRFunction
 import re
 
 
diff --git a/tripy/tests/flat_ir/ops/test_subtract.py b/tripy/tests/flat_ir/ops/test_subtract.py
@@ -18,8 +18,8 @@
 import re
 import tripy as tp
 from tripy.frontend.trace import Trace
+from tripy.flat_ir.function import FlatIRFunction
 from tripy.flat_ir.ops import SubtractOp
-from tripy.flat_ir.ops.base import FlatIRFunction
 
 
 class TestSubtractOp:
diff --git a/tripy/tests/flat_ir/test_constant_deduplication.py b/tripy/tests/flat_ir/test_constant_deduplication.py
@@ -15,7 +15,7 @@
 
 import pytest
 from tripy.flat_ir.flat_ir import FlatIR
-from tripy.flat_ir.ops.base import FlatIRFunction
+from tripy.flat_ir.function import FlatIRFunction
 from tripy.flat_ir.ops import ConstantOp
 from tripy.flat_ir.tensor import FlatIRTensor
 from tripy.common.device import device
diff --git a/tripy/tests/flat_ir/test_function_deduplication.py b/tripy/tests/flat_ir/test_function_deduplication.py
@@ -18,7 +18,8 @@
 from typing import List, Optional
 
 from tripy.flat_ir.flat_ir import FlatIR
-from tripy.flat_ir.ops.base import FlatIRFunction, BaseFlatIROp
+from tripy.flat_ir.function import FlatIRFunction
+from tripy.flat_ir.ops.base import BaseFlatIROp
 from tripy.flat_ir.ops import ConstantOp
 from tripy.flat_ir.tensor import FlatIRTensor
 from tripy.common.device import device
@@ -39,7 +40,7 @@ def __eq__(self, other):
         return True
 
     def to_mlir(self, operands):
-        assert "Not implemented"
+        raise NotImplementedError()
 
 
 def test_is_structurally_equivalent():
diff --git a/tripy/tripy/backend/mlir/memref.py b/tripy/tripy/backend/mlir/memref.py
@@ -26,7 +26,7 @@
 
 
 @lru_cache(maxsize=None)
-def _cached_create_memref(shape: Sequence[int], dtype: str, device_kind: str, stream):
+def _cached_create_empty_memref(shape: Sequence[int], dtype: str, device_kind: str, stream):
     mlirtrt_device = mlir_utils.MLIRRuntimeClient().get_devices()[0] if device_kind == "gpu" else None
     mlir_dtype = mlir_utils.convert_tripy_dtype_to_runtime_dtype(dtype)
     return mlir_utils.MLIRRuntimeClient().create_memref(
@@ -55,9 +55,10 @@ def create_empty_memref(
 
     """
     if use_cache:
-        return _cached_create_memref(tuple(shape), dtype, device.kind, stream)
+        assert common_utils.is_shape_empty(shape)
+        return _cached_create_empty_memref(tuple(shape), dtype, device.kind, stream)
     else:
-        return _cached_create_memref.__wrapped__(tuple(shape), dtype, device.kind, stream)
+        return _cached_create_empty_memref.__wrapped__(tuple(shape), dtype, device.kind, stream)
 
 
 def create_memref_view(data):
diff --git a/tripy/tripy/common/utils.py b/tripy/tripy/common/utils.py
@@ -74,6 +74,19 @@ def is_empty(data: Sequence) -> bool:
     return isinstance(data, Sequence) and all(map(is_empty, data))
 
 
+def is_shape_empty(shape: Sequence[int]) -> bool:
+    """
+    A shape is considered empty if any of its dimensions is zero.
+
+    Args:
+        shape (Tuple[int, ...]): A tuple representing the shape of a tensor.
+
+    Returns:
+        bool: True if the shape represents an empty tensor, False otherwise.
+    """
+    return any(dim == 0 for dim in shape)
+
+
 class Float16MemoryView:
     """
     A custom memory view class for handling float16 data.
diff --git a/tripy/tripy/flat_ir/flat_ir.py b/tripy/tripy/flat_ir/flat_ir.py
@@ -43,7 +43,8 @@ def __init__(self, shapes: Sequence[ShapeBounds] = None):
 
     def __str__(self) -> str:
         """Generate a string representation of the FlatIR."""
-        from tripy.flat_ir.ops.base import FlatIRFunction, BaseFlatIROp
+        from tripy.flat_ir.function import FlatIRFunction
+        from tripy.flat_ir.ops.base import BaseFlatIROp
 
         ir = []
 
@@ -96,7 +97,8 @@ def to_mlir(self):
         from mlir_tensorrt.compiler import ir
         from mlir_tensorrt.compiler.dialects import func as func_dialect
         from tripy.backend.mlir.utils import make_ir_context, make_tensor_location
-        from tripy.flat_ir.ops.base import FlatIRFunction, BaseFlatIROp
+        from tripy.flat_ir.function import FlatIRFunction
+        from tripy.flat_ir.ops.base import BaseFlatIROp
 
         def _base_op_to_mlir(op, mlir_tensor_map):
             op_inputs = [mlir_tensor_map[input_tensor.name] for input_tensor in op.inputs]
@@ -376,7 +378,8 @@ def integrate_subgraph(self, inputs: List["FlatIRTensor"], outputs: List["FlatIR
         """
         Integrate a subgraph delineated by the given inputs and outputs into this FlatIR.
         """
-        from tripy.flat_ir.ops.base import BaseFlatIROp, FlatIRFunction
+        from tripy.flat_ir.function import FlatIRFunction
+        from tripy.flat_ir.ops.base import BaseFlatIROp
         from tripy.flat_ir.tensor import FlatIRTensor
         from tripy.flat_ir.ops import ConstantOp
 
diff --git a/tripy/tripy/flat_ir/function.py b/tripy/tripy/flat_ir/function.py
@@ -0,0 +1,117 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import List
+
+from tripy.flat_ir.ops.base import BaseFlatIROp
+from tripy import utils
+
+
+class FlatIRFunction:
+    """Represents a function in the Flat IR."""
+
+    def __init__(self, name: str, inputs: List["FlatIRTensor"], outputs: List["FlatIRTensor"]):
+        """Initialize a FlatIRFunction."""
+        self.name = name
+        self.inputs = inputs
+        self.outputs = outputs
+        self.ops: List[BaseFlatIROp] = []
+        self.traced_ir_ops: List[BaseFlatIROp] = []  # Used for flat ir function deduplication
+        self.trace_input_names = None
+        self.trace_output_names = None
+        self.caller_replacements = []
+
+    def clone_with_new_io(
+        self, new_inputs: List["FlatIRTensor"], new_outputs: List["FlatIRTensor"]
+    ) -> "FlatIRFunction":
+        """
+        Create a clone of the function with new inputs and outputs.
+        """
+        new_func = FlatIRFunction(self.name, new_inputs, new_outputs)
+        new_func.ops = self.ops
+        new_func.trace_input_names = self.trace_input_names
+        new_func.trace_output_names = self.trace_output_names
+        return new_func
+
+    def set_caller_inputs(self, inputs: List["FlatIRTensor"]) -> None:
+        for callee_input, caller_input in zip(self.inputs, inputs):
+            callee_input.caller_tensor = caller_input
+
+    def get_caller_inputs(self) -> List["FlatIRTensor"]:
+        return [getattr(inp, "caller_tensor") for inp in self.inputs]
+
+    def get_caller_outputs(self) -> List["FlatIRTensor"]:
+        return [getattr(out, "caller_tensor") for out in self.outputs]
+
+    def add_op(self, op: BaseFlatIROp) -> None:
+        """
+        Add an operation to the function ops. `trace_ir_ops` are used only for function deduplication.
+        """
+        self.ops.append(op)
+        self.traced_ir_ops.append(op)
+
+    def __str__(self) -> str:
+        """Generate a string representation of the function."""
+        function_signature = [
+            f"function {self.name}(",
+            *[f"    {inp}" for inp in self.inputs],
+            ") -> (",
+            *[f"    {out}" for out in self.outputs],
+            ") {",
+        ]
+        function_body = [f"    {op}" for op in self.ops]
+        function_return = [f"    return {', '.join(out.name for out in self.outputs)}", "}"]
+        return "\n".join(function_signature + function_body + function_return)
+
+    def __repr__(self) -> str:
+        """Generate a concise string representation of the function."""
+        return f"<FlatIRFunction '{self.name}' with {len(self.inputs)} inputs, {len(self.outputs)} outputs, and {len(self.ops)} ops>"
+
+    def is_structurally_equivalent(self, other: "FlatIRFunction") -> bool:
+        """Check if two FlatIRFunction objects are structurally equivalent."""
+
+        def tensor_signature(tensor: "FlatIRTensor") -> tuple:
+            return (tensor.shape, tensor.dtype, tensor.rank, tensor.device)
+
+        def op_signature(op: "BaseFlatIROp") -> tuple:
+            from tripy.flat_ir.ops import ConstantOp
+
+            return (
+                type(op),
+                [tensor_signature(t) for t in op.inputs],
+                [tensor_signature(t) for t in op.outputs],
+                tuple(
+                    getattr(op, field.name)
+                    for field in utils.get_dataclass_fields(op, BaseFlatIROp)
+                    if field.name not in ("inputs", "outputs")
+                ),
+            )
+
+        # Check high-level structure
+        if (
+            len(self.inputs) != len(other.inputs)
+            or len(self.outputs) != len(other.outputs)
+            or len(self.traced_ir_ops) != len(other.traced_ir_ops)
+        ):
+            return False
+
+        # Check input and output signatures
+        if any(
+            tensor_signature(t1) != tensor_signature(t2)
+            for t1, t2 in zip(self.inputs + self.outputs, other.inputs + other.outputs)
+        ):
+            return False
+
+        # Check ops and its input and output signatures
+        return all(op_signature(op1) == op_signature(op2) for op1, op2 in zip(self.traced_ir_ops, other.traced_ir_ops))
diff --git a/tripy/tripy/flat_ir/ops/base.py b/tripy/tripy/flat_ir/ops/base.py
@@ -102,102 +102,3 @@ def _op_name(self) -> str:
             The name of this operation.
         """
         return self.__class__.__name__
-
-
-class FlatIRFunction(abc.ABC):
-    """Represents a function in the Flat IR."""
-
-    def __init__(self, name: str, inputs: List["FlatIRTensor"], outputs: List["FlatIRTensor"]):
-        """Initialize a FlatIRFunction."""
-        self.name = name
-        self.inputs = inputs
-        self.outputs = outputs
-        self.ops: List[BaseFlatIROp] = []
-        self.traced_ir_ops: List[BaseFlatIROp] = []  # Used for flat ir function deduplication
-        self.trace_input_names = None
-        self.trace_output_names = None
-        self.caller_replacements = []
-
-    def clone_with_new_io(
-        self, new_inputs: List["FlatIRTensor"], new_outputs: List["FlatIRTensor"]
-    ) -> "FlatIRFunction":
-        """
-        Create a clone of the function with new inputs and outputs.
-        """
-        new_func = FlatIRFunction(self.name, new_inputs, new_outputs)
-        new_func.ops = self.ops
-        new_func.trace_input_names = self.trace_input_names
-        new_func.trace_output_names = self.trace_output_names
-        return new_func
-
-    def set_caller_inputs(self, inputs: List["FlatIRTensor"]) -> None:
-        for callee_input, caller_input in zip(self.inputs, inputs):
-            callee_input.caller_tensor = caller_input
-
-    def get_caller_inputs(self) -> List["FlatIRTensor"]:
-        return [getattr(inp, "caller_tensor") for inp in self.inputs]
-
-    def get_caller_outputs(self) -> List["FlatIRTensor"]:
-        return [getattr(out, "caller_tensor") for out in self.outputs]
-
-    def add_op(self, op: BaseFlatIROp) -> None:
-        """
-        Add an operation to the function ops. `trace_ir_ops` are used only for function deduplication.
-        """
-        self.ops.append(op)
-        self.traced_ir_ops.append(op)
-
-    def __str__(self) -> str:
-        """Generate a string representation of the function."""
-        function_signature = [
-            f"function {self.name}(",
-            *[f"    {inp}" for inp in self.inputs],
-            ") -> (",
-            *[f"    {out}" for out in self.outputs],
-            ") {",
-        ]
-        function_body = [f"    {op}" for op in self.ops]
-        function_return = [f"    return {', '.join(out.name for out in self.outputs)}", "}"]
-        return "\n".join(function_signature + function_body + function_return)
-
-    def __repr__(self) -> str:
-        """Generate a concise string representation of the function."""
-        return f"<FlatIRFunction '{self.name}' with {len(self.inputs)} inputs, {len(self.outputs)} outputs, and {len(self.ops)} ops>"
-
-    def is_structurally_equivalent(self, other: "FlatIRFunction") -> bool:
-        """Check if two FlatIRFunction objects are structurally equivalent."""
-
-        def tensor_signature(tensor: "FlatIRTensor") -> tuple:
-            return (tensor.shape, tensor.dtype, tensor.rank, tensor.device)
-
-        def op_signature(op: "BaseFlatIROp") -> tuple:
-            from tripy.flat_ir.ops import ConstantOp
-
-            return (
-                type(op),
-                [tensor_signature(t) for t in op.inputs],
-                [tensor_signature(t) for t in op.outputs],
-                tuple(
-                    getattr(op, field.name)
-                    for field in utils.get_dataclass_fields(op, BaseFlatIROp)
-                    if field.name not in ("inputs", "outputs")
-                ),
-            )
-
-        # Check high-level structure
-        if (
-            len(self.inputs) != len(other.inputs)
-            or len(self.outputs) != len(other.outputs)
-            or len(self.traced_ir_ops) != len(other.traced_ir_ops)
-        ):
-            return False
-
-        # Check input and output signatures
-        if any(
-            tensor_signature(t1) != tensor_signature(t2)
-            for t1, t2 in zip(self.inputs + self.outputs, other.inputs + other.outputs)
-        ):
-            return False
-
-        # Check ops and its input and output signatures
-        return all(op_signature(op1) == op_signature(op2) for op1, op2 in zip(self.traced_ir_ops, other.traced_ir_ops))
diff --git a/tripy/tripy/frontend/trace/ops/base.py b/tripy/tripy/frontend/trace/ops/base.py
@@ -38,9 +38,6 @@ class BaseTraceOp(abc.ABC):
     outputs: List["TraceTensor"]
     """The output tensors of this operation"""
 
-    is_to_flat_ir_wrapped = False
-    """Flag indicating whether the to_flat_ir method has been wrapped"""
-
     @classmethod
     def build_internal(
         cls, inputs: List["TraceTensor"], outputs: List["TraceTensor"], *args, **kwargs
@@ -218,10 +215,11 @@ def convert_to_flat_ir(
             outputs: List of output FlatIRTensor objects
             flat_ir: Optional FlatIR object for the wrapped version
         """
-        if self.is_to_flat_ir_wrapped:
-            self.to_flat_ir(inputs, outputs, flat_ir)
-        else:
+        try:
             self.to_flat_ir(inputs, outputs)
+        except TypeError as e:
+            assert "to_flat_ir() missing 1 required positional argument: 'flat_ir'" in str(e)
+            self.to_flat_ir(inputs, outputs, flat_ir)
 
     def str_skip_fields(self) -> Set[str]:
         """
diff --git a/tripy/tripy/frontend/utils.py b/tripy/tripy/frontend/utils.py

Original file line number	Diff line number	Diff line change
`@@ -33,6 +33,6 @@ def test_str(self):`
`33`	`33`	`broadcast = func_broadcast.ops[-1]`
`34`	`34`	`assert isinstance(broadcast, DynamicBroadcastOp)`
`35`	`35`	`assert re.match(`
`36`		`- r"t_inter3: \[rank=\(2\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(t_inter4, t_inter2, broadcast_dim=\[\]\)",`
	`36`	`+ r"t_inter[0-9]+: \[rank=\(2\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(t_inter[0-9]+, t_inter[0-9]+, broadcast_dim=\[\]\)",`
`37`	`37`	`str(broadcast),`
`38`	`38`	`)`