Add flat ir function

Author: jhalakpatel

tripy/tests/backend/mlir/test_compiler.py

@@ -57,10 +57,11 @@ def test_reason_context(self):
                 b = tp.ones((1,))
                 trace = Trace([a + b])
                 flat_ir = trace.to_flat_ir()
-                producer = flat_ir.outputs[0].producer.inputs[0]
+                func_binary = flat_ir.outputs[0].producer
+                producer = func_binary.ops[-1].inputs[0]
                 flat_ir_inputs = ",".join(map(lambda i: i.name, producer.producer.inputs))
-                trace_inputs = ",".join(producer.producer.trace_input_names)
-                trace_output = producer.producer.trace_output_names[0]
+                trace_inputs = ",".join(func_binary.trace_input_names)
+                trace_output = func_binary.trace_output_names[0]
                 err_str = f'loc("{flat_ir_inputs};;<out>;;{producer.name};;<trace_in>;;{trace_inputs};;<trace_out>;;{trace_output}"): Test error'
 
                 with pytest.raises(

tripy/tests/flat_ir/ops/test_broadcast.py

@@ -29,9 +29,10 @@ def test_str(self):
         trace = Trace([out])
         flat_ir = trace.to_flat_ir()
 
-        broadcast = flat_ir.ops[-1]
+        func_broadcast = flat_ir.ops[-1]
+        broadcast = func_broadcast.ops[-1]
         assert isinstance(broadcast, DynamicBroadcastOp)
         assert re.match(
-            r"out: \[rank=\(2\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(t_inter2, t[0-9]+, broadcast_dim=\[\]\)",
+            r"t_inter3: \[rank=\(2\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(t_inter4, t_inter2, broadcast_dim=\[\]\)",
             str(broadcast),
         )

tripy/tests/flat_ir/ops/test_divide.py

@@ -31,20 +31,21 @@ def test_str(self):
         trace = Trace([out])
         flat_ir = trace.to_flat_ir()
 
-        div = flat_ir.ops[-1]
-        broadcast_a = flat_ir.ops[-3]
-        broadcast_b = flat_ir.ops[-2]
+        func_div = flat_ir.ops[-1]
+        div = func_div.ops[-1]
+        broadcast_a = func_div.ops[-3]
+        broadcast_b = func_div.ops[-2]
         assert isinstance(div, DivideOp)
 
         assert re.match(
-            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(a, t_inter[0-9]+, broadcast_dim=\[0\]\)",
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(t_inter[0-9]+, t_inter[0-9]+, broadcast_dim=\[0\]\)",
             str(broadcast_a),
         )
         assert re.match(
-            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(b, t_inter[0-9]+, broadcast_dim=\[0\]\)",
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(t_inter[0-9]+, t_inter[0-9]+, broadcast_dim=\[0\]\)",
             str(broadcast_b),
         )
         assert re.match(
-            r"out: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DivideOp\(t_inter[0-9]+, t_inter[0-9]+\)",
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DivideOp\(t_inter[0-9]+, t_inter[0-9]+\)",
             str(div),
         )

tripy/tests/flat_ir/ops/test_gather.py

@@ -19,6 +19,7 @@
 import tripy as tp
 
 from tripy.flat_ir.ops import DynamicGatherOp
+from tripy.flat_ir.ops.base import FlatIRFunction
 from tripy.frontend.trace import Trace
 import re
 
@@ -35,12 +36,15 @@ def test_gather_str(self, axis):
         trace = Trace([out])
         flat_ir = trace.to_flat_ir()
 
-        gather = flat_ir.ops[-1]
-        reshape = flat_ir.ops[-2]
+        func_gather = flat_ir.ops[-1]
+        assert isinstance(func_gather, FlatIRFunction)
+
+        gather = func_gather.ops[-1]
+        reshape = func_gather.ops[-7]
         print(str(reshape))
         assert isinstance(gather, DynamicGatherOp)
         assert re.match(
-            rf"out: \[rank=\(3\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicGatherOp\(data, indices, t_inter[0-9]+, axis={axis}\)",
+            rf"t_inter[0-9]+: \[rank=\(3\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicGatherOp\(t_inter[0-9]+, t_inter[0-9]+, t_inter[0-9]+, axis={axis}\)",
             str(gather),
         )
 
@@ -51,7 +55,7 @@ def test_gather_mlir(self, axis):
         flat_ir = trace.to_flat_ir()
         mlir_text = str(flat_ir.to_mlir())
         if axis == 0:
-            target = '"stablehlo.dynamic_gather"(%c, %c_0, %2) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>}> : (tensor<2x3xi32>, tensor<1xi32>, tensor<2xi32>) -> tensor<1x?xi32>'
+            target = '"stablehlo.dynamic_gather"(%arg0, %arg1, %2) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>}> : (tensor<2x3xi32>, tensor<1xi32>, tensor<2xi32>) -> tensor<1x?xi32>'
         else:
-            target = '"stablehlo.dynamic_gather"(%c, %c_0, %2) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1], index_vector_dim = 1>}> : (tensor<2x3xi32>, tensor<1xi32>, tensor<2xi32>) -> tensor<?x1xi32>'
+            target = '"stablehlo.dynamic_gather"(%arg0, %arg1, %2) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1], index_vector_dim = 1>}> : (tensor<2x3xi32>, tensor<1xi32>, tensor<2xi32>) -> tensor<?x1xi32>'
         assert target in mlir_text, mlir_text

tripy/tests/flat_ir/ops/test_maximum.py

@@ -19,6 +19,7 @@
 import tripy as tp
 from tripy.frontend.trace import Trace
 from tripy.flat_ir.ops import MaxOp
+from tripy.flat_ir.ops.base import FlatIRFunction
 
 
 class TestMaxOp:
@@ -31,20 +32,22 @@ def test_str(self):
         trace = Trace([out])
         flat_ir = trace.to_flat_ir()
 
-        max_op = flat_ir.ops[-1]
-        broadcast_a = flat_ir.ops[-3]
-        broadcast_b = flat_ir.ops[-2]
+        func_max = flat_ir.ops[-1]
+        assert isinstance(func_max, FlatIRFunction)
 
+        max_op = func_max.ops[-1]
+        broadcast_a = func_max.ops[-3]
+        broadcast_b = func_max.ops[-2]
         assert isinstance(max_op, MaxOp)
         assert re.match(
-            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(a, t_inter[0-9]+, broadcast_dim=\[0\]\)",
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(t_inter[0-9]+, t_inter[0-9]+, broadcast_dim=\[0\]\)",
             str(broadcast_a),
         )
         assert re.match(
-            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(b, t_inter[0-9]+, broadcast_dim=\[0\]\)",
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(t_inter[0-9]+, t_inter[0-9]+, broadcast_dim=\[0\]\)",
             str(broadcast_b),
         )
         assert re.match(
-            r"out: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = MaxOp\(t_inter[0-9]+, t_inter[0-9]+\)",
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = MaxOp\(t_inter[0-9]+, t_inter[0-9]+\)",
             str(max_op),
         )

tripy/tests/flat_ir/ops/test_minimum.py

@@ -31,21 +31,22 @@ def test_str(self):
         trace = Trace([out])
         flat_ir = trace.to_flat_ir()
 
-        min_op = flat_ir.ops[-1]
-        broadcast_a = flat_ir.ops[-3]
-        broadcast_b = flat_ir.ops[-2]
+        func_min = flat_ir.ops[-1]
+        min_op = func_min.ops[-1]
+        broadcast_a = func_min.ops[-3]
+        broadcast_b = func_min.ops[-2]
 
         assert isinstance(min_op, MinOp)
         assert re.match(
-            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(a, t_inter[0-9]+, broadcast_dim=\[0\]\)",
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(t_inter[0-9]+, t_inter[0-9]+, broadcast_dim=\[0\]\)",
             str(broadcast_a),
         )
         assert re.match(
-            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(b, t_inter[0-9]+, broadcast_dim=\[0\]\)",
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(t_inter[0-9]+, t_inter[0-9]+, broadcast_dim=\[0\]\)",
             str(broadcast_b),
         )
 
         assert re.match(
-            r"out: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = MinOp\(t_inter[0-9]+, t_inter[0-9]+\)",
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = MinOp\(t_inter[0-9]+, t_inter[0-9]+\)",
             str(min_op),
         )

tripy/tests/flat_ir/ops/test_reduce.py

@@ -18,6 +18,7 @@
 import tripy as tp
 from tripy.frontend.trace import Trace
 from tripy.flat_ir.ops import ArgMinMaxOp, ConvertOp, DivideOp, DynamicBroadcastOp, MulOp, ReduceOp
+from tripy.flat_ir.ops.base import FlatIRFunction
 import re
 
 
@@ -30,10 +31,11 @@ def test_sum_str(self):
         trace = Trace([out])
         flat_ir = trace.to_flat_ir()
 
-        reduce = flat_ir.ops[-1]
+        func_reduce = flat_ir.ops[-1]
+        reduce = func_reduce.ops[-1]
         assert isinstance(reduce, ReduceOp)
         assert re.match(
-            r"out: \[rank=\(1\), dtype=\(int32\), loc=\(gpu:0\)\] = ReduceOp\(inp, t_inter[0-9]+, reduce_mode='sum', reduce_dims=\[0\]\)",
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(int32\), loc=\(gpu:0\)\] = ReduceOp\(t_inter[0-9]+, t_inter[0-9]+, reduce_mode='sum', reduce_dims=\[0\]\)",
             str(reduce),
         )
 
