Process `grid` tensor by `builder.add_transpose` makes the function always run on CPU

[Tengxu-Sun]

Hi, @dneprDroid
I'm trying to implement a 5-D input tensor version of custom grid_sample based on your work. Now i meet two problems hoping you can help me.
First, you said "we should transpose the grid to [n, 2, w, h] shape (for encoding it to a coreml 2-channel texture)". But i found this process may lead to the function always run on CPU in my implementation. The GPU function encode works without this process. I read your code but haven't find a solution. Have you met this problem ever?

Second, since i want implement a 5-D input tensor version of custom grid_sample, how id<MTLTexture> in encode work at this time? I found the data storage is different with ordinary buffer, and i'm not sure how to read data in id<MTLTexture> so that the data can be calculated in later computation process.
Thank you!

[dneprDroid]

Hi,

1. The case without transpose you can find on this branch fix/without-transpose

2. The case with 5-D input is complicated, not sure about that. You can try to rewrite the shader with texturecube_array input instead of texture2d_array and your grid vector may contain 3 coordinates (grid_X, grid_Y, grid_Z).

[Tengxu-Sun]

Hi,

1. The case without transpose you can find on this branch fix/without-transpose
2. The case with 5-D input is complicated, not sure about that. You can try to rewrite the shader with texturecube_array input instead of texture2d_array and your grid vector may contain 3 coordinates (grid_X, grid_Y, grid_Z).

Thanks for your reply!

1. Cool, i'll see it carefully later.
2. According to guide, MTLTextureTypeCube is a cube texture with six two-dimensional images. I'm not sure weather texturecube_array is suitable for 5-D input since it seems still can't store the contiguous depth of the feature. Meanwhile, i do some tests to print the information of [MTLTexture] in [encode function]. The result shows that its default textureType = MTLTextureType2DArray and seems can't be changed manually. So i guess textureType in Metal shader also need to be texture2d_array. Do you have any suggestions to change the textureType of [MTLTexture] in encode function? Thanks again.

[Tengxu-Sun]

Hi,

1. The case without transpose you can find on this branch fix/without-transpose
2. The case with 5-D input is complicated, not sure about that. You can try to rewrite the shader with texturecube_array input instead of texture2d_array and your grid vector may contain 3 coordinates (grid_X, grid_Y, grid_Z).

1. Update: I test the without transpose branch on Iphone11, and find it not totally equal. The original implementation looks more delicate, while it looks rougher after the change. Please have a look in you spare time, Thanks a lot!

[nighting0le01]

hi @Tengxu-Sun were you able to get the 5d torch.grid_sampler to work in coreml??