Take a snapshot of current screen with Metal in swift

To make a screenshot, you need to get MTLTexture of the frame buffer.

1. If you use MTKView:

let texture = view.currentDrawable!.texture

2. If you don't use MTKView

Here's what I would do - I would have a property which holds last drawable presented to the screen:

let lastDrawableDisplayed: CAMetalDrawable?

And then when you present drawable to the screen, I would update it:

let commandBuffer = commandQueue.commandBuffer()
commandBuffer.addCompletedHandler { buffer in
  self.lastDrawableDisplayed = drawable
}

Now you whenever you need to take a screenshot, you can get a texture like this:

let texture = lastDrawableDisplayed.texture

Ok, now when you have MTLTexture you can convert it to CGImage and then to UIImage or NSImage.

Here's the code for OS X playground (MetalKit.MTLTextureLoader is not available for iOS playgrounds), in which I convert MTLTexture to CGImage

I made a small extension over MTLTexture for this.

import Metal
import MetalKit
import Cocoa

let device = MTLCreateSystemDefaultDevice()!
let textureLoader = MTKTextureLoader(device: device)

let path = "path/to/your/image.jpg"
let data = NSData(contentsOfFile: path)!

let texture = try! textureLoader.newTextureWithData(data, options: nil)

extension MTLTexture {
  
  func bytes() -> UnsafeMutablePointer<Void> {
    let width = self.width
    let height = self.height
    let rowBytes = self.width * 4
    let p = malloc(width * height * 4)
    
    self.getBytes(p, bytesPerRow: rowBytes, fromRegion: MTLRegionMake2D(0, 0, width, height), mipmapLevel: 0)
    
    return p
  }
  
  func toImage() -> CGImage? {
    let p = bytes()
    
    let pColorSpace = CGColorSpaceCreateDeviceRGB()
    
    let rawBitmapInfo = CGImageAlphaInfo.NoneSkipFirst.rawValue | CGBitmapInfo.ByteOrder32Little.rawValue
    let bitmapInfo:CGBitmapInfo = CGBitmapInfo(rawValue: rawBitmapInfo)
    
    let selftureSize = self.width * self.height * 4
    let rowBytes = self.width * 4
    let provider = CGDataProviderCreateWithData(nil, p, selftureSize, nil)
    let cgImageRef = CGImageCreate(self.width, self.height, 8, 32, rowBytes, pColorSpace, bitmapInfo, provider, nil, true, CGColorRenderingIntent.RenderingIntentDefault)!
    
    return cgImageRef
  }
}

if let imageRef = texture.toImage() {
  let image = NSImage(CGImage: imageRef, size: NSSize(width: texture.width, height: texture.height))
}

For swift 4.0, Just converting code provided by haawa

let lastDrawableDisplayed = metalView?.currentDrawable?.texture

if let imageRef = lastDrawableDisplayed?.toImage() {
    let uiImage:UIImage = UIImage.init(cgImage: imageRef)
}

extension MTLTexture {

    func bytes() -> UnsafeMutableRawPointer {
        let width = self.width
        let height   = self.height
        let rowBytes = self.width * 4
        let p = malloc(width * height * 4)

        self.getBytes(p!, bytesPerRow: rowBytes, from: MTLRegionMake2D(0, 0, width, height), mipmapLevel: 0)

        return p!
    }

    func toImage() -> CGImage? {
        let p = bytes()

        let pColorSpace = CGColorSpaceCreateDeviceRGB()

        let rawBitmapInfo = CGImageAlphaInfo.noneSkipFirst.rawValue | CGBitmapInfo.byteOrder32Little.rawValue
        let bitmapInfo:CGBitmapInfo = CGBitmapInfo(rawValue: rawBitmapInfo)

        let selftureSize = self.width * self.height * 4
        let rowBytes = self.width * 4
        let releaseMaskImagePixelData: CGDataProviderReleaseDataCallback = { (info: UnsafeMutableRawPointer?, data: UnsafeRawPointer, size: Int) -> () in
            return
        }
        let provider = CGDataProvider(dataInfo: nil, data: p, size: selftureSize, releaseData: releaseMaskImagePixelData)
        let cgImageRef = CGImage(width: self.width, height: self.height, bitsPerComponent: 8, bitsPerPixel: 32, bytesPerRow: rowBytes, space: pColorSpace, bitmapInfo: bitmapInfo, provider: provider!, decode: nil, shouldInterpolate: true, intent: CGColorRenderingIntent.defaultIntent)!

        return cgImageRef
    }
}

I didn't manage to get the accepted answer to work in Swift 4 / Metal 2 with XCode 9.1 on an iPhone 6s. Therefore I used a slightly different approach assuming lastDrawableDisplayed is saved as described in the accepted answer. Quick and dirty and without any exception handling:

let context = CIContext()
let texture = self.lastDrawableDisplayed!.texture
let cImg = CIImage(mtlTexture: texture, options: nil)!
let cgImg = context.createCGImage(cImg, from: cImg.extent)!
let uiImg = UIImage(cgImage: cgImg)

This is based on the documentation on the used CIImage Initializer:

init(mtlTexture:options:) Initializes an image object with data supplied by a Metal texture.

and CIImage Processing which describes how to create a CGImage with the use of CIContext:

CIContext() Create[s] a CIContext object (with default options) [...] context.createCGImage Render[s] the output image to a Core Graphics image that you can display or save to a file.

Hope that helps for anyone using Swift 4.

Edit: Additionally, I have multiple overlaying CAMetalLayer in my project and want to combine them into one single UIImage. Therefore it is needed to have references to the last CAMetalDrawable object of each layer. Before a new layer is added (and therefore used as the provider of nextDrawable()) I simply add the lastDrawableDisplayed to an array [CAMetalDrawable]. When "exporting" the layers I simply write all UIImages subsequently into a bitmap-based graphics context and get the final image with UIGraphicsGetImageFromCurrentImageContext().

Edit: If you are having trouble with orientation, try the following:

let uiImg = UIImage(cgImage: cgImg, scale: 1.0, orientation: UIImageOrientation.downMirrored)