SDFMapCreator/SDFMapCreator/SdfKernels.cs

using System.Diagnostics;
using System.Numerics;
using ILGPU;
using ILGPU.Runtime;
using ILGPU.Runtime.CPU;
using ILGPU.Runtime.Cuda;
using ILGPU.Runtime.OpenCL;

namespace SDFMapCreator;

public partial class SdfKernels {
    Context gpuContext;
    Accelerator accelerator;
    public SdfKernels() {
        // Initialize the GPU context
        gpuContext = Context.Create(builder => builder
            .Cuda()
            .OpenCL()
            .CPU()
            .Math(MathMode.Fast32BitOnly)
            .EnableAlgorithms()
        );

        Console.WriteLine("Reading available accelerators (CUDA only)...");
        foreach (var device in gpuContext.GetCudaDevices()) {
            accelerator = device.CreateAccelerator(gpuContext);
            Console.WriteLine($"Found accelerator: {accelerator.Name} ({accelerator.AcceleratorType})");
        }

        Console.WriteLine("Reading available accelerators (OpenCL only)...");
        foreach (var device in gpuContext.GetCLDevices()) {
            accelerator = device.CreateAccelerator(gpuContext);
            Console.WriteLine($"Found accelerator: {accelerator.Name} ({accelerator.AcceleratorType})");
        }

        Console.WriteLine("Reading available accelerators (CPU only)...");
        foreach (var device in gpuContext.GetCPUDevices()) {
            accelerator = device.CreateAccelerator(gpuContext);
            Console.WriteLine($"Found accelerator: {accelerator.Name} ({accelerator.AcceleratorType})");
        }

        accelerator = gpuContext.GetPreferredDevice(false).CreateAccelerator(gpuContext);
        Console.WriteLine($"Using accelerator: {accelerator.Name} ({accelerator.AcceleratorType})");
    }

    public SdfKernels(DeviceType device) {
        Console.WriteLine($"Looking up for {device.ToString()} accelerators...");
        var builder = Context.Create();
        switch (device) {
            case DeviceType.CPU: builder.CPU(); break;
            case DeviceType.CUDA: builder.Cuda(); break;
            case DeviceType.OpenCL: builder.OpenCL(); break;
        }
        
        gpuContext = builder
            .Math(MathMode.Fast32BitOnly)
            .EnableAlgorithms()
            .ToContext();
        
        accelerator = gpuContext.GetPreferredDevice(false).CreateAccelerator(gpuContext);
        Console.WriteLine($"Using accelerator: {accelerator.Name} ({accelerator.AcceleratorType})");
    }

    public void SelfMask(Vector3[,] input, out Vector3[,] mask) {
        var width = input.GetLength(0);
        var height = input.GetLength(1);
        mask = new Vector3[width, height];

        using var buffer = accelerator.Allocate2DDenseX<Vector3>(new(width, height));
        buffer.CopyFromCPU(input);

        using var maskBuffer = accelerator.Allocate2DDenseX<Vector3>(new(width, height));
        
        var selfMaskKernel = accelerator.LoadAutoGroupedStreamKernel<Index2D, ArrayView2D<Vector3, Stride2D.DenseX>,
            ArrayView2D<Vector3, Stride2D.DenseX>>(SelfMaskKernel);

        selfMaskKernel(new(width, height), buffer.View, maskBuffer.View);

        accelerator.Synchronize();

        mask = maskBuffer.GetAsArray2D();
    }

    public void ABMask(Vector3[,] A, Vector3[,] B, out Vector3[,] mask) {
        var width = A.GetLength(0);
        var height = A.GetLength(1);
        mask = new Vector3[width, height];

        using var bufferA = accelerator.Allocate2DDenseX<Vector3>(new(width, height));
        bufferA.CopyFromCPU(A);

        using var bufferB = accelerator.Allocate2DDenseX<Vector3>(new(width, height));
        bufferB.CopyFromCPU(B);

        using var maskBuffer = accelerator.Allocate2DDenseX<Vector3>(new(width, height));

        var abMaskKernel = accelerator.LoadAutoGroupedStreamKernel<Index2D,
            ArrayView2D<Vector3, Stride2D.DenseX>,
            ArrayView2D<Vector3, Stride2D.DenseX>,
            ArrayView2D<Vector3, Stride2D.DenseX>>(ABMaskKernel);

        abMaskKernel(new(width, height), bufferA.View, bufferB.View, maskBuffer.View);

        accelerator.Synchronize();

        mask = maskBuffer.GetAsArray2D();
    }

