Skip to content

Instantly share code, notes, and snippets.

@Khrob
Last active Mar 30, 2021
Embed
What would you like to do?
///
/// This is a 'single' function way to get a UIImage from a Metal shader
/// You still need to set up an MTLibrary in your app, but you should be
/// able to fairly readily drop this into a standard XCode app and as long
/// as you have the vertex and fragment shaders compiled into the library
/// you should be good to go.
///
/**
In your shader file:
struct Vertex
{
float4 position [[ position ]];
float2 uv;
};
vertex
Vertex
vertex_func (
const device Vertex * vertices [[ buffer(0) ]],
uint id [[ vertex_id ]] )
{
return vertices[id];
}
fragment
float4
test_func (Vertex v [[ stage_in ]])
{
// Re-map 0,0 to the centre and rescale from -1,1 vertically
float w,h = 240.
float x = v.position.x/w;
float y = v.position.y/h;
float ratio = w/h;
x = ((2.0*x)-1.0)*ratio;
y = (2.0*y)-1.0;
// Colour-code distance to the centre
return {abs(x),abs(y),0,1};
}
**/
import MetalKit
func generate_image (_ drawing_function:String, width:Int, height:Int) -> UIImage
{
/// Very rough timing of the function
let start = Date(); defer { print("\(#function) took \(-start.timeIntervalSinceNow)") }
/// Simple geometry to get fragments to work with
/// (Don't strictly need the uv variable, but hey)
struct Vertex
{
var position : simd_float4 = [0,0,0,0]
var uv : simd_float2 = [0,0]
}
let view_corner_verts : [Vertex] = [
Vertex(position: simd_float4(-1.0,-1.0, 0.0, 1.0), uv: simd_float2(0,1)),
Vertex(position: simd_float4(-1.0, 1.0, 0.0, 1.0), uv: simd_float2(0,0)),
Vertex(position: simd_float4( 1.0,-1.0, 0.0, 1.0), uv: simd_float2(1,1)),
Vertex(position: simd_float4( 1.0, 1.0, 0.0, 1.0), uv: simd_float2(1,0)),
]
/// Set up the stuff Metal needs
let device = MTLCreateSystemDefaultDevice()!
let library = device.makeDefaultLibrary()!
let command_queue = device.makeCommandQueue()!
let command_buffer = command_queue.makeCommandBuffer()!
/// Create an MTLTexture to render into
let texture_descriptor = MTLTextureDescriptor()
texture_descriptor.width = width
texture_descriptor.height = height
texture_descriptor.pixelFormat = .bgra8Unorm
texture_descriptor.usage = [.renderTarget]
let texture_output = device.makeTexture(descriptor: texture_descriptor)!
/// Create a render pipeline to do the rendering
let pass_descriptor = MTLRenderPassDescriptor()
pass_descriptor.colorAttachments[0].texture = texture_output
let pipeline_descriptor = MTLRenderPipelineDescriptor()
pipeline_descriptor.colorAttachments[0].pixelFormat = texture_output.pixelFormat
pipeline_descriptor.vertexFunction = library.makeFunction(name: "vertex_func")!
pipeline_descriptor.fragmentFunction = library.makeFunction(name: drawing_function)!
let pass_pipeline = try! device.makeRenderPipelineState(descriptor: pipeline_descriptor)
/// Do the rendering
let picture_command_encoder = command_buffer.makeRenderCommandEncoder(descriptor: pass_descriptor)!
picture_command_encoder.setRenderPipelineState(pass_pipeline)
picture_command_encoder.setVertexBytes(view_corner_verts, length: MemoryLayout<Vertex>.stride * view_corner_verts.count, index: 0)
picture_command_encoder.drawPrimitives(type: .triangleStrip, vertexStart: 0, vertexCount: view_corner_verts.count)
picture_command_encoder.endEncoding()
/// Wait until the GPU is done
command_buffer.commit()
command_buffer.waitUntilCompleted()
/// Convert the MTLTexture into a UIImage and return it
let cii = CIImage(mtlTexture: texture_output, options: nil)!
let image = UIImage(ciImage: cii)
return image
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment