Replace class decorator with a function decorator

Author: jhalakpatel

tripy/tests/backend/mlir/test_compiler.py

@@ -66,6 +66,6 @@ def test_reason_context(self):
 
                 with pytest.raises(
                     TripyException,
-                    match=".*This is the first level of context\n    This is the second level of context\n.*",
+                    match=".*This is the first level of context\n    This is the second level of context.\n.*",
                 ) as exc:
                     map_error_to_user_code_and_raise(flat_ir, exc, err_str)

tripy/tests/flat_ir/test_constant_deduplication.py

@@ -53,12 +53,9 @@ def create_subgraph(config):
         # Create a function with no inputs and a single output
         func_result_tensor = FlatIRTensor.build(shape=[2], rank=1, dtype=int32, reason_details="", device=device("gpu"))
         setattr(result_tensor, "caller_tensor", func_result_tensor)
-        func = FlatIRFunction("MockFunc", [], [result_tensor])
+        func = FlatIRFunction("MockFunc", [], [result_tensor], [op1, op2, op3, mock_op])
         func_result_tensor.producer = func
 
-        # Insert all operations in a function
-        func.ops = [op1, op2, op3, mock_op]
-
         # Return function result tensor i.e. output of a function call
         return [], [func_result_tensor]
 

tripy/tests/flat_ir/test_function_deduplication.py

