Issue creating MTLBuffer from MTLTexture used as inputs in CoreML Custom Layer for GPU execution

The code posted looks correct as far as encoding the commands goes. So we'd need a little more context to be sure your doing all the right things. You should not need to create a texture from a buffer to get the pixel values out of a texture.

Some obvious things to check:

• Have you committed the command buffer?
• Do you wait for the command buffer to complete before reading from the buffer with the CPU?

Hi there, many thanks for the quick response.

Below is an example encodeToCommandBuffer function, required by Core ML for Custom layers to run on the GPU

- (BOOL) encodeToCommandBuffer:(id<MTLCommandBuffer>)commandBuffer inputs:(NSArray<id<MTLTexture>> *)inputs outputs:(NSArray<id<MTLTexture>> *)outputs error:(NSError *__autoreleasing _Nullable *)error{
   
  id<MTLBuffer> test_buffer = [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];

  id<MTLComputeCommandEncoder> computeEncoder = [commandBuffer computeCommandEncoderWithDispatchType:MTLDispatchTypeSerial];
  assert(computeEncoder != nil);

  [computeEncoder setComputePipelineState:PSO];
  [computeEncoder setTexture:inputs[0] atIndex:0]; // Input
  [computeEncoder setTexture:outputs[0] atIndex:2]; // Output
   
  MTLSize dimThreadsBlock, dimThreadsGrid;
   
  dimThreadsGrid = MTLSizeMake(1, 1, 1);
  dimThreadsBlock = MTLSizeMake(1, 1, 1);

  [computeEncoder dispatchThreads: dimThreadsGrid threadsPerThreadgroup:dimThreadsBlock];
   
  [computeEncoder endEncoding];

Note that the above does not include [commandBuffer commit]; or [commandBuffer waitUntilCompleted];. If these are added an error occurs, stating;

failed assertion _status < MTLCommandBufferStatusCommitted at line 300 in -[IOGPUMetalCommandBuffer setCurrentCommandEncoder:]

I can run code on the GPU successfully for my Custom Layer without explicitly committing the commandBuffer within the encodeToCommandBuffer function, this make me believe this is handled by Core ML internally (or so the name of the function would imply)?

Note that all setup for GPU is done prior to above function being called, in the initWithParameterDictionary function required by Core ML;

- (instancetype) initWithParameterDictionary:(NSDictionary<NSString *,id> *)parameters error:(NSError *__autoreleasing _Nullable *)error{
 //kernel_height = [parameters[@"kernel_height"] intValue];
 //kernel_width = [parameters[@"kernel_width"] intValue];
   
 //NSLog(@"kh kw %@", [parameters allKeys]);
 //self = [super init]; // This line might introduce breaks !!!!!

 // Metal GPU setup
 id<MTLDevice> device = MTLCreateSystemDefaultDevice();

 NSError* error_pso = nil;
   
 id<MTLLibrary> defaultlibrary = [device newDefaultLibrary];
 if (defaultlibrary == nil){
  NSLog(@"Failed to find the default library");
  return nil;
 }
   
 id<MTLFunction> function = [defaultlibrary newFunctionWithName:@"kernel"];
 if (rans_decode == nil){
  NSLog(@"Failed to find the function");
  return nil;
 }

PSO = [device newComputePipelineStateWithFunction:function error:&error_pso];
 if (PSO == nil || error_pso != nil){
  NSLog(@"Failed to find the default library");
  return nil;
 }
   
 return self;
}

Thanks again!

You are correct. That assertion is letting you know that you are submitting a command buffer twice:

"Don’t commit the command buffer in this method; Core ML executes the command buffer after this method returns." https://developer.apple.com/documentation/coreml/mlcustomlayer/2936859-encodetocommandbuffer

I believe the next thing to check is if your test_buffer resource is actually used/synchronized. How are you reading the data at test_buffer? In your code snippet, it does not look like you are binding the buffer to the compute pipeline, [computeEncoder setBuffer:test_buffer offset:0 atIndex:0], nor reading its contents on the CPU, commandBuffer addCompletedHandler: to print the buffer's contents & to verify that the buffer is populated.

see also "Copying Data from a Private Texture to a Shared Buffer" in https://developer.apple.com/documentation/metal/resource_fundamentals/copying_data_to_a_private_resource?language=objc