Add ability to specify CoreML pipeline passes

[jgibson2]

Summary

Adds the ability to specify a set of CoreML passes as a CompileSpec, allowing additional customization of the model compilation.

Test plan

Converted a model and made sure it worked with a custom pipeline. Also ensured via print statements that the passes were translated correctly.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/16118

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 1 Unrelated Failure

As of commit d9174e4 with merge base 9193566 ():

NEW FAILURE - The following job has failed:

• trunk / test-arm-cortex-m-size-test (zephyr-preset) / linux-job (gh)
  RuntimeError: Command docker exec -t eafbc030667e0a1343e0e60acb79eb3892ce0d83f8b152d01db14bbd89fc993a /exec failed with exit code 1

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / android / run-emulator (gh) (trunk failure)
  Timeout waiting for emulator to boot.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Copilot]

Pull request overview

This PR adds the ability to specify a custom CoreML pass pipeline through the CompileSpec system, enabling users to customize the model compilation process beyond the default pipeline.

Key Changes:

• Added PASS_PIPELINE key to the compile spec enumeration
• Implemented methods to serialize/deserialize pass pipeline configurations
• Integrated custom pass pipeline into the CoreML conversion process

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[jgibson2]

@pytorchbot label "release notes: none"

[Copilot]

Pull request overview

Copilot reviewed 1 out of 1 changed files in this pull request and generated 3 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[metascroy]

Thanks! Can we also have a new CI test added for this?

@cymbalrush @YifanShenSZ can you also have a look?

[jgibson2]

Thanks! Can we also have a new CI test added for this?

@cymbalrush @YifanShenSZ can you also have a look?

Sure, what kind of test are you envisioning? Can you point me towards an example?

[metascroy]

Thanks! Can we also have a new CI test added for this?
@cymbalrush @YifanShenSZ can you also have a look?

Sure, what kind of test are you envisioning? Can you point me towards an example?

Basically I'm looking to automate the test you describe in this PR: "Converted a model and made sure it worked with a custom pipeline. Also ensured via print statements that the passes were translated correctly."

You can see examples of CoreML tests for compile-spec based features:

• executorch/backends/apple/coreml/test/test_coreml_partitioner.py

  Lines 187 to 230 in aff5086

  	def test_buffer(self):
  	embedding_dim = 3
  	max_seq_len = 2
  	class Model(torch.nn.Module):
  	def __init__(self):
  	super().__init__()
  	self.register_buffer(
  	"cache",
  	torch.zeros((max_seq_len, embedding_dim), dtype=torch.float32),
  	)
  	def forward(self, q, k_val, input_pos):
  	q_T = q.transpose(0, 1)
  	k = torch.ops.aten.index_put_(self.cache, [input_pos, None], k_val)
  	attn = k.mm(q_T)
  	return attn
  	model = Model()
  	model.eval()
  	q = torch.randn((1, embedding_dim))
  	k_val = torch.randn((1, embedding_dim))
  	input_pos = torch.tensor([0])
  	example_inputs = (q, k_val, input_pos)
  	exir_program_aten = torch.export.export(model, example_inputs, strict=True)
  	compile_specs = CoreMLBackend.generate_compile_specs(
  	minimum_deployment_target=ct.target.iOS18
  	)
  	partitioner = CoreMLPartitioner(compile_specs=compile_specs)
  	edge_program_manager = executorch.exir.to_edge(
  	exir_program_aten, compile_config=self.edge_compile_config
  	)
  	delegated_program_manager = edge_program_manager.to_backend(partitioner)
  	assert [
  	node.target.__name__
  	for node in delegated_program_manager.exported_program().graph.nodes
  	if node.op == "call_function"
  	] == [
  	"executorch_call_delegate",
  	"getitem",
  	]

• https://github.com/pytorch/executorch/blob/main/backends/apple/coreml/test/test_enumerated_shapes.py

There's usually an AOT portion to the test (making sure graph looks as expected), and a runtime component of the test.

[jgibson2]

Thanks! Can we also have a new CI test added for this?
@cymbalrush @YifanShenSZ can you also have a look?

Sure, what kind of test are you envisioning? Can you point me towards an example?

Basically I'm looking to automate the test you describe in this PR: "Converted a model and made sure it worked with a custom pipeline. Also ensured via print statements that the passes were translated correctly."

You can see examples of CoreML tests for compile-spec based features:

