RealityView with LowLevelTexture and Gaussian Blur (separate compute shader passes)

diagonalLineTexture.metal

#include <metal_stdlib>
using namespace metal;

kernel void diagonalLineTexture(texture2d<float, access::write> outputTexture [[texture(0)]],
                                uint2 gid [[thread_position_in_grid]]) {
    int thickness = 50;
    int jaggedness = 1;

    // Introduce jaggedness using a modulo function
    int offset = (int(gid.y) / jaggedness) % (2 * jaggedness) - jaggedness;

    // Calculate the jagged diagonal condition
    bool isJaggedDiagonal = (abs(int(gid.x) - int(gid.y) + offset) < thickness);

    // Write to the texture
    if (isJaggedDiagonal) {
        outputTexture.write(float4(0.0, 0.0, 0.0, 1.0), gid); // Black pixel
    } else {
        outputTexture.write(float4(1.0, 1.0, 1.0, 1.0), gid); // White background
    }
}

GaussianBlurParams.h

// Don't forget to include this bridging header

#ifndef GaussianBlurParams_h
#define GaussianBlurParams_h

struct GaussianBlurParams {
    int32_t kernelSize;
    float intensity;
};

#endif /* GaussianBlurParams_h */

gaussianBlurTextureProcessing.metal

#include <metal_stdlib>
using namespace metal;

#include "GaussianBlurParams.h"

kernel void gaussianBlurTextureProcessing(texture2d<half, access::sample> inputTexture [[texture(0)]],
                                          texture2d<half, access::write> outputTexture [[texture(1)]],
                                          constant GaussianBlurParams& params [[buffer(0)]],
                                          uint2 gid [[thread_position_in_grid]])
{
    // Check bounds
    if (gid.x >= outputTexture.get_width() || gid.y >= outputTexture.get_height()) {
        return;
    }
    
    constexpr sampler textureSampler(filter::linear, address::clamp_to_edge);
    float2 textureSize = float2(inputTexture.get_width(), inputTexture.get_height());
    float2 pixelSize = 1.0 / textureSize;
    float2 uv = float2(gid) / textureSize;
    
    half4 sum = half4(0.0);
    float totalWeight = 0.0;

    for(int y = -params.kernelSize; y <= params.kernelSize; y++) {
        for(int x = -params.kernelSize; x <= params.kernelSize; x++) {
            float2 offset = float2(x, y) * pixelSize;
            float weight = exp(-(x*x + y*y) / (2.0 * params.intensity));
            sum += inputTexture.sample(textureSampler, uv + offset) * weight;
            totalWeight += weight;
        }
    }
    
    outputTexture.write(sum / totalWeight, gid);
}

GaussianBlurTextureView.swift

import SwiftUI
import RealityKit

struct GaussianBlurTextureView: View {
    @State var isBlurred: Bool = false
    @State var texture: LowLevelTexture?
    let commandQueue: MTLCommandQueue
    let texturePipeline: MTLComputePipelineState
    let gaussianBlurPipeline: MTLComputePipelineState
    var intermediateTexture: MTLTexture!

    init() {
        let device = MTLCreateSystemDefaultDevice()!
        self.commandQueue = device.makeCommandQueue()!
        
        let library = device.makeDefaultLibrary()!
        
        let updateFunction = library.makeFunction(name: "diagonalLineTexture")!
        self.texturePipeline = try! device.makeComputePipelineState(function: updateFunction)
        let antiAliasFunction = library.makeFunction(name: "gaussianBlurTextureProcessing")!
        gaussianBlurPipeline = try! device.makeComputePipelineState(function: antiAliasFunction)
        
        // Create intermediate texture
        let intermediateDesc = MTLTextureDescriptor()
        intermediateDesc.width = textureDescriptor.width
        intermediateDesc.height = textureDescriptor.height
        intermediateDesc.pixelFormat = .rgba16Float
        intermediateDesc.usage = [.shaderRead, .shaderWrite]
        intermediateTexture = device.makeTexture(descriptor: intermediateDesc)!
    }
    
