apple / apple/coremltools

Is there support to run ANE accelerated loops/while_loop?

Open
#2,208 1 comment 0 reactions 0 assignees View on GitHub
question
Dominant language
Python
Stars
5.4k
Forks
850
Avg merge
4d 5h
Merged PRs (30d)
10

Description

Hi, I've been experimenting with while_loop (s), but I haven't had any success making them run accelerated on the ANE, and neither the GPU. Is it even possible for them to run with acceleration?

Here's an example code, the loop uses just a simple counter as exit condition

```python
import numpy as np

import coremltools as ct
import coremltools.converters.mil as mil
from coremltools.converters.mil import Builder as mb

bsize = 4
seqlen = 128
dim = 512

qshape = (bsize, seqlen, dim)
kshape = (dim, dim)

@mb.program(
input_specs=[
mb.TensorSpec(shape=qshape, dtype=mil.input_types.types.fp16),
mb.TensorSpec(shape=kshape, dtype=mil.input_types.types.fp16),
mb.TensorSpec(shape=(1,), dtype=mil.input_types.types.int32),
],
opset_version=mil.builder.AvailableTarget.iOS17,
)
def loop(q, k, l):
i = mb.fill(shape=np.array([1]), value=np.array(0., dtype=np.int32))
start = q
loop_vars = (i, start)

def cond(_i, state):
return mb.less(x=_i, y=l)

def body(_i, state):
_prod = mb.matmul(x=state, y=k, transpose_y=False)
state = mb.sigmoid(x=_prod)
_i = mb.add(x=_i, y=np.ones([1], dtype=np.int32))
return _i, state

loop_vars = mb.while_loop(_cond=cond, _body=body, loop_vars=loop_vars)
return loop_vars

mlmodel = ct.convert(
loop,
compute_units=ct.ComputeUnit.CPU_AND_NE,
compute_precision=ct.precision.FLOAT16,
minimum_deployment_target=ct.target.iOS17,
inputs=[
ct.TensorType(name='q', shape=ct.Shape(shape=qshape)),
ct.TensorType(name='k', shape=ct.Shape(shape=kshape)),
ct.TensorType(name='l', shape=ct.Shape(shape=(seqlen,))),
]
)

q = np.random.normal(scale=0.2, size=qshape).astype(np.float16)
k = np.random.normal(scale=0.2, size=kshape).astype(np.float16)
l = np.array([16]).astype(np.int32)

mlmodel.predict({'q': q, 'k': k, 'l': l})
```

Contributor guide

Open the contributing guide

Research direction

Start by reproducing the provided Python coremltools program using mb.while_loop, iOS17, and CPU_AND_NE, then inspect how conversion and prediction handle the loop. Done should establish whether while_loop can run on the ANE or GPU, or clearly document the current limitation and supported behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.