@@ -45,11 +47,11 @@ def test_max_str(self):
         trace = Trace([out])
         flat_ir = trace.to_flat_ir()
 
-        reduce = flat_ir.ops[-1]
+        func_reduce = flat_ir.ops[-1]
+        reduce = func_reduce.ops[-1]
         assert isinstance(reduce, ReduceOp)
-
         assert re.match(
-            r"out: \[rank=\(1\), dtype=\(int32\), loc=\(gpu:0\)\] = ReduceOp\(inp, t_inter[0-9]+, reduce_mode='max', reduce_dims=\[0\]\)",
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(int32\), loc=\(gpu:0\)\] = ReduceOp\(t_inter[0-9]+, t_inter[0-9]+, reduce_mode='max', reduce_dims=\[0\]\)",
             str(reduce),
         )
 
@@ -61,30 +63,43 @@ def test_mean_str(self):
         trace = Trace([out])
         flat_ir = trace.to_flat_ir()
 
-        div = flat_ir.ops[-1]
+        func_div = flat_ir.ops[-1]
+        div = func_div.ops[-1]
+        broadcast_a = func_div.ops[-3]
+        broadcast_b = func_div.ops[-2]
         assert isinstance(div, DivideOp)
         assert re.match(
-            r"out: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DivideOp\(t_inter[0-9]+, t_inter[0-9]+\)",
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DivideOp\(t_inter[0-9]+, t_inter[0-9]+\)",
             str(div),
         )
 
