Application+UpscaleImage.swift

#if os(Windows)
import FidelityFX
import SwiftCOM
import WinSDK
#endif

#if os(Windows)
private func createFFXSurfaceFormat(
  _ format: DXGI_FORMAT
) -> FfxApiSurfaceFormat {
  switch format {
  case DXGI_FORMAT_R10G10B10A2_UNORM:
    return FFX_API_SURFACE_FORMAT_R10G10B10A2_UNORM
  case DXGI_FORMAT_R32_FLOAT:
    return FFX_API_SURFACE_FORMAT_R32_FLOAT
  case DXGI_FORMAT_R16G16_FLOAT:
    return FFX_API_SURFACE_FORMAT_R16G16_FLOAT
  default:
    fatalError("Unrecognized DXGI format.")
  }
}

// Utility for binding DirectX resources.
private func createFFXResource(
  _ d3d12Resource: SwiftCOM.ID3D12Resource
) -> FfxApiResource {
  func createID3D12Resource() -> UnsafeMutableRawPointer {
    let iid = SwiftCOM.ID3D12Resource.IID
    
    // Fetch the underlying pointer without worrying about memory leaks.
    let interface = try! d3d12Resource.QueryInterface(iid: iid)
    _ = try! d3d12Resource.Release()
    
    guard let interface else {
      fatalError("This should never happen.")
    }
    return interface
  }
  
  // Cannot invoke ffxApiGetResourceDX12 from Clang header import.
  var output = FfxApiResource()
  output.resource = createID3D12Resource()
  output.state = UInt32(FFX_API_RESOURCE_STATE_UNORDERED_ACCESS.rawValue)
  
  let desc = try! d3d12Resource.GetDesc()
  output.description.flags = UInt32(
    FFX_API_RESOURCE_FLAGS_NONE.rawValue)
  output.description.usage = UInt32(
    FFX_API_RESOURCE_USAGE_READ_ONLY.rawValue)
  output.description.usage |= UInt32(
    FFX_API_RESOURCE_USAGE_UAV.rawValue)
  
  output.description.width = UInt32(desc.Width)
  output.description.height = UInt32(desc.Height)
  output.description.depth = UInt32(desc.DepthOrArraySize)
  output.description.mipCount = UInt32(desc.MipLevels)
  output.description.type = UInt32(
    FFX_API_RESOURCE_TYPE_TEXTURE2D.rawValue)
  
  let ffxSurfaceFormat = createFFXSurfaceFormat(desc.Format)
  output.description.format = UInt32(ffxSurfaceFormat.rawValue)
  
  return output
}

private func createEmptyFFXResource() -> FfxApiResource {
  var output = FfxApiResource()
  output.resource = nil
  output.state = 0
  output.description.type = 0
  output.description.format = 0
  output.description.width = 0
  output.description.height = 0
  output.description.depth = 0
  output.description.mipCount = 0
  output.description.flags = 0
  output.description.usage = 0
  return output
}

private func createFFXFloatCoords(
  _ input: SIMD2<Float>
) -> FfxApiFloatCoords2D {
  var output = FfxApiFloatCoords2D()
  output.x = input[0]
  output.y = input[0]
  return output
}

private func createFFXDimensions(
  _ input: SIMD2<Int>
) -> FfxApiDimensions2D {
  var output = FfxApiDimensions2D()
  output.width = UInt32(input[0])
  output.height = UInt32(input[1])
  return output
}
#endif

extension Application {
  private func createJitterOffset() -> SIMD2<Float> {
    var jitterOffsetDesc = JitterOffsetDescriptor()
    jitterOffsetDesc.index = frameID
    jitterOffsetDesc.upscaleFactor = imageResources.renderTarget.upscaleFactor
    
    //return JitterOffset.create(descriptor: jitterOffsetDesc)
    return .zero
  }
  
  public func upscale(image: Image) -> Image {
    guard imageResources.renderTarget.upscaleFactor > 1 else {
      fatalError("Upscaling is not allowed.")
    }
    guard image.scaleFactor == 1 else {
      fatalError("Received image with incorrect scale factor.")
    }
    standardUpscale()
    
    var output = Image()
    output.scaleFactor = imageResources.renderTarget.upscaleFactor
    return output
  }