    var body: some View {
        RealityView { content in
            let entity = try! getEntity()
            content.add(entity)
            updateTexture()
        }
        .gesture(
           SpatialTapGesture()
               .targetedToAnyEntity()
               .onEnded { value in
                   isBlurred.toggle()
                   updateTexture()
               }
       )
    }
    
    func getEntity() throws -> Entity {
        let width: Float = 1.0
        let height: Float = 1.0
        let mesh = MeshResource.generatePlane(width: width, height: height)
        let texture = try LowLevelTexture(descriptor: textureDescriptor)
        let resource = try TextureResource(from: texture)
        var material = UnlitMaterial()
        material.color.texture = .init(resource)
        let modelComponent = ModelComponent(mesh: mesh, materials: [material])
        let entity = ModelEntity()
        entity.components.set([
            InputTargetComponent(),
            CollisionComponent(shapes: [.generateBox(width: width, height: height, depth: 1)]),
        ])
        
        entity.components.set(modelComponent)
        entity.scale *= 5.0

        self.texture = texture
        return entity
    }
    
    func updateTexture() {
        guard let texture else { return }
        guard let commandBuffer = commandQueue.makeCommandBuffer(),
              let computeEncoder = commandBuffer.makeComputeCommandEncoder()
        else { return }

        commandBuffer.enqueue()
        
        let targetTexture = texture.replace(using: commandBuffer)

        computeEncoder.setComputePipelineState(texturePipeline)

        computeEncoder.setTexture(isBlurred ? intermediateTexture : targetTexture, index: 0)

        let w = texturePipeline.threadExecutionWidth
        let h = texturePipeline.maxTotalThreadsPerThreadgroup / w
        let threadGroupSize = MTLSizeMake(w, h, 1)
        let threadGroupCount = MTLSizeMake(
            (textureDescriptor.width + threadGroupSize.width - 1) / threadGroupSize.width,
            (textureDescriptor.height + threadGroupSize.height - 1) / threadGroupSize.height,
            1)

        computeEncoder.dispatchThreadgroups(threadGroupCount, threadsPerThreadgroup: threadGroupSize)

        if isBlurred {
            applyGaussianBlur(computeEncoder: computeEncoder,
                              inputTexture: intermediateTexture,
                              outputTexture: targetTexture,
                              pipelineState: gaussianBlurPipeline
            )
        }

        computeEncoder.endEncoding()
        commandBuffer.commit()
    }
    
    func applyGaussianBlur(computeEncoder: MTLComputeCommandEncoder,
                           inputTexture: MTLTexture,
                           outputTexture: MTLTexture,
                           pipelineState: MTLComputePipelineState) {
        computeEncoder.setComputePipelineState(pipelineState)
        computeEncoder.setTexture(inputTexture, index: 0)
        computeEncoder.setTexture(outputTexture, index: 1)

        var params = GaussianBlurParams(kernelSize: Int32(3),
                                        intensity: 1.0)
        computeEncoder.setBytes(&params,
                               length: MemoryLayout<GaussianBlurParams>.size,
                               index: 0)
        
        let w = pipelineState.threadExecutionWidth
        let h = pipelineState.maxTotalThreadsPerThreadgroup / w
        let threadGroupSize = MTLSizeMake(w, h, 1)
        
        let threadGroupCount = MTLSizeMake(
            (outputTexture.width + threadGroupSize.width - 1) / threadGroupSize.width,
            (outputTexture.height + threadGroupSize.height - 1) / threadGroupSize.height,
            1)
        
        computeEncoder.dispatchThreadgroups(threadGroupCount,
                                          threadsPerThreadgroup: threadGroupSize)
    }
    
    var textureDescriptor: LowLevelTexture.Descriptor {
        var desc = LowLevelTexture.Descriptor()

        desc.textureType = .type2D
        desc.arrayLength = 1

        desc.width = 2048
        desc.height = 2048
        desc.depth = 1

        desc.mipmapLevelCount = 1
        desc.pixelFormat = .bgra8Unorm
        desc.textureUsage = [.shaderRead, .shaderWrite]
        desc.swizzle = .init(red: .red, green: .green, blue: .blue, alpha: .alpha)

        return desc
    }
}

#Preview {
    GaussianBlurTextureView()
}