-        broadcast = flat_ir.ops[-2]
-        assert isinstance(broadcast, DynamicBroadcastOp)
+        assert isinstance(broadcast_a, DynamicBroadcastOp)
+        assert re.match(
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(t_inter[0-9]+, t_inter[0-9]+, broadcast_dim=\[[0-9]*\]\)",
+            str(broadcast_a),
+        )
+
+        assert isinstance(broadcast_b, DynamicBroadcastOp)
         assert re.match(
             r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(t_inter[0-9]+, t_inter[0-9]+, broadcast_dim=\[[0-9]*\]\)",
-            str(broadcast),
+            str(broadcast_b),
         )
 
-        mul = flat_ir.ops[-15]
+        mul = flat_ir.ops[-3].ops[-1]
         assert isinstance(mul, MulOp)
         assert re.match(
-            r"t[0-9]+: \[rank=\(0\), dtype=\(int32\), loc=\(gpu:0\)\] = MulOp\(t_inter[0-9]+, t_inter[0-9]+\)",
+            r"t_inter[0-9]+: \[rank=\(0\), dtype=\(int32\), loc=\(gpu:0\)\] = MulOp\(t_inter[0-9]+, t_inter[0-9]+\)",
             str(mul),
         )
-        reduce = flat_ir.ops[2]
+
+        func_reduce = flat_ir.ops[1]
+        assert isinstance(func_reduce, FlatIRFunction)
+
+        reduce = func_reduce.ops[-1]
         assert isinstance(reduce, ReduceOp)
