Fix accuracy issue for double alias

python/aitemplate/backend/codegen.py

@@ -370,7 +370,6 @@ def __init__(
         self.visited_func = set()
         self.visited_dims = set()
         self.set_up_constant_names = []
-        self.param_name_to_ptr_idx = {}
 
         self.num_constants = 0
         self.constants_data_size = 0
@@ -584,7 +583,6 @@ def _codegen_input_tensor(self, tensor: Tensor) -> None:
             )
         )
         self.set_inputs.append(check_not_null(tensor))
-        self.param_name_to_ptr_idx[name] = self.input_idx
         self._record_param_tensor_info(tensor, self.input_idx)
         self.input_idx += 1
 
@@ -599,29 +597,27 @@ def _get_output_idx(self, name: str) -> int:
     def _codegen_output_aliases_tensor(self, tensor: Tensor) -> None:
         name = tensor._attrs["name"]
         view = tensor._attrs["is_view_of"]
-        if tensor._attrs["external_tensor"] is not None:
+        external_tensor = tensor._attrs["external_tensor"]
+        if external_tensor is not None:
+            assert not external_tensor._attrs[
+                "is_param"
+            ], "Views of constants should be folded."
             self.set_inputs.append(set_value(name, view._attrs["name"]))
             return
-        is_view = view is not None
-        if is_view and (view._attrs["name"] in self.param_name_to_ptr_idx):
-            ptr_idx = self.param_name_to_ptr_idx[view._attrs["name"]]
+
+        if view:
+            # View is already initialized, assign to view.
             self.set_inputs.append(set_value(name, view._attrs["name"]))
         else:
-            ptr_idx = self._get_output_idx(name)
+            # Original tensor, initialize it.
+            output_idx = self._get_output_idx(name)
             self.set_inputs.append(
                 set_value(
                     name,
-                    f"static_cast<decltype({name})>(params_[{ptr_idx}].ptr)",
+                    f"static_cast<decltype({name})>(params_[{output_idx}].ptr)",
                 )
             )
 
-        self.param_name_to_ptr_idx[name] = ptr_idx
-        if tensor._attrs["is_output"]:
-            self._record_param_tensor_info(tensor, ptr_idx)
-            self.set_inputs.append(
-                check_not_null(tensor, skip_if_lower_bound_is_zero=True)
-            )
-
     def _codegen_output_tensor(self, tensor: Tensor) -> None:
         is_param = tensor._attrs["is_param"]
         is_input = tensor._attrs["is_input"]
@@ -634,22 +630,26 @@ def _codegen_output_tensor(self, tensor: Tensor) -> None:
 
         if is_param:
             self._codegen_param_setup(tensor)
-            self._record_param_tensor_info(tensor, output_idx)
             self.device_to_device_copies.append(device_copy(tensor, tensor, output_idx))
-        elif external_tensor is not None:
-            # Special view cases for outputs; we can hit this case if the output
-            # is a view of a constant, input, or another output.
+        elif is_view or external_tensor is not None:
             assert (
                 is_view
-            ), f"orig_tensor is not None, but node {name} is not marked as a view! Node: {tensor}"
+            ), f"External tensor is not None, but node {name} is not marked as a view! Node: {tensor}"
+            view_name = view._attrs["name"]
+            self.set_inputs.append(set_value(name, view_name))
             self.set_inputs.append(
-                check_not_null(tensor, output_idx, skip_if_lower_bound_is_zero=True)
+                check_not_null(tensor, skip_if_lower_bound_is_zero=True)
             )
-            self.set_inputs.append(set_value(name, view._attrs["name"]))
-            self.device_to_device_copies.append(
-                device_copy(tensor, external_tensor, output_idx)
+
+            view_assigned_to_another_output = (
+                self._get_output_idx(view_name) != output_idx
             )
-            self._record_param_tensor_info(tensor, output_idx)
+            if external_tensor or view_assigned_to_another_output:
+                # Copy from original tensor so this output can also have the data.
+                original_tensor = external_tensor if external_tensor else view
+                self.device_to_device_copies.append(
+                    device_copy(tensor, original_tensor, output_idx)
+                )
         elif is_input:
             # Inputs that are also outputs require an extra copy
             self.set_inputs.append(
@@ -659,11 +659,20 @@ def _codegen_output_tensor(self, tensor: Tensor) -> None:
                 )
             )
             self._record_param_tensor_info(tensor, self.input_idx)
-            self._record_param_tensor_info(tensor, output_idx)
             self.device_to_device_copies.append(device_copy(tensor, tensor, output_idx))
             self.input_idx += 1
         else:
-            self._codegen_output_aliases_tensor(tensor)
+            self.set_inputs.append(
+                set_value(
+                    name,
+                    f"static_cast<decltype({name})>(params_[{output_idx}].ptr)",
+                )
+            )
+            self.set_inputs.append(
+                check_not_null(tensor, skip_if_lower_bound_is_zero=True)
+            )
+
+        self._record_param_tensor_info(tensor, output_idx)
 
     def _process_dims(self, shape: List[IntVar]) -> None:
         for dim in shape:

