Skip to content

Instantly share code, notes, and snippets.

@amygoodchild
Created January 31, 2020 13:37
Show Gist options
  • Save amygoodchild/70b58f527d14bc061413d047aecad5f9 to your computer and use it in GitHub Desktop.
Save amygoodchild/70b58f527d14bc061413d047aecad5f9 to your computer and use it in GitHub Desktop.
// Reset
#pragma kernel ResetTextureKernel
#pragma kernel ResetAgentsKernel
// Step
#pragma kernel MoveAgentsKernel
#pragma kernel WriteTrailsKernel
// Render
#pragma kernel AgentsDebugKernel
#pragma kernel RenderKernel
Texture2D<float4> readTex;
SamplerState sampler_readTex;
RWTexture2D<float4> writeTex;
RWTexture2D<float4> outTex;
struct Agent
{
float2 position;
float2 direction;
};
RWStructuredBuffer<Agent> agentsBuffer;
uint rez;
uint stepn;
uint time;
/*
*
*
*
* UTIL
*
*
*
*/
// via "Art of Code" on youtube
float2 Random(float2 p)
{
float3 a = frac(p.xyx * float3(123.34, 234.34, 345.65));
a+= dot(a, a + 34.45);
return frac(float2(a.x * a.y, a.y * a.z));
}
float2 RandomDirection(float2 p)
{
return(normalize(2.0 * (Random(p) - 0.5)));
}
/*
*
*
*
* RESET
*
*
*
*/
[numthreads(1, 1, 1)]
void ResetTextureKernel(uint3 id : SV_DispatchThreadID)
{
writeTex[id.xy] = 0;
}
[numthreads(1, 1, 1)]
void ResetAgentsKernel(uint3 id : SV_DispatchThreadID)
{
Agent a;
a.position = Random(id.x * .0001 + time * .001) * rez;
a.direction = RandomDirection(id.xx * .01 + sin(time));
agentsBuffer[id.x] = a;
}
/*
*
*
*
* STEP
*
*
*
*/
[numthreads(1, 1, 1)]
void WriteTrailsKernel(uint3 id : SV_DispatchThreadID)
{
Agent a = agentsBuffer[id.x];
writeTex[round(a.position)] = clamp(writeTex[round(a.position)] + .1, 0, 1);
}
float2 SimpleTurns(uint3 id, Agent a) {
//We will use round whenever reading from or writing to the grid tex
float4 f = readTex[round(a.position + a.direction * 2)];
float2 d = a.direction;
if (f.x > 0){
d = RandomDirection(id.xx * 0.1 + sin(time));
}
return d;
}
[numthreads(1, 1, 1)]
void MoveAgentsKernel(uint3 id : SV_DispatchThreadID)
{
Agent a = agentsBuffer[id.x];
// Choose next direction
a.direction = SimpleTurns(id, a);
// Move Forward
a.position = a.position + a.direction;
// Boundaries: Wrap
if (a.position.x < 0){
a.position.x - rez - 1;
}
if (a.position.y < 0){
a.position.y - rez - 1;
}
a.position %= float2(rez, rez);
agentsBuffer[id.x] = a;
}
/*
*
*
*
* RENDER
*
*
*
*/
[numthreads(1, 1, 1)]
void RenderKernel(uint3 id : SV_DispatchThreadID)
{
outTex[id.xy] = readTex[id.xy];
}
[numthreads(1, 1, 1)]
void AgentsDebugKernel(uint3 id : SV_DispatchThreadID)
{
Agent a = agentsBuffer[id.x];
outTex[round(a.position)] += float4(0, .1, 0, 0);
}
using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using EasyButtons;
[ExecuteAlways]
public class TrailAgent : MonoBehaviour
{
[Header("Trail Agent Params")]
[Range(1, 50000)]
public int agentsCount = 1;
private ComputeBuffer agentsBuffer;
[Header("Setup")]
[Range(8, 2048)]
public int rez = 8;
[Range(0,50)]
public int stepsPerFrame = 0;
[Range(1, 50)]
public int stepMod = 1;
public Material outMat;
public ComputeShader cs;
private RenderTexture readTex;
private RenderTexture writeTex;
private RenderTexture outTex;
private int agentsDebugKernel;
private int moveAgentsKernel;
private int writeTrailsKernel;
private int renderKernel;
protected List<ComputeBuffer> buffers;
protected List<RenderTexture> textures;
protected int stepn = -1;
/*
*
*
* RESET
*
*
*/
void Start()
{
Reset();
}
[Button]
public void Reset()
{
Release();
agentsDebugKernel = cs.FindKernel("AgentsDebugKernel");
moveAgentsKernel = cs.FindKernel("MoveAgentsKernel");
renderKernel = cs.FindKernel("RenderKernel");
writeTrailsKernel = cs.FindKernel("WriteTrailsKernel");
readTex = CreateTexture(rez, FilterMode.Point);
writeTex = CreateTexture(rez, FilterMode.Point);
outTex = CreateTexture(rez, FilterMode.Point);
agentsBuffer = new ComputeBuffer(agentsCount, sizeof(float) * 4);
buffers.Add(agentsBuffer);
GPUResetKernel();
Render();
}
private void GPUResetKernel()
{
int kernel;
cs.SetInt("rez", rez);
cs.SetInt("time", Time.frameCount);
kernel = cs.FindKernel("ResetTextureKernel");
cs.SetTexture(kernel, "writeTex", writeTex);
cs.Dispatch(kernel, rez, rez, 1);
cs.SetTexture(kernel, "writeTex", readTex);
cs.Dispatch(kernel, rez, rez, 1);
kernel = cs.FindKernel("ResetAgentsKernel");
cs.SetBuffer(kernel, "agentsBuffer", agentsBuffer);
cs.Dispatch(kernel, agentsCount, 1, 1);
}
/*
*
*
* STEP
*
*
*/
void Update()
{
if (Time.frameCount % stepMod == 0)
{
for (int i = 0; i < stepsPerFrame; i++)
{
Step();
}
}
}
[Button]
public void Step()
{
stepn+=1;
cs.SetInt("time", Time.frameCount);
GPUMoveAgentsKernel();
GPUWriteTrailsKernel();
Render();
}
private void GPUMoveAgentsKernel()
{
cs.SetBuffer(moveAgentsKernel, "agentsBuffer", agentsBuffer);
cs.SetTexture(moveAgentsKernel, "readTex", readTex);
cs.Dispatch(moveAgentsKernel, agentsCount, 1, 1);
}
private void GPUWriteTrailsKernel()
{
cs.SetBuffer(writeTrailsKernel, "agentsBuffer", agentsBuffer);
// Note: We will have to change this later:
cs.SetTexture(writeTrailsKernel, "writeTex", readTex);
cs.Dispatch(writeTrailsKernel, agentsCount, 1, 1);
}
/*
*
*
* RENDER
*
*
*/
private void Render()
{
GPURenderKernel();
GPUAgentsDebugKernel();
outMat.SetTexture("_UnlitColorMap", outTex);
if (!Application.isPlaying)
{
UnityEditor.SceneView.RepaintAll();
}
}
private void GPURenderKernel()
{
cs.SetTexture(renderKernel, "readTex", readTex);
cs.SetTexture(renderKernel, "outTex", outTex);
cs.Dispatch(renderKernel, rez, rez, 1);
}
private void GPUAgentsDebugKernel()
{
cs.SetBuffer(agentsDebugKernel, "agentsBuffer", agentsBuffer);
cs.SetTexture(agentsDebugKernel, "outTex", outTex);
cs.Dispatch(agentsDebugKernel, agentsCount, 1, 1);
}
/*
*
*
* RENDER
*
*
*/
public void Release()
{
if (buffers != null)
{
foreach (ComputeBuffer buffer in buffers)
{
if (buffer != null)
{
buffer.Release();
}
}
}
buffers = new List<ComputeBuffer>();
if (textures != null)
{
foreach (RenderTexture tex in textures)
{
if (tex != null)
{
tex.Release();
}
}
}
textures = new List<RenderTexture>();
}
private void onDestroy()
{
Release();
}
private void onEnable()
{
Release();
}
private void onDisable()
{
Release();
}
protected RenderTexture CreateTexture(int r, FilterMode filterMode)
{
RenderTexture texture = new RenderTexture(r, r, 1, RenderTextureFormat.ARGBFloat);
texture.name = "out";
texture.enableRandomWrite = true;
texture.dimension = UnityEngine.Rendering.TextureDimension.Tex2D;
texture.volumeDepth = 1;
texture.filterMode = filterMode;
texture.wrapMode = TextureWrapMode.Repeat;
texture.autoGenerateMips = false;
texture.useMipMap = false;
texture.Create();
textures.Add(texture);
return texture;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment