LensWarp.swift (11941B)
1 // Polynomial-3K equations adapted from Monado u_distortion_mesh.c. 2 // Copyright 2019-2021, Collabora, Ltd. 3 // SPDX-License-Identifier: BSL-1.0 (see LICENSES/BSL-1.0.txt) 4 import Foundation 5 import Metal 6 import simd 7 8 struct HeadsetOptics { 9 let calibration: HeadsetCalibration 10 let bounds: [Poly3KBounds] 11 static var cacheURL: URL { 12 FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0] 13 .appendingPathComponent("ReverbM1Demo/factory.json") 14 } 15 16 init(json: Data) throws { 17 calibration = try HeadsetCalibration(json: json) 18 bounds = calibration.eyes.enumerated().map { $0.element.distortionBounds(view: $0.offset) } 19 for (index, eye) in calibration.eyes.enumerated() { 20 let b = bounds[index] 21 guard eye.displaySize == SIMD2(4320, 2160), 22 b.texXRange.x.isFinite, b.texXRange.y.isFinite, 23 b.texYRange.x.isFinite, b.texYRange.y.isFinite, 24 b.texXRange.x < 0, b.texXRange.y > 0, 25 b.texYRange.x < 0, b.texYRange.y > 0 else { 26 throw NSError(domain: "Optics", code: 1, userInfo: [NSLocalizedDescriptionKey: "Invalid Reverb G1 optical calibration for eye \(index)"]) 27 } 28 } 29 } 30 31 static func loadCached() throws -> HeadsetOptics { 32 let url = option("--calibration").map { URL(fileURLWithPath: $0) } ?? cacheURL 33 return try HeadsetOptics(json: Data(contentsOf: url)) 34 } 35 } 36 37 private struct WarpUniforms { 38 var inverseAffine: simd_float3x3 39 var red: SIMD4<Float> 40 var green: SIMD4<Float> 41 var blue: SIMD4<Float> 42 var k3: SIMD4<Float> 43 var ranges: SIMD4<Float> 44 var panel: SIMD4<Float> 45 var tuning: SIMD4<Float> 46 } 47 48 /// Render the scene to an oversampled target, then correct geometry and each color separately. 49 final class LensWarp { 50 let color: MTLTexture 51 let depth: MTLTexture 52 private let pipeline: MTLRenderPipelineState 53 private let gridPipeline: MTLRenderPipelineState 54 private var uniforms: [WarpUniforms] = [] 55 56 init(device: MTLDevice, optics: HeadsetOptics) throws { 57 // 1.4x per axis preserves more central detail through the nonlinear resampling. 58 let width = 6048, height = 3024 59 func texture(_ format: MTLPixelFormat, _ usage: MTLTextureUsage) throws -> MTLTexture { 60 let d = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: format, width: width, height: height, mipmapped: false) 61 d.storageMode = .private; d.usage = usage 62 guard let t = device.makeTexture(descriptor: d) else { 63 throw NSError(domain: "Optics", code: 2, userInfo: [NSLocalizedDescriptionKey: "Cannot allocate optical render target"]) 64 } 65 return t 66 } 67 color = try texture(.bgra8Unorm, [.renderTarget, .shaderRead]) 68 depth = try texture(.depth32Float, .renderTarget) 69 let library = try device.makeLibrary(source: lensWarpShader, options: nil) 70 func makePipeline(_ fragment: String) throws -> MTLRenderPipelineState { 71 let p = MTLRenderPipelineDescriptor() 72 p.vertexFunction = library.makeFunction(name: "warp_vertex") 73 p.fragmentFunction = library.makeFunction(name: fragment) 74 p.colorAttachments[0].pixelFormat = .bgra8Unorm 75 if fragment == "optics_grid" { p.depthAttachmentPixelFormat = .depth32Float } 76 return try device.makeRenderPipelineState(descriptor: p) 77 } 78 pipeline = try makePipeline("warp_rgb") 79 gridPipeline = try makePipeline("optics_grid") 80 for index in 0..<2 { 81 let eye = optics.calibration.eyes[index], b = optics.bounds[index] 82 let m = eye.inverseAffine 83 let inverse = simd_float3x3(columns: (SIMD3<Float>(m.columns.0), SIMD3<Float>(m.columns.1), SIMD3<Float>(m.columns.2))) 84 func channel(_ i: Int) -> SIMD4<Float> { 85 let c = eye.channels[i] 86 return SIMD4(Float(c.eyeCenter.x), Float(c.eyeCenter.y), Float(c.k.x), Float(c.k.y)) 87 } 88 uniforms.append(WarpUniforms(inverseAffine: inverse, red: channel(0), green: channel(1), blue: channel(2), 89 k3: SIMD4(Float(eye.channels[0].k.z), Float(eye.channels[1].k.z), Float(eye.channels[2].k.z), 0), 90 ranges: SIMD4(Float(b.texXRange.x), Float(b.texXRange.y), Float(b.texYRange.x), Float(b.texYRange.y)), 91 panel: SIMD4(4320, 2160, Float(index), Float(width / 2)), tuning: SIMD4(1, 1, 1, 0))) 92 } 93 } 94 95 #if LENS_WARP_CHECKS 96 /// Test-only float output from the exact production shader function, before sampling/quantization. 97 func coordinateSamples(device: MTLDevice, points: [SIMD2<Float>], eye: Int, redGain: Float = 1, blueHorizontalGain: Float = 1, blueVerticalGain: Float = 1) throws -> [SIMD2<Float>] { 98 let kernel = """ 99 kernel void check_coordinates(constant WarpUniforms &u [[buffer(0)]], 100 device const float2 *points [[buffer(1)]], device float2 *out [[buffer(2)]], uint i [[thread_position_in_grid]]) { 101 ColorCoordinates tc = corrected_coordinates(points[i], u); 102 out[3*i] = tc.r; out[3*i+1] = tc.g; out[3*i+2] = tc.b; 103 } 104 """ 105 let lib = try device.makeLibrary(source: lensWarpShader + kernel, options: nil) 106 let pipeline = try device.makeComputePipelineState(function: lib.makeFunction(name: "check_coordinates")!) 107 let input = device.makeBuffer(bytes: points, length: points.count * MemoryLayout<SIMD2<Float>>.stride)! 108 let output = device.makeBuffer(length: points.count * 3 * MemoryLayout<SIMD2<Float>>.stride, options: .storageModeShared)! 109 let command = device.makeCommandQueue()!.makeCommandBuffer()! 110 let enc = command.makeComputeCommandEncoder()! 111 var u = uniforms[eye] 112 u.tuning = SIMD4(redGain, blueHorizontalGain, blueVerticalGain, 0) 113 enc.setComputePipelineState(pipeline) 114 enc.setBytes(&u, length: MemoryLayout<WarpUniforms>.stride, index: 0) 115 enc.setBuffer(input, offset: 0, index: 1) 116 enc.setBuffer(output, offset: 0, index: 2) 117 enc.dispatchThreads(MTLSize(width: points.count, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(32, points.count), height: 1, depth: 1)) 118 enc.endEncoding(); command.commit(); command.waitUntilCompleted() 119 precondition(command.status == .completed) 120 return Array(UnsafeBufferPointer(start: output.contents().assumingMemoryBound(to: SIMD2<Float>.self), count: points.count * 3)) 121 } 122 #endif 123 124 func encodeGrid(_ cb: MTLCommandBuffer, pass: MTLRenderPassDescriptor, greenOnly: Bool = false) { 125 guard let enc = cb.makeRenderCommandEncoder(descriptor: pass) else { return } 126 enc.setRenderPipelineState(gridPipeline) 127 for eye in 0..<2 { 128 enc.setViewport(MTLViewport(originX: Double(color.width / 2 * eye), originY: 0, 129 width: Double(color.width / 2), height: Double(color.height), znear: 0, zfar: 1)) 130 var green: UInt32 = greenOnly ? 1 : 0 131 enc.setFragmentBytes(&green, length: MemoryLayout<UInt32>.stride, index: 1) 132 var ranges = uniforms[eye].ranges 133 enc.setFragmentBytes(&ranges, length: MemoryLayout<SIMD4<Float>>.stride, index: 0) 134 enc.drawPrimitives(type: .triangle, vertexStart: 0, vertexCount: 3) 135 } 136 enc.endEncoding() 137 } 138 139 func encode(_ cb: MTLCommandBuffer, target: MTLTexture, redGain: Float = 1, blueHorizontalGain: Float = 1, blueVerticalGain: Float = 1) { 140 let pass = MTLRenderPassDescriptor() 141 pass.colorAttachments[0].texture = target 142 pass.colorAttachments[0].loadAction = .clear 143 pass.colorAttachments[0].storeAction = .store 144 pass.colorAttachments[0].clearColor = MTLClearColorMake(0, 0, 0, 1) 145 guard let enc = cb.makeRenderCommandEncoder(descriptor: pass) else { return } 146 enc.setRenderPipelineState(pipeline) 147 enc.setFragmentTexture(color, index: 0) 148 for eye in 0..<2 { 149 enc.setViewport(MTLViewport(originX: Double(target.width / 2 * eye), originY: 0, 150 width: Double(target.width / 2), height: Double(target.height), znear: 0, zfar: 1)) 151 var u = uniforms[eye] 152 u.tuning = SIMD4(redGain, blueHorizontalGain, blueVerticalGain, 0) 153 enc.setFragmentBytes(&u, length: MemoryLayout<WarpUniforms>.stride, index: 0) 154 enc.drawPrimitives(type: .triangle, vertexStart: 0, vertexCount: 3) 155 } 156 enc.endEncoding() 157 } 158 } 159 160 private let lensWarpShader = """ 161 #include <metal_stdlib> 162 using namespace metal; 163 struct WarpUniforms { 164 float3x3 inverseAffine; 165 float4 red, green, blue, k3, ranges, panel, tuning; 166 }; 167 struct WarpVertex { float4 position [[position]]; float2 uv; }; 168 vertex WarpVertex warp_vertex(uint id [[vertex_id]]) { 169 float2 p = float2(id == 1 ? 3.0 : -1.0, id == 2 ? 3.0 : -1.0); 170 WarpVertex o; o.position = float4(p, 0, 1); o.uv = float2((p.x + 1) * .5, (1 - p.y) * .5); return o; 171 } 172 float2 channel_uv(float2 uv, float4 channel, float k3, constant WarpUniforms &u) { 173 float2 pix = float2((uv.x + u.panel.z) * u.panel.x * .5, uv.y * u.panel.y) - channel.xy; 174 float r2 = dot(pix, pix); 175 float d = 1 + r2 * (channel.z + r2 * (channel.w + r2 * k3)); 176 float3 vp = u.inverseAffine * float3(pix * d + channel.xy, 1); 177 return (vp.xy / vp.z - u.ranges.xz) / (u.ranges.yw - u.ranges.xz); 178 } 179 float color_sample(texture2d<float> source, float2 tc, uint channel, constant WarpUniforms &u) { 180 if (any(tc < 0) || any(tc > 1)) return 0; 181 // Clamp within this eye's target so interpolation cannot bleed across the stereo seam. 182 float2 halfPixel = float2(.5 / u.panel.w, .5 / source.get_height()); 183 tc = clamp(tc, halfPixel, 1 - halfPixel); 184 tc.x = (tc.x + u.panel.z) * .5; 185 constexpr sampler linearSampler(coord::normalized, address::clamp_to_edge, filter::linear); 186 return source.sample(linearSampler, tc)[channel]; 187 } 188 struct ColorCoordinates { float2 r, g, b; }; 189 ColorCoordinates corrected_coordinates(float2 uv, constant WarpUniforms &u) { 190 ColorCoordinates tc; 191 tc.g = channel_uv(uv, u.green, u.k3.y, u); 192 tc.r = mix(tc.g, channel_uv(uv, u.red, u.k3.x, u), u.tuning.x); 193 tc.b = mix(tc.g, channel_uv(uv, u.blue, u.k3.z, u), u.tuning.yz); 194 return tc; 195 } 196 fragment float4 warp_rgb(WarpVertex in [[stage_in]], texture2d<float> source [[texture(0)]], constant WarpUniforms &u [[buffer(0)]]) { 197 ColorCoordinates tc = corrected_coordinates(in.uv, u); 198 float2 r = tc.r, g = tc.g, b = tc.b; 199 return float4(color_sample(source, r, 0, u), color_sample(source, g, 1, u), color_sample(source, b, 2, u), 1); 200 } 201 fragment float4 optics_grid(WarpVertex in [[stage_in]], constant float4 &ranges [[buffer(0)]], constant uint &greenOnly [[buffer(1)]]) { 202 float2 tangent = mix(ranges.xz, ranges.yw, in.uv); 203 float2 cell = abs(fract(tangent * 10 + .5) - .5); 204 float2 aa = fwidth(tangent * 10); 205 float line = 1 - min(smoothstep(aa.x, aa.x * 2, cell.x), smoothstep(aa.y, aa.y * 2, cell.y)); 206 float cross = min(abs(tangent.x), abs(tangent.y)); 207 float3 color = mix(float3(.01), float3(.8), line); 208 if (cross < .001) color = float3(0, .8, .2); 209 if (greenOnly) color *= float3(0, 1, 0); 210 return float4(color, 1); 211 } 212 """ 213 214 func perspective(bounds: Poly3KBounds, near: Float, far: Float) -> simd_float4x4 { 215 let left = Float(bounds.texXRange.x) * near, right = Float(bounds.texXRange.y) * near 216 let top = -Float(bounds.texYRange.x) * near, bottom = -Float(bounds.texYRange.y) * near 217 return simd_float4x4(columns: ( 218 SIMD4(2 * near / (right - left), 0, 0, 0), 219 SIMD4(0, 2 * near / (top - bottom), 0, 0), 220 SIMD4((right + left) / (right - left), (top + bottom) / (top - bottom), far / (near - far), -1), 221 SIMD4(0, 0, near * far / (near - far), 0) 222 )) 223 }