tests/unittest/backend/test_codegen_output_aliases.py

@@ -0,0 +1,267 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+#  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.
+#
+
+import os
+import unittest
+from typing import Sequence
+
+from aitemplate.backend.codegen import set_value
+
+from aitemplate.compiler import compile_model, ops
+from aitemplate.compiler.base import Tensor
+from aitemplate.compiler.ops.common.epilogue import FuncEnum
+from aitemplate.testing import detect_target
+
+from aitemplate.testing.test_utils import gen_input_tensor, get_random_torch_tensor
+
+
+class TestCodegenOutputAliases(unittest.TestCase):
+    MODEL_GENERATED_FILE = "model-generated.h"
+    WORKDIR = "/tmp/ait/test_codegen_output_aliases"
+    SHAPE = (2, 3, 4)
+
+    def _check_string_in_file(self, file_path: str, target_string: str) -> bool:
+        with open(file_path, "r") as f:
+            content = f.read()
+            return target_string in content
+
+    def _assert_codegen_exists(
+        self,
+        test_name: str,
+        expected_codegen: Sequence[str],
+        filename: str,
+        exists: bool = True,
+    ) -> bool:
+        model_path = os.path.join(self.WORKDIR, test_name, filename)
+        for line in expected_codegen:
+            self.assertEqual(
+                exists,
+                self._check_string_in_file(model_path, line),
+                f"Expected this to be codegen'd in {model_path}: \n {line}",
+            )
+
+    """
+    ### Test Cases:
+    NOTE: <-- represents a "is_view_of" edge
+    1. gelu <-- output
+    2. intermediate0 <-- intermediate1 <-- output
+    3. gelu <-- output_0
+          ^---- output_1
+    4. input <-- intermediate <-- output
+    5. constant <-- intermediate <-- output
+    6. output_0 <-- intermediate <-- output_1
+    """
+
+    def test_simple_output_alias(self, test_name="simple_output_alias"):
+        """
+        Case: Simple case where an output aliases an internal tensor.
+        Graph: ( gelu ) <--view-- ( output )
+        """
+        x = gen_input_tensor(shape=self.SHAPE, name="input_x")
+        gelu = ops.elementwise(FuncEnum.GELU)(x)  # has an output alias
+        output0 = ops.flatten()(gelu)
+        output0._attrs["is_output"] = True
+        output0._attrs["name"] = "output_0"
+
+        compile_model(
+            output0,
+            detect_target(),
+            self.WORKDIR,
+            test_name,
+            do_optimize_graph=False,
+        )
+
+        view_name = output0._attrs["is_view_of"]._attrs["name"]
+        expected_codegen = (
+            set_value(
+                view_name,
+                f"static_cast<decltype({view_name})>(params_[1].ptr)",
+            ),
+        )
+        self._assert_codegen_exists(
+            test_name, expected_codegen, self.MODEL_GENERATED_FILE
+        )
+
+    def test_view_of_view(self, test_name="view_of_view"):
+        """
+        Case: There's two intermediate views, where one is a view of the other.
+        Graph: ( input ) <--view-- ( intermediate0 ) <--view-- ( intermediate1 ) <--view-- ( output_0 )
+        """
+        x = gen_input_tensor(shape=self.SHAPE, name="input_x")
+        intermediate0 = ops.unsqueeze(dim=0)(x)  # has an output alias
+        intermediate1 = ops.unsqueeze(dim=0)(intermediate0)  # has an output alias
+        output0 = ops.flatten()(intermediate1)
+        output0._attrs["is_output"] = True
+        output0._attrs["name"] = "output_0"
+
+        compile_model(
+            output0,
+            detect_target(),
+            self.WORKDIR,
+            test_name,
+            do_optimize_graph=False,
+        )
+
+        intermediate0_name = intermediate0._attrs["name"]
+        intermediate1_name = intermediate1._attrs["name"]
+        expected_codegen = (
+            set_value(intermediate0_name, "input_x"),
+            set_value(intermediate1_name, intermediate0_name),
+        )
+        self._assert_codegen_exists(
+            test_name, expected_codegen, self.MODEL_GENERATED_FILE
+        )
+
+    def test_double_alias(self, test_name="double_alias"):
+        """
+        Case: Two outputs are a view of the same tensor.
+        Graph: ( gelu ) <--view-- ( output_0 )
+                        <--view-- ( output_1 )
+        Expected: gelu will be initialized with params_[2]
+        """
+        # AIT, two outputs.
+        x = gen_input_tensor(shape=self.SHAPE, name="input_x")
+        gelu = ops.elementwise(FuncEnum.GELU)(x)  # has an output alias
+        output0 = ops.unsqueeze(dim=0)(gelu)
+        output0._attrs["is_output"] = True
+        output0._attrs["name"] = "output_0"
+
+        output1 = ops.flatten()(gelu)
+        output1._attrs["is_output"] = True
+        output1._attrs["name"] = "output_1"
+
+        compile_model(
+            [output0, output1],
+            detect_target(),
+            self.WORKDIR,
+            test_name,
+            do_optimize_graph=True,
+        )
+
+        view_name = output0._attrs["is_view_of"]._attrs["name"]
+        expected_codegen = (
+            # input_x assigned to params_[0].ptr
+            set_value(
+                view_name,
+                f"static_cast<decltype({view_name})>(params_[2].ptr)",
+                # see _construct_output_name_to_index_map for params indexing
+            ),
+        )
+        self._assert_codegen_exists(
+            test_name, expected_codegen, self.MODEL_GENERATED_FILE
+        )
+
+    # ======================================================================== #
+    #    Cases where the has_output_alias tensor also has an external tensor   #
+    # ======================================================================== #
+    def test_external_tensor_is_input(self, test_name="external_tensor_is_input"):
+        """
+        Case: The intermediate tensor is a view of an input.
+        Graph: ( input ) <--view-- ( intermediate ) <--view-- ( output_0 )
+        """
+
+        x = gen_input_tensor(shape=self.SHAPE, name="input_x")
+        intermediate_view = ops.unsqueeze(dim=0)(x)  # has an output alias
+        output0 = ops.flatten()(intermediate_view)
+        output0._attrs["is_output"] = True
+        output0._attrs["name"] = "output_0"
+
+        compile_model(
+            output0,
+            detect_target(),
+            self.WORKDIR,
+            test_name,
+            do_optimize_graph=False,
+        )
+        intermediate_view_name = intermediate_view._attrs["name"]
+        expected_codegen = (
+            # input_x is assigned to params_[0].ptr
+            set_value(intermediate_view_name, "input_x"),
+        )
+        self._assert_codegen_exists(
+            test_name, expected_codegen, self.MODEL_GENERATED_FILE
+        )
+
+    def test_external_tensor_is_constant(self, test_name="external_tensor_is_constant"):
+        """
+        Case: The intermediate tensor is a view of a constant.
+        Graph: ( constant ) <--view-- ( intermediate ) <--view-- ( output_0 )
+        Expect: The intermediate tensor is constant folded.
+        """
+        constant = Tensor(
+            shape=self.SHAPE, dtype="float16", name="constant"
+        )  # has an output alias
+        intermediate_view = ops.unsqueeze(dim=0)(constant)  # has an output alias
+        output0 = ops.flatten()(intermediate_view)
+        output0._attrs["is_output"] = True
+        output0._attrs["name"] = "output_0"
+
+        constant_pt = get_random_torch_tensor(self.SHAPE, dtype="float16")
+
+        compile_model(
+            output0,
+            detect_target(),
+            self.WORKDIR,
+            test_name,
+            constants={"constant": constant_pt},
+            do_optimize_graph=False,
+        )
+        intermediate_view_name = intermediate_view._attrs["name"]
+        expected_codegen = (
+            set_value(intermediate_view_name, "constant"),
+            set_value("output_0", intermediate_view_name),
+        )
+        # We expect the intermediate tensor and output_0 to be constant folded.
+        # So, we don't expect _codegen_output_aliases_tensor to be called.
+        self._assert_codegen_exists(
+            test_name, expected_codegen, self.MODEL_GENERATED_FILE, exists=False
+        )
+
+    def test_external_tensor_is_output(self, test_name="external_tensor_is_output"):
+        """
+        Case: The intermediate tensor is a view of an output.
+        Graph: ( output_0 ) <--view-- ( intermediate ) <--view-- ( output_1 )
+        """
+        x = gen_input_tensor(shape=self.SHAPE, name="input_x")
+        z = gen_input_tensor(shape=self.SHAPE, name="input_z")
+        output0 = ops.elementwise(FuncEnum.RELU)(x + z)
+        output0._attrs["is_output"] = True
+        output0._attrs["name"] = "output_0"
+
+        intermediate_view = ops.unsqueeze(dim=0)(output0)  # has an output alias
+        output1 = ops.flatten()(intermediate_view)
+        output1._attrs["is_output"] = True
+        output1._attrs["name"] = "output_1"
+
+        compile_model(
+            [output0, output1],
+            detect_target(),
+            self.WORKDIR,
+            test_name,
+            do_optimize_graph=False,
+        )
+        intermediate_name = intermediate_view._attrs["name"]
+        expected_codegen = (
+            # input_z is assigned to params_[0].ptr
+            # input_x is assigned to params_[1].ptr
+            set_value(
+                "output_0",
+                "static_cast<decltype(output_0)>(params_[2].ptr)",
+            ),
+            set_value(intermediate_name, "output_0"),
+        )
+        self._assert_codegen_exists(
+            test_name, expected_codegen, self.MODEL_GENERATED_FILE
+        )