fix coreml partitioner for llama

Differential Revision: D55663281

Author: cccclai

backends/apple/coreml/partition/coreml_partitioner.py

@@ -18,6 +18,7 @@
     PartitionResult,
 )
 from executorch.exir.backend.utils import tag_constant_data
+from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
 from torch.export.exported_program import ExportedProgram
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
 from torch.fx.passes.operator_support import OperatorSupportBase
@@ -88,7 +89,25 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 node.meta["delegation_tag"] = tag
                 partition_tags[tag] = self.delegation_spec
 
-        tag_constant_data(exported_program)
+                is_attr = (
+                    node.op == "placeholder"
+                    and (
+                        is_param(exported_program, node)
+                        or is_buffer(exported_program, node)
+                        or is_lifted_tensor_constant(exported_program, node)
+                    )
+                )
+                # if all users of const/param/buffer nodes are partitioned then partition
+                if is_attr:
+                    user_tags = set()
+                    for user in node.users:
+                        user_tag = user.meta.get("delegation_tag", None)
+                        if user_tag is not None:
+                            user_tags.add(user_tag)
+                    if len(user_tags) >= 1:
+                        # There are more than one user tag, just pick one and there is a pass later to copy the constant data
+                        node.meta["delegation_tag"] = user_tags.pop()
+
 
         return PartitionResult(
             tagged_exported_program=exported_program, partition_tags=partition_tags