apple / apple/coremltools

slice_by_index operation not supported on GPU, defaulting to CPU

Open
#2,550 0 comments 1 reaction 0 assignees View on GitHub
bug
Dominant language
Python
Stars
5.4k
Forks
850
Avg merge
4d 5h
Merged PRs (30d)
10

Description

## 🐞Describing the bug
Running into an issue where my entire model is able to compile via ct.convert and run on GPU except for the slice_by_index operations:

![Image](https://github.com/user-attachments/assets/8d1ba890-ec7f-4f70-97eb-21ed29c9059a)

Interestingly, this only occurs when I run on my iPhone on device and does not when I run on my M4 Macbook Pro, which assigns whole model onto GPU successfully.

The reason I'm asking about this here is because I wonder if there's a way to do the same operation without necessarily slicing or if there's some core code I can edit in the package to fix this. Tried things like calling z.contiguous() prior to slicing, adding static slicing where indices are int constants determined by if statements. Also tried F.interpolate but doesn't result in the same output, and torch.narrow calls mb.slice_by_index() so results in no changes.

## To Reproduce
The issue arises from `z = z[..., self.pad:-self.pad, :]` and `z = z[..., self.pad:self.pad + length]`, which are just [2:-2], [2:-3], [2:-4] dependent on the layer number.
```
class HDecLayerNew(nn.Module):
def __init__(self, chin, chout, last=False, kernel_size=8, stride=4, norm_groups=1, empty=False,
freq=True, dconv=True, norm=True, context=1, dconv_kw={}, pad=True,
context_freq=True, rewrite=True):

super().__init__()
self.pad = kernel_size // 4 if pad else 0
self.last = last
self.freq = freq
self.chin = chin
self.empty = empty
self.stride = stride
self.kernel_size = kernel_size
self.norm = norm
self.context_freq = context_freq

norm_fn = lambda d: nn.Identity() # noqa
if norm:
norm_fn = lambda d: nn.GroupNorm(norm_groups, d) # noqa
self.conv_tr = nn.ConvTranspose1d(chin, chout, kernel_size, stride)
self.norm2 = norm_fn(chout)

if not self.empty:
klass = nn.Conv2d if freq else nn.Conv1d
if rewrite:
self.rewrite = klass(chin, 2 * chin, 1 + 2 * context, 1, context) if context_freq else klass(chin, 2 * chin, [1, 1 + 2 * context], 1, [0, context])
self.norm1 = norm_fn(2 * chin)
else:
self.rewrite = None

self.dconv = DConv(chin, **dconv_kw) if dconv else None

def forward(self, x, skip, length):
if self.freq and x.dim() == 3:
B, C, T = x.shape
x = x.view(B, self.chin, -1, T)

if not self.empty:
x = x + skip
y = F.glu(self.norm1(self.rewrite(x)), dim=1) if self.rewrite else x
if self.dconv:
if self.freq:
B, C, Fr, T = y.shape
y = y.permute(0, 2, 1, 3).reshape(-1, C, T)
y = self.dconv(y)
if self.freq:
y = y.view(B, Fr, C, T).permute(0, 2, 1, 3)
else:
y = x
assert skip is None

if self.freq:
B, C, Fr, T = y.shape
y_reshaped = y.permute(0, 3, 1, 2).reshape(B * T, C, Fr)
z_reshaped = self.conv_tr(y_reshaped)
C_out, Fr_out = z_reshaped.shape[1], z_reshaped.shape[2]
z = z_reshaped.view(B, T, C_out, Fr_out).permute(0, 2, 3, 1)
z = self.norm2(z)
if self.pad:
z = z[..., self.pad:-self.pad, :]
if not self.last:
z = F.gelu(z)
else:
z = self.norm2(self.conv_tr(y))
if self.pad:
z = z[..., self.pad:self.pad + length]
assert z.shape[-1] == length, (z.shape[-1], length)
if not self.last:
z = F.gelu(z)

return z, y
```

## System environment (please complete the following information):
- coremltools version: 8.3.0
- Source dialect: torchscript (torch.jit.trace)
- OS: MacOS 15.3.2, iOS 18.5
- PyTorch version: 2.4.1

Contributor guide

Open the contributing guide

Research direction

Start by reproducing the torch.jit.trace and ct.convert flow with the reported slice expressions, then compare execution on the iPhone GPU and M4 Mac GPU. Trace how slice_by_index is assigned during conversion and verify that the model preserves its output without falling back to CPU on iOS.

Written by the indexing model from the issue text.

Assessment

Tech stack
ios, python, pytorch
Domain
machine-learning, mobile-dev
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.