Compute Shader + Render Feature Heatmap

HeatmapCompute.compute

#pragma kernel CSMain

RWTexture2D<float> heatmapTexture;
float2 texSize;

StructuredBuffer<float2> enemyPositions;
int enemyCount;

[numthreads(8, 8, 1)]
void CSMain(uint3 id : SV_DispatchThreadID)
{
    int2 pixel = int2(id.xy);
    float2 uv = pixel;
    float heat = 0;

    for (int i = 0; i < enemyCount; i++) {
        float2 enemyPos = enemyPositions[i];
        float dist = distance(uv, enemyPos);
        float radius = 20.0;
        heat += saturate(1.0 - dist / radius); // linear falloff
    }

    heat = saturate(heat);
    heatmapTexture[pixel] = heat;
}

HeatmapRendererFeature.cs

using UnityEngine;
using UnityEngine.Rendering;
using UnityEngine.Rendering.RenderGraphModule;
using UnityEngine.Rendering.Universal;
using UnityEngine.Experimental.Rendering;

public class HeatmapRendererFeature : ScriptableRendererFeature {
    public static HeatmapRendererFeature Instance { get; private set; }

    class HeatmapPass : ScriptableRenderPass {
        ComputeShader computeShader;
        int kernel;

        GraphicsBuffer enemyBuffer;
        Vector2[] enemyPositions;
        int enemyCount = 64;

        RTHandle heatmapHandle;
        int width = 256, height = 256;

        public RTHandle Heatmap => heatmapHandle;

        public void Setup(ComputeShader cs) {
            computeShader = cs;
            kernel = cs.FindKernel("CSMain");

            if (heatmapHandle == null || heatmapHandle.rt.width != width || heatmapHandle.rt.height != height) {
                heatmapHandle?.Release();
                var desc = new RenderTextureDescriptor(width, height, GraphicsFormat.R32_SFloat, 0) {
                    enableRandomWrite = true,
                    msaaSamples = 1,
                    sRGB = false,
                    useMipMap = false
                };
                heatmapHandle = RTHandles.Alloc(desc, name: "_HeatmapRT");
            }

            if (enemyBuffer == null || enemyBuffer.count != enemyCount) {
                enemyBuffer?.Release();
                enemyBuffer = new GraphicsBuffer(GraphicsBuffer.Target.Structured, enemyCount, sizeof(float) * 2);
                enemyPositions = new Vector2[enemyCount];
            }
        }

        class PassData {
            public ComputeShader compute;
            public int kernel;
            public TextureHandle output;
            public Vector2 texSize;
            public BufferHandle enemyHandle;
            public int enemyCount;
        }

        public override void RecordRenderGraph(RenderGraph graph, ContextContainer context) {
            for (int i = 0; i < enemyCount; i++) {
                float t = Time.time * 0.5f + i * 0.1f;
                float x = Mathf.PerlinNoise(t, i * 1.31f) * width;
                float y = Mathf.PerlinNoise(i * 0.91f, t) * height;
                enemyPositions[i] = new Vector2(x, y);
            }
            
            enemyBuffer.SetData(enemyPositions);

            TextureHandle texHandle = graph.ImportTexture(heatmapHandle);
            BufferHandle enemyHandle = graph.ImportBuffer(enemyBuffer);

            using IComputeRenderGraphBuilder builder = graph.AddComputePass("HeatmapPass", out PassData data);
            data.compute = computeShader;
            data.kernel = kernel;
            data.output = texHandle;
            data.enemyHandle = enemyHandle;
            data.enemyCount = enemyCount;

            builder.UseTexture(texHandle, AccessFlags.Write);
            builder.UseBuffer(enemyHandle, AccessFlags.Read);

            builder.SetRenderFunc((PassData d, ComputeGraphContext ctx) => {
                ctx.cmd.SetComputeIntParam(d.compute, "enemyCount", d.enemyCount);
                ctx.cmd.SetComputeBufferParam(d.compute, d.kernel, "enemyPositions", d.enemyHandle);
                ctx.cmd.SetComputeTextureParam(d.compute, d.kernel, "heatmapTexture", d.output);
                ctx.cmd.DispatchCompute(d.compute, d.kernel, Mathf.CeilToInt(width / 8f), Mathf.CeilToInt(height / 8f), 1);
            });
        }


        public void Cleanup() {
            heatmapHandle?.Release();
            heatmapHandle = null;

            enemyBuffer?.Release();
            enemyBuffer = null;
        }
    }

    [SerializeField] ComputeShader computeShader;
    HeatmapPass pass;

    public override void Create() {
        pass = new HeatmapPass {
            renderPassEvent = RenderPassEvent.BeforeRendering
        };
        Instance = this;
    }

    public override void AddRenderPasses(ScriptableRenderer renderer, ref RenderingData renderingData) {
        if (!SystemInfo.supportsComputeShaders || computeShader == null)
            return;

        pass.Setup(computeShader);
        renderer.EnqueuePass(pass);
    }

    protected override void Dispose(bool disposing) {
        pass?.Cleanup();
    }

    public RTHandle GetHeatmapTexture() => pass?.Heatmap;
}

HeatmapVisualizer.cs

using UnityEngine;
using UnityEngine.UI;

public class HeatmapVisualizer : MonoBehaviour {
    public Material material;
    Image heatmapImage;

    void Start() {
        heatmapImage = GetComponent<Image>();
    }

    void Update() {
        var feature = HeatmapRendererFeature.Instance;
        if (feature == null) return;

        var texture = feature.GetHeatmapTexture();
        if (texture != null) {
            material.SetTexture("_MainTex", texture);
        }

        if (heatmapImage && texture != null) {
            Texture2D texture2D = new Texture2D(texture.rt.width, texture.rt.height, TextureFormat.RFloat, false);

            RenderTexture.active = texture;
            texture2D.ReadPixels(new Rect(0, 0, texture.rt.width, texture.rt.height), 0, 0);
            texture2D.Apply();
            heatmapImage.sprite = Sprite.Create(texture2D, new Rect(0, 0, texture.rt.width, texture.rt.height), new Vector2(0.5f, 0.5f));
        }
    }
}