Статьи

Металлические функции ядра / вычислительные шейдеры в Swift

Декабрь 31, 2017

Как часть проекта по созданию симуляции реакции диффузии на основе графического процессора 
, я заявил, что на  этих выходных я расскажу об использовании 
Metal  в 
Swift .

В прошлом я проделывал аналогичную работу, ориентируясь на Flash Player и используя 
AGAL . Металл — это язык намного более высокого уровня, чем AGAL: он основан на C ++ с более богатым синтаксисом и включает в себя 
вычислительные функции . В то время как в AGAL для запуска клеточных автоматов я создавал прямоугольник из двух треугольников с помощью вершинного шейдера и выполнял функции диффузии реакции в отдельном фрагментном шейдере, вычислительный шейдер более прямой: я могу получать и устанавливать текстуры и может работать с отдельными пикселями этой текстуры без необходимости использования вершинного шейдера.

Код Swift, который я обсуждаю в этой статье, основан на двух статьях 
Metal By Example : 
Введение в компьютерное программирование в металле  и 
основы обработки изображений в металле . Оба из них включают исходный код Objective-C, так что, надеюсь, моя реализация Swift поможет некоторым. 
Мое приложение  имеет четыре основных шага: инициализировать Metal, создать текстуру Metal из 
UIImage , применить функцию ядра к этой текстуре, преобразовать вновь сгенерированную текстуру обратно в 
UIImage  и отобразить ее. Я использую простой пример шейдера, который изменяет насыщенность входного изображения. поэтому я также добавил ползунок, который изменяет значение насыщенности.

Давайте посмотрим на каждый шаг один за другим:
Инициализация Металла

Инициализация Metal довольно проста: внутри переопределенного viewDidLoad () моего контроллера представления 
я создаю указатель на устройство Metal по умолчанию:
var device: MTLDevice! = nil
    [...]
    device = MTLCreateSystemDefaultDevice()

Мне также нужно создать библиотеку и очередь команд:

defaultLibrary = device.newDefaultLibrary()
    commandQueue = device.newCommandQueue()

Наконец, я добавляю ссылку на свою функцию Metal в библиотеку и синхронно создаю и компилирую состояние конвейерного вычисления:

et kernelFunction = defaultLibrary.newFunctionWithName("kernelShader")
    pipelineState = device.newComputePipelineStateWithFunction(kernelFunction!, error: nil)

KernelShader указывает на функцию обработки изображения насыщенности, написанную на Metal, которая находится в моем   файле Shaders.metal :

kernel void kernelShader(texture2d<float, access::read> inTexture [[texture(0)]],
                         texture2d<float, access::write> outTexture [[texture(1)]],
                         constant AdjustSaturationUniforms &uniforms [[buffer(0)]],
                         uint2 gid [[thread_position_in_grid]])
    {
        float4 inColor = inTexture.read(gid);
        float value = dot(inColor.rgb, float3(0.299, 0.587, 0.114));
        float4 grayColor(value, value, value, 1.0);
        float4 outColor = mix(grayColor, inColor, uniforms.saturationFactor);
        outTexture.write(outColor, gid);
    }
Создание металлической текстуры из UIIMage

Есть несколько шагов по преобразованию 
UIImage  в 
 экземпляр MTLTexture . Я создаю массив 
UInt8  для хранения пустого
CGBitmapInfo , затем использую 
CGContextDrawImage ()  для копирования изображения в растровый контекст 
let image = UIImage(named: "grand_canyon.jpg")
    let imageRef = image.CGImage
        
    let imageWidth = CGImageGetWidth(imageRef)
    let imageHeight = CGImageGetHeight(imageRef)

    let bytesPerRow = bytesPerPixel * imageWidth
        
    var rawData = [UInt8](count: Int(imageWidth * imageHeight * 4), repeatedValue: 0)
  
    let bitmapInfo = CGBitmapInfo(CGBitmapInfo.ByteOrder32Big.toRaw() | CGImageAlphaInfo.PremultipliedLast.toRaw())

    let context = CGBitmapContextCreate(&rawData, imageWidth, imageHeight, bitsPerComponent, bytesPerRow, rgbColorSpace, bitmapInfo)
        
    CGContextDrawImage(context, CGRectMake(0, 0, CGFloat(imageWidth), CGFloat(imageHeight)), imageRef)

После того, как все эти шаги выполнены, я могу создать новую текстуру, используя   метод replaceRegion (), чтобы записать в нее изображение:

let textureDescriptor = MTLTextureDescriptor.texture2DDescriptorWithPixelFormat(MTLPixelFormat.RGBA8Unorm, width: Int(imageWidth), height: Int(imageHeight), mipmapped: true)
        
    texture = device.newTextureWithDescriptor(textureDescriptor)

    let region = MTLRegionMake2D(0, 0, Int(imageWidth), Int(imageHeight))
    texture.replaceRegion(region, mipmapLevel: 0, withBytes: &rawData, bytesPerRow: Int(bytesPerRow))