+
         assert re.match(
-            r"t[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = ReduceOp\(inp, t_inter[0-9]+, reduce_mode='sum', reduce_dims=\[0\]\)",
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = ReduceOp\(t_inter[0-9]+, t_inter[0-9]+, reduce_mode='sum', reduce_dims=\[0\]\)",
             str(reduce),
         )
 
@@ -96,11 +111,14 @@ def test_argmax_str(self):
         trace = Trace([out])
         flat_ir = trace.to_flat_ir()
 
-        reduce = flat_ir.ops[-1]
-        assert isinstance(reduce, ArgMinMaxOp)
+        func_argminmax = flat_ir.ops[-1]
+        assert isinstance(func_argminmax, FlatIRFunction)
+
+        argminmax = func_argminmax.ops[-1]
+        assert isinstance(argminmax, ArgMinMaxOp)
         assert re.match(
-            r"out: \[rank=\(1\), dtype=\(int32\), loc=\(gpu:0\)\] = ArgMinMaxOp\(inp, t[0-9]+, t_inter[0-9]+, t_inter[0-9]+, reduce_mode='argmax', reduce_dims=\[0\]\)",
-            str(reduce),
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(int32\), loc=\(gpu:0\)\] = ArgMinMaxOp\(t_inter[0-9]+, t_inter[0-9]+, t_inter[0-9]+, t_inter[0-9]+, reduce_mode='argmax', reduce_dims=\[0\]\)",
+            str(argminmax),
         )
 
     def test_argmin_str(self):
@@ -111,9 +129,12 @@ def test_argmin_str(self):
         trace = Trace([out])
         flat_ir = trace.to_flat_ir()
 
-        reduce = flat_ir.ops[-1]
-        assert isinstance(reduce, ArgMinMaxOp)
+        func_argminmax = flat_ir.ops[-1]
+        assert isinstance(func_argminmax, FlatIRFunction)
+
+        argminmax = func_argminmax.ops[-1]
+        assert isinstance(argminmax, ArgMinMaxOp)
         assert re.match(
-            r"out: \[rank=\(1\), dtype=\(int32\), loc=\(gpu:0\)\] = ArgMinMaxOp\(inp, t[0-9]+, t_inter[0-9]+, t_inter[0-9]+, reduce_mode='argmin', reduce_dims=\[0\]\)",
-            str(reduce),
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(int32\), loc=\(gpu:0\)\] = ArgMinMaxOp\(t_inter[0-9]+, t_inter[0-9]+, t_inter[0-9]+, t_inter[0-9]+, reduce_mode='argmin', reduce_dims=\[0\]\)",
+            str(argminmax),
         )

tripy/tests/flat_ir/ops/test_subtract.py

@@ -19,6 +19,7 @@
 import tripy as tp
 from tripy.frontend.trace import Trace
 from tripy.flat_ir.ops import SubtractOp
+from tripy.flat_ir.ops.base import FlatIRFunction
 
 
 class TestSubtractOp:
@@ -31,21 +32,24 @@ def test_str(self):
         trace = Trace([out])
         flat_ir = trace.to_flat_ir()
 
