Is there support to run ANE accelerated loops/while_loop?
- 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
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