  private func standardUpscale() {
    guard let upscaler = imageResources.upscaler else {
      fatalError("Upscaler was not present.")
    }
    let colorTexture = imageResources.renderTarget
      .colorTextures[frameID % 2]
    let depthTexture = imageResources.renderTarget
      .depthTextures[frameID % 2]
    let motionTexture = imageResources.renderTarget
      .motionTextures[frameID % 2]
    let upscaledTexture = imageResources.renderTarget
      .upscaledTextures[frameID % 2]
    
    #if os(macOS)
    if frameID == 0 {
      upscaler.scaler.reset = true
    } else {
      upscaler.scaler.reset = false
    }
    
    upscaler.scaler.colorTexture = colorTexture
    upscaler.scaler.depthTexture = depthTexture
    upscaler.scaler.motionTexture = motionTexture
    upscaler.scaler.outputTexture = upscaledTexture
    
    let jitterOffset = createJitterOffset()
    upscaler.scaler.jitterOffsetX = -jitterOffset[0]
    upscaler.scaler.jitterOffsetY = -jitterOffset[1]
    
    device.commandQueue.withCommandList { commandList in
      commandList.mtlCommandEncoder.endEncoding()
      
      upscaler.scaler.encode(commandBuffer: commandList.mtlCommandBuffer)
      
      commandList.mtlCommandEncoder =
      commandList.mtlCommandBuffer.makeComputeCommandEncoder()!
      
      nonisolated(unsafe)
      let selfReference = self
      let inFlightFrameID = frameID % 3
      commandList.mtlCommandBuffer.addCompletedHandler { commandBuffer in
        selfReference.bvhBuilder.counters.queue.sync {
          var executionTime = commandBuffer.gpuEndTime
          executionTime -= commandBuffer.gpuStartTime
          let latencyMicroseconds = Int(executionTime * 1e6)
          selfReference.bvhBuilder.counters
            .upscaleLatencies[inFlightFrameID] = latencyMicroseconds
        }
      }
    }
    #else
    device.commandQueue.withCommandList { commandList in
      try! commandList.d3d12CommandList.EndQuery(
        bvhBuilder.counters.queryHeap,
        D3D12_QUERY_TYPE_TIMESTAMP,
        6)
      
      func createID3D12CommandList() -> UnsafeMutableRawPointer {
        let d3d12CommandList = commandList.d3d12CommandList
        let iid = SwiftCOM.ID3D12GraphicsCommandList.IID
        
        // Fetch the underlying pointer without worrying about memory leaks.
        let interface = try! d3d12CommandList.QueryInterface(iid: iid)
        _ = try! d3d12CommandList.Release()
        
        guard let interface else {
          fatalError("This should never happen.")
        }
        return interface
      }
      
      let dispatch = FFXDescriptor<ffxDispatchDescUpscale>()
      dispatch.type = FFX_API_DISPATCH_DESC_TYPE_UPSCALE
      dispatch.value.commandList = createID3D12CommandList()
      
      dispatch.value.color = createFFXResource(colorTexture)
      dispatch.value.depth = createFFXResource(depthTexture)
      dispatch.value.motionVectors = createFFXResource(motionTexture)
      dispatch.value.exposure = createEmptyFFXResource()
      dispatch.value.reactive = createEmptyFFXResource()
      dispatch.value.transparencyAndComposition = createEmptyFFXResource()
      dispatch.value.output = createFFXResource(upscaledTexture)
      
      // It takes some effort to investigate, but we are indeed getting better
      // results from jitterOffset * -1 than jitterOffset.
      let jitterOffset = createJitterOffset()
      let motionVectorScale = SIMD2<Float>(1, 1)
      dispatch.value.jitterOffset = createFFXFloatCoords(jitterOffset * -1)
      dispatch.value.motionVectorScale = createFFXFloatCoords(motionVectorScale)
      
      let upscaleFactor = imageResources.renderTarget.upscaleFactor
      let renderSize = display.frameBufferSize / Int(upscaleFactor)
      let upscaleSize = display.frameBufferSize
      dispatch.value.renderSize = createFFXDimensions(renderSize)
      dispatch.value.upscaleSize = createFFXDimensions(upscaleSize)
      
      // Sharpening harms the quality of bright shiny light reflections
      // (specular effect), making it look pixelated instead of smooth.
      dispatch.value.enableSharpening = false
      dispatch.value.sharpness = 0
      dispatch.value.frameTimeDelta = 2 // this doesn't do anything
      dispatch.value.preExposure = 1
      
      if frameID == 0 {
        dispatch.value.reset = true
      } else {
        dispatch.value.reset = false
      }
      
      dispatch.value.cameraNear = Float.greatestFiniteMagnitude
      dispatch.value.cameraFar = 0.075 // 75 pm, circumvents debug warning
      dispatch.value.cameraFovAngleVertical = camera.fovAngleVertical
      dispatch.value.viewSpaceToMetersFactor = 1
      dispatch.value.flags = 0
      
      // Encode the GPU commands for upscaling.
      upscaler.ffxContext.dispatch(descriptor: dispatch)
      
      try! commandList.d3d12CommandList.EndQuery(
        bvhBuilder.counters.queryHeap,
        D3D12_QUERY_TYPE_TIMESTAMP,
        7)
      
      let destinationBuffer = bvhBuilder.counters
        .queryDestinationBuffers[frameID % 3]
      try! commandList.d3d12CommandList.ResolveQueryData(
        bvhBuilder.counters.queryHeap,
        D3D12_QUERY_TYPE_TIMESTAMP,
        6,
        2,
        destinationBuffer.d3d12Resource,
        48)
    }
    #endif
  }

