Simple indexing prevents ios18.gather from being schedule-able on the ANE

🐛 Describe the bug

I'm lowering my PyTorch model to CoreML + Executorch, an everything delegates (or is delegate-able) to ANE, except some ios18.gather op that is not delegate-able to ANE.

I boiled the issue down to this small repro:

import torch
import torch.nn as nn

from typing import Any

import coremltools as ct
import torch

from coremltools.optimize.torch.quantization.quantization_config import (
    LinearQuantizerConfig,
    QuantizationScheme,
)
from executorch.backends.apple.coreml.compiler import CoreMLBackend

from executorch.backends.apple.coreml.partition import CoreMLPartitioner
from executorch.backends.apple.coreml.quantizer import CoreMLQuantizer
from executorch.exir import (
    EdgeCompileConfig,
    EdgeProgramManager,
    ExecutorchBackendConfig,
    ExecutorchProgramManager,
    to_edge_transform_and_lower,
)

from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e

from torch.fx import GraphModule


# TODO: check these two params
_EDGE_COMPILE_CONFIG = EdgeCompileConfig(
    _check_ir_validity=False,
    _skip_dim_order=True,
)

# TODO: check these two types
# Using CoreMLBackend.MODEL_TYPE.COMPILED_MODEL and doing compilation ahead of time
# should improve the first time on-device model load time
MODEL_TYPE: CoreMLBackend.MODEL_TYPE = CoreMLBackend.MODEL_TYPE.MODEL
# MODEL_TYPE = CoreMLBackend.MODEL_TYPE.COMPILED_MODEL

COMPUTE_UNIT: ct.ComputeUnit = ct.ComputeUnit.ALL
COMPUTE_PRECISION: ct.precision = ct.precision.FLOAT16


def lower_to_coreml_quantized(
    module: torch.nn.Module,
    example_inputs: tuple[Any, ...],  # pyre-ignore
    min_deployment_target: ct.target,
) -> EdgeProgramManager:
    module.eval()

    compile_specs = CoreMLBackend.generate_compile_specs(
        compute_unit=COMPUTE_UNIT,
        minimum_deployment_target=min_deployment_target,
        compute_precision=COMPUTE_PRECISION,
        model_type=MODEL_TYPE,
    )

    coreml_partitioner = CoreMLPartitioner(
        compile_specs=compile_specs,
    )

    # Export the model for training (pre-autograd ATen dialect)
    # pyre-fixme[9]: graph_module is declared to have type `GraphModule` but is used as type `Module`.
    graph_module: GraphModule = torch.export.export_for_training(
        module, example_inputs, strict=True
    ).module()

    # Define a LinearQuantizerConfig and create an instance of a CoreMLQuantizer
    quantization_config = LinearQuantizerConfig.from_dict(
        {
            "global_config": {
                "quantization_scheme": QuantizationScheme.affine,
                "activation_dtype": torch.quint8,
                "weight_dtype": torch.qint8,
                "weight_per_channel": True,
            }
        }
    )

    quantizer = CoreMLQuantizer(quantization_config)

    # Prepare the model for quantization
    prepared_graph = prepare_pt2e(graph_module, quantizer)

    # Calibrate the model
    # TODO(grinvald): Replace with representative calibration data
    prepared_graph(*example_inputs)

    # Convert the calibrated model to a quantized model
    quantized_model = convert_pt2e(prepared_graph)

    # Single-graph, ATen dialect
    exported_program: torch.export.ExportedProgram = torch.export.export(
        quantized_model, example_inputs, strict=False
    )

    print("Exported program:", exported_program)

    # Lower to CoreML (Edge dialect)
    edge_program: EdgeProgramManager = to_edge_transform_and_lower(
        programs=exported_program,
        partitioner=[coreml_partitioner],
        # compile_config=_EDGE_COMPILE_CONFIG,
    )

    return edge_program


def edge_program_to_et(edge_program: EdgeProgramManager) -> ExecutorchProgramManager:
    return edge_program.to_executorch(
        config=ExecutorchBackendConfig(extract_delegate_segments=True)
    )


