From 25fbf1d8ce03473be09f9c0891d78ad671045217 Mon Sep 17 00:00:00 2001
From: bjorn <tie.372@gmail.com>
Date: Fri, 4 Jun 2021 22:11:21 -0600
Subject: [PATCH] love.graphics.dispatchThreadgroups;

---
 src/modules/graphics/Graphics.cpp        | 19 +++++++++++++++++++
 src/modules/graphics/Graphics.h          |  4 ++++
 src/modules/graphics/opengl/Graphics.cpp |  5 +++++
 src/modules/graphics/opengl/Graphics.h   |  2 ++
 src/modules/graphics/wrap_Graphics.cpp   | 12 ++++++++++++
 5 files changed, 42 insertions(+)

diff --git a/src/modules/graphics/Graphics.cpp b/src/modules/graphics/Graphics.cpp
index 2bb03affe..1cffbebb2 100644
--- a/src/modules/graphics/Graphics.cpp
+++ b/src/modules/graphics/Graphics.cpp
@@ -1066,6 +1066,25 @@ void Graphics::copyBuffer(Buffer *source, Buffer *dest, size_t sourceoffset, siz
 	source->copyTo(dest, sourceoffset, destoffset, size);
 }
 
+void Graphics::dispatchThreadgroups(Shader* shader, int x, int y, int z)
+{
+	if (!shader->hasStage(SHADERSTAGE_COMPUTE))
+		throw love::Exception("Only compute shaders can have threads dispatched.");
+
+	if (x <= 0 || y <= 0 || z <= 0)
+		throw love::Exception("Threadgroup dispatch size must be positive.");
+
+	if (x > capabilities.limits[LIMIT_THREADGROUPS_X]
+		|| y > capabilities.limits[LIMIT_THREADGROUPS_Y]
+		|| z > capabilities.limits[LIMIT_THREADGROUPS_Z])
+	{
+		throw love::Exception("Too many threadgroups dispatched.");
+	}
+
+	shader->attach();
+	dispatch(x, y, z);
+}
+
 Graphics::BatchedVertexData Graphics::requestBatchedDraw(const BatchedDrawCommand &cmd)
 {
 	BatchedDrawState &state = batchedDrawState;
diff --git a/src/modules/graphics/Graphics.h b/src/modules/graphics/Graphics.h
index e7f3d5cdb..57254289e 100644
--- a/src/modules/graphics/Graphics.h
+++ b/src/modules/graphics/Graphics.h
@@ -674,6 +674,8 @@ public:
 
 	void copyBuffer(Buffer *source, Buffer *dest, size_t sourceoffset, size_t destoffset, size_t size);
 
+	void dispatchThreadgroups(Shader* shader, int x, int y, int z);
+
 	void draw(Drawable *drawable, const Matrix4 &m);
 	void draw(Texture *texture, Quad *quad, const Matrix4 &m);
 	void drawLayer(Texture *texture, int layer, const Matrix4 &m);
@@ -935,6 +937,8 @@ protected:
 	virtual Shader *newShaderInternal(StrongRef<ShaderStage> stages[SHADERSTAGE_MAX_ENUM]) = 0;
 	virtual StreamBuffer *newStreamBuffer(BufferUsage type, size_t size) = 0;
 
+	virtual void dispatch(int x, int y, int z) = 0;
+
 	virtual void setRenderTargetsInternal(const RenderTargets &rts, int w, int h, int pixelw, int pixelh, bool hasSRGBtexture) = 0;
 
 	virtual void initCapabilities() = 0;
diff --git a/src/modules/graphics/opengl/Graphics.cpp b/src/modules/graphics/opengl/Graphics.cpp
index 2cd180626..385d4aec4 100644
--- a/src/modules/graphics/opengl/Graphics.cpp
+++ b/src/modules/graphics/opengl/Graphics.cpp
@@ -469,6 +469,11 @@ void Graphics::setActive(bool enable)
 	active = enable;
 }
 
+void Graphics::dispatch(int x, int y, int z)
+{
+	glDispatchCompute(x, y, z);
+}
+
 void Graphics::draw(const DrawCommand &cmd)
 {
 	gl.prepareDraw(this);
diff --git a/src/modules/graphics/opengl/Graphics.h b/src/modules/graphics/opengl/Graphics.h
index 9ec95d3ac..10ae808ec 100644
--- a/src/modules/graphics/opengl/Graphics.h
+++ b/src/modules/graphics/opengl/Graphics.h
@@ -68,6 +68,8 @@ public:
 
 	void setActive(bool active) override;
 
+	void dispatch(int x, int y, int z) override;
+
 	void draw(const DrawCommand &cmd) override;
 	void draw(const DrawIndexedCommand &cmd) override;
 	void drawQuads(int start, int count, const VertexAttributes &attributes, const BufferBindings &buffers, love::graphics::Texture *texture) override;
diff --git a/src/modules/graphics/wrap_Graphics.cpp b/src/modules/graphics/wrap_Graphics.cpp
index 169b7852c..1e12dc1d0 100644
--- a/src/modules/graphics/wrap_Graphics.cpp
+++ b/src/modules/graphics/wrap_Graphics.cpp
@@ -3255,6 +3255,16 @@ int w_polygon(lua_State *L)
 	return 0;
 }
 
+int w_dispatchThreadgroups(lua_State* L)
+{
+	Shader *shader = luax_checkshader(L, 1);
+	int x = (int) luaL_checkinteger(L, 2);
+	int y = (int) luaL_optinteger(L, 3, 1);
+	int z = (int) luaL_optinteger(L, 4, 1);
+	luax_catchexcept(L, [&](){ instance()->dispatchThreadgroups(shader, x, y, z); });
+	return 0;
+}
+
 int w_copyBuffer(lua_State *L)
 {
 	Buffer *source = luax_checkbuffer(L, 1);
@@ -3472,6 +3482,8 @@ static const luaL_Reg functions[] =
 	{ "print", w_print },
 	{ "printf", w_printf },
 
+	{ "dispatchThreadgroups", w_dispatchThreadgroups },
+
 	{ "copyBuffer", w_copyBuffer },
 
 	{ "isCreated", w_isCreated },