• executorch/backends/apple/coreml/test/test_coreml_partitioner.py

  Lines 187 to 230 in aff5086

  	def test_buffer(self):
  	embedding_dim = 3
  	max_seq_len = 2
  	class Model(torch.nn.Module):
  	def __init__(self):
  	super().__init__()
  	self.register_buffer(
  	"cache",
  	torch.zeros((max_seq_len, embedding_dim), dtype=torch.float32),
  	)
  	def forward(self, q, k_val, input_pos):
  	q_T = q.transpose(0, 1)
  	k = torch.ops.aten.index_put_(self.cache, [input_pos, None], k_val)
  	attn = k.mm(q_T)
  	return attn
  	model = Model()
  	model.eval()
  	q = torch.randn((1, embedding_dim))
  	k_val = torch.randn((1, embedding_dim))
  	input_pos = torch.tensor([0])
  	example_inputs = (q, k_val, input_pos)
  	exir_program_aten = torch.export.export(model, example_inputs, strict=True)
  	compile_specs = CoreMLBackend.generate_compile_specs(
  	minimum_deployment_target=ct.target.iOS18
  	)
  	partitioner = CoreMLPartitioner(compile_specs=compile_specs)
  	edge_program_manager = executorch.exir.to_edge(
  	exir_program_aten, compile_config=self.edge_compile_config
  	)
  	delegated_program_manager = edge_program_manager.to_backend(partitioner)
  	assert [
  	node.target.__name__
  	for node in delegated_program_manager.exported_program().graph.nodes
  	if node.op == "call_function"
  	] == [
  	"executorch_call_delegate",
  	"getitem",
  	]

• https://github.com/pytorch/executorch/blob/main/backends/apple/coreml/test/test_enumerated_shapes.py

There's usually an AOT portion to the test (making sure graph looks as expected), and a runtime component of the test.

I do think those tests are a little different, as they are mostly testing the partitioner. I can add a test for translating a list of strings into a CompileSpec and back again if you want, but as far as what happens under the hood of coremltools I don't know how to adequately test that -- an MLModel comes out, not something from executorch.

[metascroy]

artitioner. I can add a test for translating a list of strings into a CompileSpec and back again if you want, but as far as what happens under the hood of coremltools I don't know how to adequately test that -- an MLModel comes out, not something from executorch.

I understand that the pass applies on the MLModel, but how do you know the pass you provide to the compile spec is actually being applied? Do you see a way to test this?

Once simple test might be to grep the output of the to_edge_lower_and_transform and look for something like "Applying user passes" if Core ML logs something like that.

Another possibility is to extract the mlmodel from the pte and inspect it has been transformed as you expect from the pass.

Perhaps you could tell me more about the custom pass you're applying and what it does to help brainstorm test ideas?

[jgibson2]

artitioner. I can add a test for translating a list of strings into a CompileSpec and back again if you want, but as far as what happens under the hood of coremltools I don't know how to adequately test that -- an MLModel comes out, not something from executorch.

I understand that the pass applies on the MLModel, but how do you know the pass you provide to the compile spec is actually being applied? Do you see a way to test this?

Once simple test might be to grep the output of the to_edge_lower_and_transform and look for something like "Applying user passes" if Core ML logs something like that.

Another possibility is to extract the mlmodel from the pte and inspect it has been transformed as you expect from the pass.

Perhaps you could tell me more about the custom pass you're applying and what it does to help brainstorm test ideas?

One way to tell is that coremltools will output the following log when specifying a user pipeline:

Running MIL executorch_user_pipeline pipeline: 100%|██████████| 9/9 [00:00<00:00, 16.49 passes/s]

but inspecting the mlmodel to make sure that e.g. op fusion was not performed is a whole other can of worms.

The pipeline I tested with is:

        pass_names=[
            "common::const_elimination",
            "common::const_deduplication",
            "common::remove_symbolic_reshape",
            "common::noop_elimination",
            "common::merge_consecutive_relus",
            "common::merge_consecutive_reshapes",
            "common::merge_consecutive_transposes",
            "common::dedup_op_and_var_names",
            "common::dead_code_elimination",  # always end with dce
        ]

The reason that I am using this pipeline is that models exported with the default pipeline produce wildly different results on iPhone 13s vs iPhone 15s.

[metascroy]

artitioner. I can add a test for translating a list of strings into a CompileSpec and back again if you want, but as far as what happens under the hood of coremltools I don't know how to adequately test that -- an MLModel comes out, not something from executorch.

I understand that the pass applies on the MLModel, but how do you know the pass you provide to the compile spec is actually being applied? Do you see a way to test this?
Once simple test might be to grep the output of the to_edge_lower_and_transform and look for something like "Applying user passes" if Core ML logs something like that.
Another possibility is to extract the mlmodel from the pte and inspect it has been transformed as you expect from the pass.
Perhaps you could tell me more about the custom pass you're applying and what it does to help brainstorm test ideas?

One way to tell is that coremltools will output the following log when specifying a user pipeline:

Running MIL executorch_user_pipeline pipeline: 100%|██████████| 9/9 [00:00<00:00, 16.49 passes/s]

but inspecting the mlmodel to make sure that e.g. op fusion was not performed is a whole other can of worms.

The pipeline I tested with is:

        pass_names=[
            "common::const_elimination",
            "common::const_deduplication",
            "common::remove_symbolic_reshape",
            "common::noop_elimination",
            "common::merge_consecutive_relus",
            "common::merge_consecutive_reshapes",
            "common::merge_consecutive_transposes",
            "common::dedup_op_and_var_names",
            "common::dead_code_elimination",  # always end with dce
        ]

The reason that I am using this pipeline is that models exported with the default pipeline produce wildly different results on iPhone 13s vs iPhone 15s.

Sounds good. I guess this PR is OK. If possible, let's have @YifanShenSZ from Apple weigh in.