  // Fallback for debugging if the upscaler goes wrong, or for easily
  // visualizing the 3 inputs to the upscaler.
  private func fallbackUpscale() {
    device.commandQueue.withCommandList { commandList in
      // Bind the descriptor heap.
      #if os(Windows)
      commandList.setDescriptorHeap(descriptorHeap)
      #endif
      
      // Encode the compute command.
      commandList.withPipelineState(imageResources.upscaleShader) {
        // Bind the textures.
        #if os(macOS)
        let colorTexture = renderTarget.colorTextures[frameID % 2]
        let upscaledTexture = renderTarget.upscaledTextures[frameID % 2]
        commandList.mtlCommandEncoder
          .setTexture(colorTexture, index: 0)
        commandList.mtlCommandEncoder
          .setTexture(upscaledTexture, index: 1)
        #else
        commandList.setDescriptor(
          handleID: frameID % 2, index: 0)
        commandList.setDescriptor(
          handleID: 6 + frameID % 2, index: 1)
        #endif
        
        // Determine the dispatch grid size.
        func createGroupCount32() -> SIMD3<UInt32> {
          var groupCount = display.frameBufferSize
          
          let groupSize = SIMD2<Int>(8, 8)
          groupCount &+= groupSize &- 1
          groupCount /= groupSize
          
          return SIMD3<UInt32>(
            UInt32(groupCount[0]),
            UInt32(groupCount[1]),
            UInt32(1))
        }
        commandList.dispatch(groups: createGroupCount32())
      }
    }
  }
}

ImageResources.swift

struct ImageResourcesDescriptor {
  var device: Device?
  var display: Display?
  var memorySlotCount: Int?
  var upscaleFactor: Float?
  var worldDimension: Float?
}

class ImageResources {
  let renderShader: Shader
  let upscaleShader: Shader

  let renderTarget: RenderTarget
  let upscaler: Upscaler?
  
  var cameraArgsBuffer: RingBuffer
  var previousCameraArgs: CameraArgs?
  
  init(descriptor: ImageResourcesDescriptor) {
    guard let device = descriptor.device,
          let display = descriptor.display,
          let upscaleFactor = descriptor.upscaleFactor else {
      fatalError("Descriptor was incomplete.")
    }
    self.renderShader = Self.createRenderShader(descriptor: descriptor)
    self.upscaleShader = Self.createUpscaleShader(descriptor: descriptor)

    var renderTargetDesc = RenderTargetDescriptor()
    renderTargetDesc.device = device
    renderTargetDesc.display = display
    renderTargetDesc.upscaleFactor = upscaleFactor
    self.renderTarget = RenderTarget(descriptor: renderTargetDesc)
    
    if upscaleFactor > 1 {
      var upscalerDesc = UpscalerDescriptor()
      upscalerDesc.device = device
      upscalerDesc.display = display
      upscalerDesc.upscaleFactor = upscaleFactor
      self.upscaler = Upscaler(descriptor: upscalerDesc)
    } else {
      self.upscaler = nil
    }
    
    self.cameraArgsBuffer = Self.createCameraArgsBuffer(device: device)
    self.previousCameraArgs = nil
  }

  private static func createRenderShader(
    descriptor: ImageResourcesDescriptor
  ) -> Shader {
    guard let device = descriptor.device,
          let display = descriptor.display,
          let memorySlotCount = descriptor.memorySlotCount,
          let upscaleFactor = descriptor.upscaleFactor,
          let worldDimension = descriptor.worldDimension else {
      fatalError("Descriptor was incomplete.")
    }

    var renderShaderDesc = RenderShaderDescriptor()
    renderShaderDesc.isOffline = display.isOffline
    renderShaderDesc.memorySlotCount = memorySlotCount
    renderShaderDesc.supports16BitTypes = device.supports16BitTypes
    renderShaderDesc.upscaleFactor = upscaleFactor
    renderShaderDesc.worldDimension = worldDimension
    let renderShaderSource = RenderShader.createSource(
      descriptor: renderShaderDesc)
    
    var shaderDesc = ShaderDescriptor()
    shaderDesc.device = device
    shaderDesc.name = "render"
    #if os(macOS)
    shaderDesc.maxTotalThreadsPerThreadgroup = 1024
    #endif
    shaderDesc.threadsPerGroup = SIMD3(8, 8, 1)
    shaderDesc.source = renderShaderSource
    return Shader(descriptor: shaderDesc)
  }

  private static func createUpscaleShader(
    descriptor: ImageResourcesDescriptor
  ) -> Shader {
    guard let device = descriptor.device,
          let upscaleFactor = descriptor.upscaleFactor else {
      fatalError("Descriptor was incomplete.")
    }

    var shaderDesc = ShaderDescriptor()
    shaderDesc.device = device
    shaderDesc.name = "upscale"
    shaderDesc.threadsPerGroup = SIMD3(8, 8, 1)
    shaderDesc.source = UpscaleShader.createSource(
      upscaleFactor: upscaleFactor)
    return Shader(descriptor: shaderDesc)
  }
  
  private static func createCameraArgsBuffer(
    device: Device
  ) -> RingBuffer {
    var ringBufferDesc = RingBufferDescriptor()
    ringBufferDesc.accessLevel = .constant
    ringBufferDesc.device = device
    ringBufferDesc.size = MemoryLayout<CameraArgs>.stride * 2
    return RingBuffer(descriptor: ringBufferDesc)
  }
}

UpscaleShader.swift

// Fallback for debugging if the upscaler goes wrong, or for easily
// visualizing the 3 inputs to the upscaler.
struct UpscaleShader {
  static func createSource(upscaleFactor: Float) -> String {
    func importStandardLibrary() -> String {
      #if os(macOS)
      """
      #include <metal_stdlib>
      using namespace metal;
      """
      #else
      ""
      #endif
    }
    
    func functionSignature() -> String {
      #if os(macOS)
      """
      kernel void upscale(
        texture2d<float, access::read> colorTexture [[texture(0)]],
        texture2d<float, access::write> upscaledTexture [[texture(1)]],
        uint2 pixelCoords [[thread_position_in_grid]])
      """
      #else
      """
      RWTexture2D<float4> colorTexture : register(u0);
      RWTexture2D<float4> upscaledTexture : register(u1);
      
      [numthreads(8, 8, 1)]
      [RootSignature(
        "DescriptorTable(UAV(u0, numDescriptors = 1)),"
        "DescriptorTable(UAV(u1, numDescriptors = 1)),"
      )]
      void upscale(
        uint2 pixelCoords : SV_DispatchThreadID)
      """
      #endif
    }
    
    func readColor() -> String {
      #if os(macOS)
      "float4 color = colorTexture.read(inputCoords);"
      #else
      "float4 color = colorTexture[inputCoords];"
      #endif
    }
    
    func writeColor() -> String {
      #if os(macOS)
      "upscaledTexture.write(color, pixelCoords);"
      #else
      "upscaledTexture[pixelCoords] = color;"
      #endif
    }
    
    return """
    \(importStandardLibrary())
    
    // Utility for visualizing grayscale quantities
    // with more distinction between very close numbers.
    //
    // n = 0: red
    // n = 8: green
    // n = 4: blue
    float convertToChannel(
      float hue,
      float saturation,
      float lightness,
      uint n
    ) {
      float k = float(n) + hue / 30;
      k -= 12 * floor(k / 12);
      
      float a = saturation;
      a *= min(lightness, 1 - lightness);
      
      float output = min(k - 3, 9 - k);
      output = max(output, float(-1));
      output = min(output, float(1));
      output = lightness - a * output;
      return output;
    }
    
    \(functionSignature())
    {
      // Read from the input texture.
      uint2 inputCoords = pixelCoords / \(Int(upscaleFactor));
      \(readColor())
      
      // Write to the output texture.
      \(writeColor())
    }
    """
  }
}