Search code examples
swiftmetalmetalkitxcassetmtktextureloader

How do I use MetalKit texture loader with Metal heaps?


I have a set of Metal textures that are stored in an Xcode Assets Catalog as Texture Sets. I'm loading these using MTKTextureLoader.newTexture(name:scaleFactor:bundle:options).

I then use a MTLArgumentEncoder to encode all of the textures into a Metal 2 argument buffer.

This works great. However, the Introducing Metal 2 WWDC 2017 session recommends combining argument buffers with resource heaps for even better performance, and I'm quite keen to try this. According to the Argument Buffer documentation, instead of having to call MTLRenderCommandEncoder.useResource on each texture in the argument buffer, you just call useHeap on the heap that the textures were allocated from.

However, I haven't found a straightforward way to use MTKTextureLoader together with MTLHeap. It doesn't seem to have a loading option to allocate the texture from a heap.

I'm guessing that the approach would be:

  • load the textures with MTKTextureLoader
  • reverse-engineer a set of MTLTextureDescriptor objects for each texture
  • use the texture descriptors to create an appropriately sized MTLHeap
  • assign a new set of textures from the MTLHeap
  • use some method to copy the textures across, perhaps replaceBytes or maybe even a MTLBlitCommandEncoder
  • deallocate the original textures loaded with the MTKTextureLoader

It seems like a fairly long-winded approach, and i've not seen any examples of this, so I thought I'd ask here first in case I'm missing something obvious.

Should I abandon MTKTextureLoader, and search out some pre-MetalKit art on loading textures from asset catalogs?

I'm using Swift, but happy to accept Objective-C answers.


Solution

  • Well, the method I outlined above seems to work. As predicted, it's pretty long-winded. I'd be very interested to know if anyone has anything more elegant.

    enum MetalError: Error {
        case anErrorOccured
    }
    
    extension MTLTexture {
        var descriptor: MTLTextureDescriptor {
            let descriptor = MTLTextureDescriptor()
            descriptor.width = width
            descriptor.height = height
            descriptor.depth = depth
            descriptor.textureType = textureType
            descriptor.cpuCacheMode = cpuCacheMode
            descriptor.storageMode = storageMode
            descriptor.pixelFormat = pixelFormat
            descriptor.arrayLength = arrayLength
            descriptor.mipmapLevelCount = mipmapLevelCount
            descriptor.sampleCount = sampleCount
            descriptor.usage = usage
            return descriptor
        }
    
        var size: MTLSize {
            return MTLSize(width: width, height: height, depth: depth)
        }
    }
    
    extension MTKTextureLoader {
        func newHeap(withTexturesNamed names: [String], queue: MTLCommandQueue, scaleFactor: CGFloat, bundle: Bundle?, options: [MTKTextureLoader.Option : Any]?, onCompletion: (([MTLTexture]) -> Void)?) throws -> MTLHeap {
            let device = queue.device
            let sourceTextures = try names.map { name in
                return try newTexture(name: name, scaleFactor: scaleFactor, bundle: bundle, options: options)
            }
            let storageMode: MTLStorageMode = .private
            let descriptors: [MTLTextureDescriptor] = sourceTextures.map { source in
                let desc = source.descriptor
                desc.storageMode = storageMode
                return desc
            }
            let sizeAligns = descriptors.map { device.heapTextureSizeAndAlign(descriptor: $0) }
            let heapDescriptor = MTLHeapDescriptor()
            heapDescriptor.size = sizeAligns.reduce(0) { $0 + $1.size }
            heapDescriptor.cpuCacheMode = descriptors[0].cpuCacheMode
            heapDescriptor.storageMode = storageMode
            guard let heap = device.makeHeap(descriptor: heapDescriptor),
                let buffer = queue.makeCommandBuffer(),
                let blit = buffer.makeBlitCommandEncoder()
                else {
                throw MetalError.anErrorOccured
            }
            let destTextures = descriptors.map { descriptor in
                return heap.makeTexture(descriptor: descriptor)
            }
            let origin = MTLOrigin()
            zip(sourceTextures, destTextures).forEach {(source, dest) in
                blit.copy(from: source, sourceSlice: 0, sourceLevel: 0, sourceOrigin: origin, sourceSize: source.size, to: dest, destinationSlice: 0, destinationLevel: 0, destinationOrigin: origin)
                blit.generateMipmaps(for: dest)
            }
            blit.endEncoding()
            buffer.addCompletedHandler { _ in
                onCompletion?(destTextures)
            }
            buffer.commit()
            return heap
        }
    }