Neural Engine compiler fails ("Graph traversal failure", Code=22) on a slice of a concat; the model silently falls back to the CPU

Environment: Mac mini, Apple M2 Pro (Mac14,12), macOS 27.2 (26B123). coremltools 9.0, mlprogram, iOS17 target.

Disclosure. I used AI to find and triage this: to bisect the model, isolate the
pattern and draft this post. The reproducer and every result above were run on the machine described.

Summary. I load the model with computeUnits = .cpuAndNeuralEngine. If the program slices a concat's output along the last axis (slice_by_index) and then joins the slice with other tensors (concat or stack), the Neural Engine compiler refuses the program. MLModel still loads without an error, but runs every operation on the CPU.

On a large model, a streaming speech encoder, this costs:

  • a load of 20 s to 6 minutes, at every launch;
  • 1.1 to 1.8 GB of extra memory, because the weights are copied to float32;
  • inference about 15 times slower.

MLComputePlan puts 0 of the model's 1,674 operations on the Neural Engine. If the program computes the same values without slicing the concat, the compiler accepts it, and the whole model runs on the Neural Engine.

The pattern is common. It is the cache update of a streaming convolution: next_cache = concat(cache, x)[..., -K:].

Reproducer. The script needs coremltools 9.0. It writes refused.mlpackage, or fixed.mlpackage with --fix:

import sys
import numpy as np
import coremltools as ct
from coremltools.converters.mil import Builder as mb
from coremltools.converters.mil.mil import types
from coremltools.converters.mil.mil.passes.pass_pipeline import PassPipeline

N, C, T, K = 24, 1024, 7, 8  # layers, channels, new frames per step, cached frames
FIX = "--fix" in sys.argv
rng = np.random.default_rng(0)
W = [(rng.standard_normal((C, 1, K + 1)) * 0.1).astype(np.float16) for _ in range(N)]

@mb.program(input_specs=[mb.TensorSpec(shape=(1, N, C, K), dtype=types.fp16),
                         mb.TensorSpec(shape=(1, C, T), dtype=types.fp16)],
            opset_version=ct.target.iOS17)
def prog(cache, x):
    v = mb.transpose(x=cache, perm=[1, 0, 2, 3])
    outs, h = [], x
    for i in range(N):
        c = mb.slice_by_index(x=v, begin=[i, 0, 0, 0], end=[i + 1, 1, 0, 0],
                              end_mask=[False, False, True, True],
                              squeeze_mask=[True, False, False, False])      # [1, C, K]
        new = mb.concat(values=[c, h], axis=-1)                              # [1, C, K+T]
        if FIX:  # the same values, without cutting the concat
            o = mb.concat(values=[mb.slice_by_index(x=c, begin=[0, 0, T], end=[1, C, K]), h], axis=-1)
        else:    # the next cache: the last K columns of the concat (start 7, width 8)
            o = mb.slice_by_index(x=new, begin=[0, 0, T], end=[1, C, T + K])
        outs.append(o)
        h = mb.conv(x=new, weight=W[i], groups=C, pad_type="valid")          # [1, C, T]
    return mb.transpose(x=mb.stack(values=outs, axis=0), perm=[1, 0, 2, 3]), h

pipe = PassPipeline.EMPTY  # keep the ops as written
pipe.append_pass("common::dead_code_elimination")
m = ct.convert(prog, convert_to="mlprogram", minimum_deployment_target=ct.target.iOS17,
               pass_pipeline=pipe, compute_precision=ct.precision.FLOAT32)
m.save("fixed.mlpackage" if FIX else "refused.mlpackage")

How to see the failure. Start the log, then load each model with .cpuAndNeuralEngine:

log stream --level info --predicate 'process == "ANECompilerService"'
  • refused: the log shows allocation hint propogation fails …, MirInfo propogation fails … and Calling ANE compiler done ret(22). The model then runs on the CPU, and MLComputePlan shows no Neural Engine operations.
  • fixed: the log shows ret(0), and the model runs on the Neural Engine.

With private data logging on, the compiler names the operation: allocation hint propogation fails in y and MirInfo propogation fails in y/src, then Compiler internal error: Graph traversal failure (target h14g). Here y is the concat that joins the slice, and y/src is the slice.

What causes the failure. I changed one property at a time. The compiler fails only when all four conditions hold:

  • a slice_by_index slices a concat's output along the last axis;
  • the slice starts at an offset that is not a multiple of 8 elements (here 7);
  • the slice's width is a multiple of 8 but not of 32 (here 8);
  • a concat or stack then joins the slice with other tensors.

These cases compile: starts 0, 8, 16, 24 and 32, and widths 1, 2, 4, 9, 12, 31, 32, 33 and 64. The result does not
depend on the number of channels (1 to 1,024), the number of layers or slices (2 to 24), or the weights. Slicing the inputs instead (--fix) compiles.

I applied the rewrite to a real model that has this pattern in 24 layers. On the CPU, its outputs are identical to
the original's. On the Neural Engine, its word error rate is the same as the original model's on the CPU.

Questions.

  1. Is this a known limitation of the M2-generation (h14g) compiler, or a bug? Other reports describe similar
    "propogation" or register-allocation failures on M2- and M3-family chips, with streaming encoders that compile on M5.
  2. Could Core ML report the refusal instead of silently loading the model on the CPU? Today the only signals are the log and the memory footprint. Core ML does not cache the failed compile, so every launch repeats it.
  3. Is there a recommended way to write a streaming cache update for the Neural Engine?