Я также создаю пустую текстуру, в которую функция ядра будет записывать:

  let outTextureDescriptor = MTLTextureDescriptor.texture2DDescriptorWithPixelFormat(texture.pixelFormat, width: texture.width, height: texture.height, mipmapped: false)
    outTexture = device.newTextureWithDescriptor(outTextureDescriptor)
Вызов функции ядра

Следующий блок работы — установить текстуры и другую переменную в функции kerne и запустить шейдер. Первым шагом является создание экземпляра буфера команд и кодировщика команд:
let commandBuffer = commandQueue.commandBuffer()
    let commandEncoder = commandBuffer.computeCommandEncoder()

… затем установите состояние конвейера (мы получили из device.newComputePipelineStateWithFunction () ранее) и текстуры в кодировщике команд:

commandEncoder.setComputePipelineState(pipelineState)
    commandEncoder.setTexture(texture, atIndex: 0)
    commandEncoder.setTexture(outTexture, atIndex: 1)

Фильтру требуется дополнительный параметр, который определяет величину насыщенности. Это передается в шейдер через MTLBuffer . Чтобы заполнить буфер, я создал небольшую структуру:

struct AdjustSaturationUniforms 
    {
        var saturationFactor: Float
    }

Затем  newBufferWithBytes ()  для передачи моего   значения с плавающей точкой saturationFactor

var saturationFactor = AdjustSaturationUniforms(saturationFactor: self.saturationFactor)
    var buffer: MTLBuffer = device.newBufferWithBytes(&saturationFactor, length: sizeof(AdjustSaturationUniforms), options: nil)
    commandEncoder.setBuffer(buffer, offset: 0, atIndex: 0)

Теперь это доступно внутри шейдера в качестве аргумента его функции ядра:

constant AdjustSaturationUniforms &uniforms [[buffer(0)]]

Теперь я готов вызвать саму функцию. Металлические функции ядра используют группы потоков, чтобы разбить их рабочую нагрузку на куски. В моем примере я создаю 64 группы потоков, а затем отправляю их в графический процессор:

let threadGroupCount = MTLSizeMake(8, 8, 1)
    let threadGroups = MTLSizeMake(texture.width / threadGroupCount.width, texture.height / threadGroupCount.height, 1)
        
    commandQueue = device.newCommandQueue()
        
    commandEncoder.dispatchThreadgroups(threadGroups, threadsPerThreadgroup: threadGroupCount)
    commandEncoder.endEncoding()
    commandBuffer.commit()
    commandBuffer.waitUntilCompleted()
Преобразование текстуры в UIImage

Наконец, теперь, когда функция ядра выполнена, нам нужно сделать все наоборот, чтобы получить изображение, хранящееся в
outTexture,  в 
UIImage,  чтобы его можно было отобразить. Опять же, я использую регион для определения размера и getBytes () текстуры 
 для заполнения массива в 
UInt8 :
let imageSize = CGSize(width: texture.width, height: texture.height)
    let imageByteCount = Int(imageSize.width * imageSize.height * 4)
        
    let bytesPerRow = bytesPerPixel * UInt(imageSize.width)
    var imageBytes = [UInt8](count: imageByteCount, repeatedValue: 0)
    let region = MTLRegionMake2D(0, 0, Int(imageSize.width), Int(imageSize.height))
        
    outTexture.getBytes(&imageBytes, bytesPerRow: Int(bytesPerRow), fromRegion: region, mipmapLevel: 0)

Теперь, когда  imageBytes  содержит необработанные данные, для создания CGImage требуется несколько строк  :

let providerRef = CGDataProviderCreateWithCFData(
            NSData(bytes: &imageBytes, length: imageBytes.count * sizeof(UInt8))
        )
        
    let bitmapInfo = CGBitmapInfo(CGBitmapInfo.ByteOrder32Big.toRaw() | CGImageAlphaInfo.PremultipliedLast.toRaw())
    let renderingIntent = kCGRenderingIntentDefault
        
    let imageRef = CGImageCreate(UInt(imageSize.width), UInt(imageSize.height), bitsPerComponent, bitsPerPixel, bytesPerRow, rgbColorSpace, bitmapInfo, providerRef, nil, false, renderingIntent)
        
    imageView.image = UIImage(CGImage: imageRef)

… и мы закончили! 

Для Metal требуется процессор A7 или A8, и этот код был собран и протестирован под Xcode 6. Весь исходный код доступен в моем 
репозитории GitHub здесь .

50
Share: