How to pass texture buffer data to Shader with Metal?

Question

I would like to work with texture data as a 1D array in a compute shader. I read that the best way is to pass it as a buffer instead of a 1D texture.

I am loading the texture with:

let textureLoader = MTKTextureLoader(device: device)

do {
    if let image = UIImage(named: "testImage") {
        let options = [ MTKTextureLoaderOptionSRGB : NSNumber(value: false) ]
        try kernelSourceTexture = textureLoader.newTexture(with: image.cgImage!, options: options)
            kernelDestTexture = device.makeTexture(descriptor: kernelSourceTexture!.matchingDescriptor())
    } else {
        print("Failed to load texture image from main bundle")
    }
}
catch let error {
    print("Failed to create texture from image, error \(error)")
}

And I am creating the buffer with (not sure if this is correct):

var textureBuffer: MTLBuffer! = nil
var currentVertPtr = kernelSourceTexture!.buffer!.contents()
textureBuffer = device.makeBuffer(bytes: &currentVertPtr, length: kernelSourceTexture!.buffer!.length, options: [])
uniformBuffer.label = "textureData"

How do I pass the buffer to a compute shader? Do I pass it as an argument or as a uniform? What would the buffer's data type be?

Sorry if these are dumb questions, I am just getting started with Metal and I can't find much for reading. I bought and read "Metal by Example: High-performance graphics and data-parallel programming for iOS". Side question, can anyone recommend more books on Metal?

Answer 1

Whether you should pass the data as a buffer or texture depends somewhat on what you want to do with it in your kernel function. If you use a buffer, you won't get several of the benefits of textures: defined behavior when sampling out of bounds, interpolation, and automatic conversion of components from the source pixel format to the component type requested in the shader.

But since you asked about buffers, let's talk about how to create a buffer that contains image data and how to pass it to a kernel.

I'll assume for the sake of discussion that we want our data in the equivalent of .rgba8unorm format, where each component is a single byte.

Creating a texture just for the sake of doing this conversion is wasteful (and as Ken noted in the comments, textures aren't backed by a buffer by default, which complicates how we get their data), so let's set MTKTextureLoader aside and do it ourselves.

Suppose we have an image in our bundle for which we have a URL. Then we can use a method like the following to load it, ensure it's in the desired format, and wrap the data in an MTLBuffer with a minimal number of copies:

func bufferWithImageData(at url: URL, resourceOptions: MTLResourceOptions, device: MTLDevice) -> MTLBuffer? {
    guard let imageSource = CGImageSourceCreateWithURL(url as CFURL, nil) else { return nil }
    if CGImageSourceGetCount(imageSource) != 1 { return nil }
    guard let image = CGImageSourceCreateImageAtIndex(imageSource, 0, nil) else { return nil }
    guard let colorspace = CGColorSpace(name: CGColorSpace.genericRGBLinear) else { return nil }

    let bitsPerComponent = UInt32(8)
    let bytesPerComponent = bitsPerComponent / 8
    let componentCount = UInt32(4)
    let bytesPerPixel = bytesPerComponent * componentCount
    let rowBytes = UInt32(image.width) * bytesPerPixel
    let imageSizeBytes = rowBytes * UInt32(image.height)

    let pageSize = UInt32(getpagesize())
    let allocSizeBytes = (imageSizeBytes + pageSize - 1) & (~(pageSize - 1))

    var dataBuffer: UnsafeMutableRawPointer? = nil
    let allocResult = posix_memalign(&dataBuffer, Int(pageSize), Int(allocSizeBytes))
    if allocResult != noErr { return nil }

    var targetFormat = vImage_CGImageFormat()
    targetFormat.bitsPerComponent = bitsPerComponent
    targetFormat.bitsPerPixel = bytesPerPixel * 8
    targetFormat.colorSpace = Unmanaged.passUnretained(colorspace)
    targetFormat.bitmapInfo = CGBitmapInfo(rawValue: CGImageAlphaInfo.premultipliedLast.rawValue)

    var imageBuffer = vImage_Buffer(data: dataBuffer, height: UInt(image.height), width: UInt(image.width), rowBytes: Int(rowBytes))
    let status = vImageBuffer_InitWithCGImage(&imageBuffer, &targetFormat, nil, image, vImage_Flags(kvImageNoAllocate))
    if status != kvImageNoError {
        free(dataBuffer)
        return nil
    }

    return device.makeBuffer(bytesNoCopy: imageBuffer.data, length: Int(allocSizeBytes), options: resourceOptions, deallocator: { (memory, size) in
        free(memory)
    })
}

(Note that you'll need to import Accelerate in order to use vImage functions.)

Here's an example of how to call this method:

let resourceOptions: MTLResourceOptions = [ .storageModeShared ]
let imageURL = Bundle.main.url(forResource: "my_image", withExtension: "png")!
let inputBuffer = bufferWithImageData(at: imageURL, resourceOptions: resourceOptions, device: device)

This may seem unnecessarily complex, but the beauty of this is that for a huge variety of input formats, we can use vImage to efficiently convert to our desired layout and color space. By changing only a couple of lines, we could go from RGBA8888 to BGRAFFFF, or many other formats.

Create your compute pipeline state and any other resources you want to work with in the usual way. You can pass the buffer you just created by assigning it to any buffer argument slot:

computeCommandEncoder.setBuffer(inputBuffer, offset: 0, at: 0)

Dispatch your compute grid, also in the usual way.

For completeness, here's a kernel function that operates on our buffer. It's by no means the most efficient way to compute this result, but this is just for illustration:

kernel void threshold(constant uchar4 *imageBuffer [[buffer(0)]],
                      device uchar *outputBuffer [[buffer(1)]],
                      uint gid [[thread_position_in_grid]])
{
    float3 p = float3(imageBuffer[gid].rgb);
    float3 k = float3(0.299, 0.587, 0.114);
    float luma = dot(p, k);
    outputBuffer[gid] = (luma > 127) ? 255 : 0;
}

Note:

1. We take the buffer as a uchar4, since each sequence of 4 bytes represents one pixel.
2. We index into the buffer using a parameter attributed with thread_position_in_grid, which indicates the global index into the grid we dispatched with our compute command encoder. Since our "image" is 1D, this position is also one-dimensional.
3. In general, integer arithmetic operations are very expensive on GPUs. It's possible that the time spent doing the integer->float conversions in this function dominates the extra bandwidth of operating on a buffer containing floats, at least on some processors.

Hope that helps. If you tell us more about what you're trying to do, we can make better suggestions about how to load and process your image data.