Clarify placeholder IR that delegates need to support in ExecuTorch

[metascroy]

📚 The doc issue

We do not specify the IR for placeholders/graph inputs that delegates need to support.

For example, today I think placeholders will usually be tensors, with floats and ints inputs being lifted to rank0 tensors. But raw symints also show up as placeholders in the presence of graph breaks. For example,

Lowered module 0:

def forward(...):
    sym_size: "Sym(s53)" = torch.ops.aten.sym_size.int(attention_mask, 1);  attention_mask = None
    return [sym_size, aten_expand_copy_default, aten_add_tensor_1]

Lowered module 1:

def forward(aten_embedding_default_2: "f32[1, s53, 768]", sym_size: "Sym(s53)"):
    aten_view_copy_default_3: "f32[1, s53, 768]" = executorch_exir_dialects_edge__ops_aten_view_copy_default(aten_permute_copy_default_3, [1, sym_size, 768]); 

We should clarify placeholder IR that delegates are expected to support.

Suggest a potential alternative/fix

No response

cc @JacobSzwejbka @angelayi

[JacobSzwejbka]

Delegates dont have to support sym ints if they refuse to eat nodes that consume sym ints.

I guess Im not sure what documentation would help here besides the implicit, delegates need to know how to handle anything that could be an input or output of the nodes they consume.

Which is tensors and symInt, symBool, symFloat.

Tensors can be of several kind which you can see by looking at the graph signature.

https://docs.pytorch.org/docs/main/export.ir_spec.html is the spec for export ir

Syms notably are not constants. Which is why they are not promoted to tensor. They originate from things like sym_size (in python tensor.size(dim)), and you can run all the normal python prim operators on them like >, <, +, -, etc

[JacobSzwejbka]

Also on discord you mentioned

In terms of how they show up, it happens because of partitioning. If sym_int is first defined in lowered_module0, then used in lowered_module1, then ET passes the sym_int to lowered_module1 as a placeholder. So the delegate needs to handle EPs with sym_ints as placeholders to avoid erroring out. If the entire graph delegates, they don't show up

I just want to point out that they absolutely can appear if the entire graph delegates. Export supports dynamic prim user inputs now. Its not well tested in ET and I'm sure itll violate an assert in the runtime today (we can relax this now), but backends should not be assuming they cant appear as inputs.

[metascroy]

Delegates dont have to support sym ints if they refuse to eat nodes that consume sym ints.

CoreML does support symints, but not when they are placeholders in the graph. They must be created and consumed in the same lowered graph, and I don't know a great way to ensure ET's partitioner does not split the creation of a symint from its use.

How are symint placeholders handled at runtime in a way that is different than ints? It seems like symints would just be regular ints that get passed in as inputs? But generally, integer inputs are wrapped as rank0 tensors by export/ET.

Should the ET partitioner actually do this symint -> int wrapping when creating graph breaks? For example, instead of creating something like this:

def forward(aten_embedding_default_2: "f32[1, s53, 768]", sym_size: "Sym(s53)"):
    aten_view_copy_default_3: "f32[1, s53, 768]" = executorch_exir_dialects_edge__ops_aten_view_copy_default(aten_permute_copy_default_3, [1, sym_size, 768]); 

should ET have created something like this:

def forward(aten_embedding_default_2: "f32[1, s53, 768]", sym_size_rank0: "i32[]"):
    sym_size = torch.ops.aten.sym_size.int(sym_size_rank0, 0);
    aten_view_copy_default_3: "f32[1, s53, 768]" = executorch_exir_dialects_edge__ops_aten_view_copy_default(aten_permute_copy_default_3, [1, sym_size, 768]); 

[angelayi]

I don't know a great way to ensure ET's partitioner does not split the creation of a symint from its use.

I wonder if the partitioner can just check that none of the annotated nodes contain arguments that are symints and are not annotated. Then the partitioner can decide to annotate the node producing the symint, or don't annotate the node consuming the symint.

Should the ET partitioner actually do this symint -> int wrapping when creating graph breaks? For example, instead of creating something like this:

I'm not sure where would be the proper place to put this pass. You would need to modify the toplevel program to first wrap the symint as a scalar tensor, and then you could call .item on this value once you're within the delegated graph. Additionally, I recall one of the contracts of the partitioner is that it should not modify the original graph.

[metascroy]

I wonder if the partitioner can just check that none of the annotated nodes contain arguments that are symints and are not annotated. Then the partitioner can decide to annotate the node producing the symint, or don't annotate the node consuming the symint.

If I understand correct, if an op contains a symint arg, we'd not annotate it in this case to avoid having a symint show up in a placeholder. This is overly exclusive, though, because CoreML can handle ops with symint args. The restriction is that the symints cannot be placeholders, so they just need to be created and used in the same partition.

I also don't really understand how ET handles symint placeholders at runtime, and how they are different than regular ints at runtime (except regular ints get wrapped as rank0 tensors). cc @JacobSzwejbka

With that said, I think I will add an option to the CoreML partitioner like (lower_full_graph: bool). When lower_full_graph is True, it will error out if there is an unsupported node. lower_full_graph mode will take ops with symint args. If lower_full_graph is False, it will not take ops with symint nodes.