@@ -77,12 +77,11 @@ def create_function(
             flat_ir.register_tensor(callee_out)
             setattr(callee_out, "caller_tensor", original_out)
 
-        func = FlatIRFunction(name, [callee_input], callee_outputs)
         mock_op = MockOp([callee_input], [callee_outputs[0]])
         const_op = ConstantOp.build([], [callee_outputs[1]], data=[3, 4, 5])
         callee_outputs[1].producer = const_op
 
-        func.ops.extend([mock_op, const_op])
+        func = FlatIRFunction(name, [callee_input], callee_outputs, [mock_op, const_op])
         for out in output_tensors:
             out.producer = func
 

tripy/tripy/backend/mlir/utils.py

@@ -405,7 +405,7 @@ def find_inputs(graph_nodes):
             inputs = [subgraph.register_tensor(inp.to_flat_ir()) for inp in op.inputs]
             outputs = [subgraph.register_tensor(out.to_flat_ir()) for out in op.outputs]
             # Pass shallow copies of inputs/outputs so that the op is free to modify them
-            op.convert_to_flat_ir(copy.copy(inputs), copy.copy(outputs), subgraph)
+            op.to_flat_ir(copy.copy(inputs), copy.copy(outputs))
             subgraph.integrate_subgraph(inputs, outputs)
 
         mlir = subgraph.to_mlir()

tripy/tripy/flat_ir/function.py

@@ -21,13 +21,13 @@
 class FlatIRFunction:
     """Represents a function in the Flat IR."""
 
-    def __init__(self, name: str, inputs: List["FlatIRTensor"], outputs: List["FlatIRTensor"]):
+    def __init__(self, name: str, inputs: List["FlatIRTensor"], outputs: List["FlatIRTensor"], ops):
         """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.ops = ops
+        self.traced_ir_ops = ops  # Used only for function deduplication.
         self.trace_input_names = None
         self.trace_output_names = None
         self.caller_replacements = []
@@ -38,8 +38,7 @@ def clone_with_new_io(
         """
         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 = FlatIRFunction(self.name, new_inputs, new_outputs, self.ops)
         new_func.trace_input_names = self.trace_input_names
         new_func.trace_output_names = self.trace_output_names
         return new_func
@@ -54,13 +53,6 @@ def get_caller_inputs(self) -> List["FlatIRTensor"]:
     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 = [

tripy/tripy/frontend/trace/ops/base.py

@@ -201,26 +201,6 @@ def to_flat_ir(self, inputs: List["FlatIRTensor"], outputs: List["FlatIRTensor"]
         """
         ...
 
-    def convert_to_flat_ir(
-        self, inputs: List["FlatIRTensor"], outputs: List["FlatIRTensor"], flat_ir: Optional["FlatIR"] = None
-    ) -> None:
-        """
-        Convert the trace operation to Flat IR representation.
-
-        This method decides whether to call the wrapped or unwrapped version
-        of to_flat_ir based on the is_flat_ir_conversion_wrapped flag.
-
-        Args:
-            inputs: List of input FlatIRTensor objects
-            outputs: List of output FlatIRTensor objects
-            flat_ir: Optional FlatIR object for the wrapped version
-        """
-        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]:
         """
         Returns names of dataclass fields to skip when generating a string representation of the op.

tripy/tripy/frontend/trace/ops/binary_elementwise.py

@@ -28,7 +28,6 @@
 
 
 @dataclass(repr=False)
-@frontend_utils.wraps_to_flat_ir_to_func
 class BinaryElementwise(BaseTraceOp):
     class Kind:
         SUM = " + "
@@ -127,6 +126,7 @@ def broadcast_inputs(self, inputs):
 
         return [broadcasted_input_0, broadcasted_input_1]
 
+    @frontend_utils.make_function
     def to_flat_ir(self, inputs, outputs):
         from tripy.flat_ir.ops import AddOp, DivideOp, FloorOp, MaxOp, MinOp, MulOp, PowOp, SubtractOp
         from tripy.flat_ir.tensor import FlatIRTensor
@@ -193,7 +193,6 @@ def to_flat_ir(self, inputs, outputs):
 
 
 @dataclass(repr=False)
-@frontend_utils.wraps_to_flat_ir_to_func
 class Comparison(BinaryElementwise):
     class Kind:
         class KindElem(str):
@@ -217,6 +216,7 @@ def __new__(cls, content, compare_direction):
     def infer_dtypes(self):
         self.outputs[0].dtype = datatype.bool
 
+    @frontend_utils.make_function
     def to_flat_ir(self, inputs, outputs):
         from tripy.flat_ir.ops import CompareOp
 

tripy/tripy/frontend/trace/ops/cast.py

@@ -23,7 +23,6 @@
 
 
 @dataclass(repr=False)
-@frontend_utils.wraps_to_flat_ir_to_func
 class Cast(BaseTraceOp):
     dtype: "tripy.common.dtype"
 
@@ -46,6 +45,7 @@ def infer_shape_output_idxs(self, inputs):
     def infer_dtypes(self):
         self.outputs[0].dtype = self.dtype
 
+    @frontend_utils.make_function
     def to_flat_ir(self, inputs, outputs):
         from tripy.common.datatype import int32, int64, float32, bool as tp_bool
         from tripy.flat_ir.ops import CompareOp, ConvertOp, ConstantOp, DynamicBroadcastOp

tripy/tripy/frontend/trace/ops/dequantize.py

@@ -27,7 +27,6 @@
 
 
 @dataclass(repr=False)
-@frontend_utils.wraps_to_flat_ir_to_func
 class Dequantize(BaseTraceOp):
 
     dtype: datatype.dtype
@@ -36,6 +35,7 @@ class Dequantize(BaseTraceOp):
     def infer_dtypes(self):
         self.outputs[0].dtype = self.dtype
 
+    @frontend_utils.make_function
     def to_flat_ir(self, inputs, outputs):
         from tripy.common.datatype import int32
         from tripy.flat_ir.ops import (

tripy/tripy/frontend/trace/ops/fill.py

@@ -29,7 +29,6 @@
 
 
 @dataclass(repr=False)
-@frontend_utils.wraps_to_flat_ir_to_func
 class Fill(BaseTraceOp):
     value: float
     output_rank: int
@@ -55,6 +54,7 @@ def infer_rank(self):
             self.output_rank = input_shape[0]
         self.outputs[0].rank = self.output_rank
 
+    @frontend_utils.make_function
     def to_flat_ir(self, inputs, outputs):
         from tripy.flat_ir.ops import ConstantOp, ConvertOp, DynamicBroadcastOp
         from tripy.flat_ir.tensor import FlatIRTensor