    public void EdgeDetect(Vector3[,] mask, out Vector3[,] edge) {
        var width = mask.GetLength(0);
        var height = mask.GetLength(1);
        edge = new Vector3[width, height];

        using var buffer = accelerator.Allocate2DDenseX<Vector3>(new(width, height));
        buffer.CopyFromCPU(mask);

        var edgeKernel = accelerator.LoadAutoGroupedStreamKernel<Index2D,
            ArrayView2D<Vector3, Stride2D.DenseX>, uint, uint>(EdgeKernel);

        edgeKernel(new(width, height), buffer.View, (uint)width, (uint)height);

        accelerator.Synchronize();

        edge = buffer.GetAsArray2D();
    }

    public void Sdf(Vector2[] edges, int width, int height, out Vector3[,] sdf) {
        using var buffer = accelerator.Allocate2DDenseX<Vector3>(new(width, height));

        using var edgeBuffer = accelerator.Allocate1D<Vector2>(edges.Length);
        edgeBuffer.CopyFromCPU(edges);

        var sdfKernel = accelerator.LoadAutoGroupedStreamKernel<Index2D,
            ArrayView2D<Vector3, Stride2D.DenseX>,
            ArrayView1D<Vector2, Stride1D.Dense>,
            int, int>(SdfKernel);

        sdfKernel(new(width, height), buffer.View, edgeBuffer.View, width, height);

        accelerator.Synchronize();

        sdf = buffer.GetAsArray2D();
    }

    public void Gradient(Vector3[,] mask, Vector3[,] sdfa, Vector3[,] sdfb, out Vector3[,] gradient) {
        var width = mask.GetLength(0);
        var height = mask.GetLength(1);
        gradient = new Vector3[width, height];

        using var bufferMask = accelerator.Allocate2DDenseX<Vector3>(new(width, height));
        bufferMask.CopyFromCPU(mask);

        using var bufferSdfa = accelerator.Allocate2DDenseX<Vector3>(new(width, height));
        bufferSdfa.CopyFromCPU(sdfa);

        using var bufferSdfb = accelerator.Allocate2DDenseX<Vector3>(new(width, height));
        bufferSdfb.CopyFromCPU(sdfb);

        using var bufferGradient = accelerator.Allocate2DDenseX<Vector3>(new(width, height));

        var gradientKernel = accelerator.LoadAutoGroupedStreamKernel<Index2D,
            ArrayView2D<Vector3, Stride2D.DenseX>,
            ArrayView2D<Vector3, Stride2D.DenseX>,
            ArrayView2D<Vector3, Stride2D.DenseX>,
            ArrayView2D<Vector3, Stride2D.DenseX>>(GradientKernel);

        gradientKernel(new(width, height), bufferMask.View, bufferSdfa.View, bufferSdfb.View, bufferGradient.View);

        accelerator.Synchronize();

        gradient = bufferGradient.GetAsArray2D();
    }

    public void DirectionalBlur(Vector3[,] image, Vector3[,] mask, out Vector3[,] output, float radius = 3f,
        float step = .5f, float sigma = 1f) {
        var width = image.GetLength(0);
        var height = image.GetLength(1);
        output = new Vector3[width, height];

        using var imageBuffer = accelerator.Allocate2DDenseX<Vector3>(new(width, height));
        imageBuffer.CopyFromCPU(image);

        using var maskBuffer = accelerator.Allocate2DDenseX<Vector3>(new(width, height));
        maskBuffer.CopyFromCPU(mask);

        using var outputBuffer = accelerator.Allocate2DDenseX<Vector3>(new(width, height));

        var blurKernel = accelerator.LoadAutoGroupedStreamKernel<Index2D,
            ArrayView2D<Vector3, Stride2D.DenseX>,
            ArrayView2D<Vector3, Stride2D.DenseX>,
            ArrayView2D<Vector3, Stride2D.DenseX>,
            float, float, float, int, int>(DirectionalBlurKernel);

        blurKernel(new(width, height), imageBuffer.View, maskBuffer.View, outputBuffer.View,
            radius, step, sigma, width, height);

        accelerator.Synchronize();

        output = outputBuffer.GetAsArray2D();
    }

    static string GetInfoString(Accelerator a) {
        var infoString = new StringWriter();
        a.PrintInformation(infoString);
        return infoString.ToString();
    }
    
    ~SdfKernels() {
        accelerator?.Dispose();
        gpuContext?.Dispose();
    }
}