0

I am trying to create a CoreML Custom layer that runs on the GPU, using Objective-C for CoreML setup and Metal for GPU programming.

I have created the CoreML model with the custom layer and can successfully execute on the GPU, I wish to create an MTLBuffer from an input MTLTexture in my setup actual GPU execution, although I can't seem to do so, or get access to the memory address to the MTLTexture memory.

When defining a custom layer in CoreML to run on the GPU, the following function needs to be defined, with the given prototype;

(BOOL) encodeToCommandBuffer:(id<MTLCommandBuffer>)commandBuffer inputs:(NSArray<id<MTLTexture>> *)inputs outputs:(NSArray<id<MTLTexture>> *)outputs error:(NSError *__autoreleasing  _Nullable *)error{

    // GPU Setup, moving data, encoding, execution and so on here

}

Here, the inputs are passed as an NSArray of MTLTexture's, we then pass these texture's on to the Metal Shader for computation. My problem is that I want to pass an MTLBuffer to the Metal Shader, which points to the input data, say inputs[0], but I am having troubling copying the input MTLTexture to an MTLBuffer.

I have tried using the MTLBlitCommandEncoder to copy the data from the MTLTexture to an MTLBuffer like so;

id<MTLBuffer> test_buffer = [command_PSO.device newBufferWithLength:(8) options:MTLResourceStorageModeShared];
id <MTLBlitCommandEncoder> blitCommandEncoder = [commandBuffer blitCommandEncoder];
[blitCommandEncoder copyFromTexture:inputs[0]
                            sourceSlice:0
                            sourceLevel:0
                           sourceOrigin:MTLOriginMake(0, 0, 0)
                             sourceSize:MTLSizeMake(1, 1, 1)
                               toBuffer:test_buffer
                      destinationOffset:0
                 destinationBytesPerRow:8
               destinationBytesPerImage:8];
[blitCommandEncoder endEncoding];

The above example should copy a single pixel from the MTLTexture, inputs[0], to the MTLBuffer, test_buffer, but this is not the case.

MTLTextures, getBytes also doesn't work as the inputs have MTLResourceStorageModePrivate set.

When I inspect the input MTLTexture I note that the attribute buffer = <null> and I'm wondering if this could be an issue since the texture was not created from a buffer, and perhaps doesn't store the address to memory easily, but surely we should be able to get the memory address somewhere?

For further reference, here is the input MTLTexture definition;

<CaptureMTLTexture: 0x282469500> -> <AGXA14FamilyTexture: 0x133d9bb00>
    label = <none> 
    textureType = MTLTextureType2DArray 
    pixelFormat = MTLPixelFormatRGBA16Float 
    width = 8 
    height = 1 
    depth = 1 
    arrayLength = 1 
    mipmapLevelCount = 1 
    sampleCount = 1 
    cpuCacheMode = MTLCPUCacheModeDefaultCache 
    storageMode = MTLStorageModePrivate 
    hazardTrackingMode = MTLHazardTrackingModeTracked 
    resourceOptions = MTLResourceCPUCacheModeDefaultCache MTLResourceStorageModePrivate MTLResourceHazardTrackingModeTracked  
    usage = MTLTextureUsageShaderRead MTLTextureUsageShaderWrite 
    shareable = 0 
    framebufferOnly = 0 
    purgeableState = MTLPurgeableStateNonVolatile 
    swizzle = [MTLTextureSwizzleRed, MTLTextureSwizzleGreen, MTLTextureSwizzleBlue, MTLTextureSwizzleAlpha] 
    isCompressed = 0 
    parentTexture = <null> 
    parentRelativeLevel = 0 
    parentRelativeSlice = 0 
    buffer = <null> 
    bufferOffset = 0 
    bufferBytesPerRow = 0 
    iosurface = 0x0 
    iosurfacePlane = 0 
    allowGPUOptimizedContents = YES
    label = <none>

1 Answers1

0

In your snippet, destinationBytesPerRow and destinationBytesPerImage are calculated incorrectly. Referring to documentation

destinationBytesPerRow The stride in bytes between rows of the source texture memory. The value must be a multiple of the source texture's pixel size, in bytes. The value must be less than or equal to 32,767 multiplied by the source texture’s pixel size.

destinationBytesPerImage For 3D textures and 2D array textures, the stride in bytes between 2D images of the source buffer memory. The value must be a multiple of the source texture's pixel size, in bytes.

Your texture is 8 pixels wide and has .rgba16Float format, so destinationBytesPerRow should be 8 * 8. Same for destinationBytesPerImage, it seems like it should be 64 in your case.

JustSomeGuy
  • 3,677
  • 1
  • 23
  • 31
  • Thank you! Yes, correct, the texture in this case is 8 pixels, but, I am only trying to copy a single pixel (the first one) into the buffer to verify it works correctly, I assume this is achievable and we don't have to copy the entire texture. – aaron_dees Apr 29 '22 at 09:56
  • I'm pretty sure it doesn't copy the entire texture, but it still needs you to calculate the right `bytesPerRow` and `perImage`. – JustSomeGuy Apr 29 '22 at 18:15
  • Actually, I might be wrong. – JustSomeGuy Apr 29 '22 at 18:26
  • By the way, `-[MTLTexture buffer]` property is `nil` cause your texture is created from an `MTLDevice` and not an `MTLBuffer`, so it's okay. If you are gonna be using this texture just to copy it to `MTLBuffer` anyway you can try allocating it from a buffer instead, so you can deal with it as a buffer. – JustSomeGuy Apr 29 '22 at 18:29
  • Another thing you can try is capturing it with Frame Capture in Xcode and see if the app is doing what you think it should be doing. – JustSomeGuy Apr 29 '22 at 18:30
  • Thanks again @JustSomeGuy, unfortunately it seems the allocation is handled by CoreML and we can only access the texture through `inputs` to the `encodeToCommandBuffer` function, https://developer.apple.com/documentation/coreml/mlcustomlayer/2936859-encodetocommandbuffer?language=objc, unless you know some way of explicitly dealing with this allocation in CoreML? I will also take a look at the Frame Capture to inspect the behaviour, thanks for your help! – aaron_dees May 10 '22 at 08:31
  • Not sure how it's handled by CoreML if you are passing your textures as `inputs` and `outputs`. You should be able to allocate any texture you want. – JustSomeGuy May 10 '22 at 17:16