-        sub = flat_ir.ops[-1]
-        broadcast_a = flat_ir.ops[-3]
-        broadcast_b = flat_ir.ops[-2]
+        func_sub = flat_ir.ops[-1]
+        assert isinstance(func_sub, FlatIRFunction)
+
+        sub = func_sub.ops[-1]
+        broadcast_a = func_sub.ops[-3]
+        broadcast_b = func_sub.ops[-2]
 
         assert isinstance(sub, SubtractOp)
 
         assert re.match(
-            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(a, t_inter[0-9]+, broadcast_dim=\[0\]\)",
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(t_inter[0-9]+, t_inter[0-9]+, broadcast_dim=\[0\]\)",
             str(broadcast_a),
         )
         assert re.match(
-            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(b, t_inter[0-9]+, broadcast_dim=\[0\]\)",
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicBroadcastOp\(t_inter[0-9]+, t_inter[0-9]+, broadcast_dim=\[0\]\)",
             str(broadcast_b),
         )
         assert re.match(
-            r"out: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = SubtractOp\(t_inter[0-9]+, t_inter[0-9]+\)",
+            r"t_inter[0-9]+: \[rank=\(1\), dtype=\(float32\), loc=\(gpu:0\)\] = SubtractOp\(t_inter[0-9]+, t_inter[0-9]+\)",
             str(sub),
         )

tripy/tests/flat_ir/test_constant_deduplication.py

@@ -15,6 +15,7 @@
 
 import pytest
 from tripy.flat_ir.flat_ir import FlatIR
+from tripy.flat_ir.ops.base import FlatIRFunction
 from tripy.flat_ir.ops import ConstantOp
 from tripy.flat_ir.tensor import FlatIRTensor
 from tripy.common.device import device
@@ -29,7 +30,7 @@ def __init__(self, inputs, outputs):
             output.producer = self
 
 
-def create_subgraph():
+def create_subgraph(config):
     # Create constant tensors and ops
     const1 = FlatIRTensor.build(shape=[2], rank=1, dtype=int32, reason_details="", device=device("gpu"))
     op1 = ConstantOp.build([], [const1], data=[1, 2])
@@ -48,18 +49,34 @@ def create_subgraph():
     result_tensor = FlatIRTensor.build(shape=[2], rank=1, dtype=int32, reason_details="", device=device("gpu"))
     mock_op = MockOp([const1, const2, const3], [result_tensor])
 
+    if config == "func":
+        # 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_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]
+
     return [], [result_tensor]
 
 
-def test_integrate_subgraph_constant_deduplication():
+@pytest.mark.parametrize("config", ["main", "func"])
+def test_integrate_subgraph_constant_deduplication(config):
     flat_ir = FlatIR()
-    inputs, outputs = create_subgraph()
+    inputs, outputs = create_subgraph(config)
 
     # Integrate the subgraph
     flat_ir.integrate_subgraph(inputs, outputs)
 
     # Verify that only two ConstantOps remain
-    constant_ops = [op for op in flat_ir.ops if isinstance(op, ConstantOp)]
+    ops = flat_ir.ops[0].ops if isinstance(flat_ir.ops[0], FlatIRFunction) else flat_ir.ops
+
+    constant_ops = [op for op in ops if isinstance(op, ConstantOp)]
     assert len(constant_ops) == 2, "There should be only two ConstantOps after integration"
 
     # Verify that the remaining ConstantOps have different data
@@ -72,12 +89,6 @@ def test_integrate_subgraph_constant_deduplication():
     assert set(constant_data) == expected_data, f"Expected constant data {expected_data}, but got {set(constant_data)}"
 
     # Verify that the mock op now uses the same tensor for its first two inputs
-    mock_op = [op for op in flat_ir.ops if isinstance(op, MockOp)][0]
+    mock_op = [op for op in ops if isinstance(op, MockOp)][0]
     assert mock_op.inputs[0] is mock_op.inputs[1], "The mock op should use the same tensor for its first two inputs"
     assert mock_op.inputs[0] is not mock_op.inputs[2], "The mock op should still have a different third input"
-
-    # Verify that tensor replacements were applied
-    assert len(flat_ir.tensor_replacements) > 0, "There should be tensor replacements after integration"
-
-    # Verify that the constant map has the correct number of entries
-    assert len(flat_ir.constant_map) == 2, "Constant map should have 2 entries"