class DummyModel(nn.Module):
    def __init__(self):
        super().__init__()

        # This variant makes io18.gather able to be scheduled on ANE
        # index = [1, 1, 1, 1, 1, 1, 1, 1, 1]

        # This variant prevents ios18.gather from being scheduled on ANE
        index = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
        self.register_buffer("index", torch.LongTensor(index), persistent=False)

    def forward(self, x):
        ret = x[:, :, self.index]
        return ret


model = DummyModel()
data_loader = [(torch.randn(1, 80, 200),) for _ in range(10)]
model(*data_loader[0])

example_inputs = data_loader[0]

edge_program = lower_to_coreml_quantized(
    model,
    example_inputs,
    min_deployment_target=ct.target.iOS18,
)

et_program = edge_program_to_et(edge_program)

I'm then saving the ET program and then I'm extracting the CoreML *.mlpackage files (https://docs.pytorch.org/executorch/stable/backends-coreml.html#extracting-the-mlpackage) and opening them with XCode and running a Perfomance Report on my iPhone 15 Pro.

Even though the above repro model exhibits a very simple slicing/indexing operation, the ios18.gather op is not delegate-able on ANE (see image below)

yet, when I change the dummy model to shorten the indexing list ( # index = [1, 1, 1, 1, 1, 1, 1, 1, 1] as described also in the comments in the model code) then everything is delegate-able on ANE (see image below)

It's seems weird to me that such a small difference can make or break delegation of this op to ANE. Do you have any idea why that could be? I think solving this mystery will help me make my original, more complex model also delegate fully to ANE.

Versions

[I'm running from Meta's internal environment, but sharing what I can about the env anyway]

Collecting environment information...
PyTorch version: 2.8.0a0+fb
Is debug build: False
CUDA used to build PyTorch: 12.4.0
ROCM used to build PyTorch: N/A

OS: CentOS Stream 9 (x86_64)
GCC version: (GCC) 11.5.0 20240719 (Red Hat 11.5.0-5)
Clang version: Could not collect
CMake version: version 3.26.5
Libc version: glibc-2.34

Python version: 3.10.5+cinder (cinder/3.10:dcba7ea, May 14 2024, 14:29:26) [Clang 17.0.4 (mononoke://mononoke.internal.tfbnw.net/fbsource b40a4deb90605a472 (64-bit runtime)
Python platform: Linux-6.4.3-0_fbk15_hardened_2630_gf27365f948db-x86_64-with-glibc2.34
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA PG509-210
Nvidia driver version: 550.90.07
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 46 bits physical, 48 bits virtual
Byte Order: Little Endian
CPU(s): 22
On-line CPU(s) list: 0-21
Vendor ID: GenuineIntel
Model name: Intel(R) Xeon(R) Platinum 8339HC CPU @ 1.80GHz
CPU family: 6
Model: 85
Thread(s) per core: 1
Core(s) per socket: 22
Socket(s): 1
Stepping: 11
BogoMIPS: 3591.72
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl xtopology cpuid tsc_known_freq pni pclmulqdq vmx ssse3 fma cx16 pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx512_bf16 arat vnmi umip pku ospke avx512_vnni md_clear flush_l1d arch_capabilities
Virtualization: VT-x
Hypervisor vendor: KVM
Virtualization type: full
L1d cache: 704 KiB (22 instances)
L1i cache: 704 KiB (22 instances)
L2 cache: 88 MiB (22 instances)
L3 cache: 16 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-21
Vulnerability Gather data sampling: Unknown: Dependent on hypervisor status
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Vulnerable
Vulnerability Retbleed: Vulnerable
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1: Vulnerable: __user pointer sanitization and usercopy barriers only; no swapgs barriers
Vulnerability Spectre v2: Vulnerable, IBPB: disabled, STIBP: disabled, PBRSB-eIBRS: Vulnerable
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Mitigation; TSX disabled

Versions of relevant libraries:
[pip3] Could not collect
[conda] Could not collect

Start with the DummyModel repro and the lower_to_coreml_quantized and edge_program_to_et entry points, then inspect CoreMLPartitioner and to_edge_transform_and_lower behavior for the two index lengths. Extract the generated mlpackage files and compare the Xcode Performance Reports; done should explain why ios18.gather changes ANE schedulability and identify a verified fix or limitation.