[Fix] (MG_State/GLState, MG_Util/ShaderTranspiler): validate compute local size [skip ci]

MobileGL/MG_State/GLState/ProgramState/ShaderObject.cpp

@@ -12,6 +12,132 @@
 #include <MG_Util/Converters/MGToGL/ProgramEnumConverter.h>
 #include <MG_Util/ShaderTranspiler/ShaderSourceProcessor.h>
 #include <MG_Util/ShaderTranspiler/glslang/UniformTraverser.h>
+#include <MG_Backend/BackendObjects.h>
+
+namespace {
+    struct ComputeLocalSize {
+        MobileGL::Uint x = 1;
+        MobileGL::Uint y = 1;
+        MobileGL::Uint z = 1;
+        bool declared = false;
+    };
+
+    static MobileGL::String StripGlslComments(const MobileGL::String& source) {
+        MobileGL::String result;
+        result.reserve(source.length());
+
+        bool inLineComment = false;
+        bool inBlockComment = false;
+        for (MobileGL::SizeT i = 0; i < source.length(); ++i) {
+            if (inLineComment) {
+                if (source[i] == '\n') {
+                    inLineComment = false;
+                    result.push_back(source[i]);
+                } else {
+                    result.push_back(' ');
+                }
+                continue;
+            }
+
+            if (inBlockComment) {
+                if (source[i] == '*' && i + 1 < source.length() && source[i + 1] == '/') {
+                    inBlockComment = false;
+                    result.append("  ");
+                    ++i;
+                } else {
+                    result.push_back(source[i] == '\n' ? '\n' : ' ');
+                }
+                continue;
+            }
+
+            if (source[i] == '/' && i + 1 < source.length()) {
+                if (source[i + 1] == '/') {
+                    inLineComment = true;
+                    result.append("  ");
+                    ++i;
+                    continue;
+                }
+                if (source[i + 1] == '*') {
+                    inBlockComment = true;
+                    result.append("  ");
+                    ++i;
+                    continue;
+                }
+            }
+
+            result.push_back(source[i]);
+        }
+
+        return result;
+    }
+
+    static ComputeLocalSize ParseComputeLocalSize(const MobileGL::String& source) {
+        ComputeLocalSize localSize;
+        const MobileGL::String uncommentedSource = StripGlslComments(source);
+        const std::regex localSizePattern(R"(local_size_([xyz])\s*=\s*([0-9]+))");
+
+        for (std::sregex_iterator it(uncommentedSource.begin(), uncommentedSource.end(), localSizePattern), end;
+             it != end; ++it) {
+            const char axis = (*it)[1].str()[0];
+            const auto value = static_cast<unsigned long long>(std::stoull((*it)[2].str()));
+            const MobileGL::Uint clampedValue = value > UINT_MAX ? UINT_MAX : static_cast<MobileGL::Uint>(value);
+
+            // TODO: Replace this literal layout scanner with parser/AST-backed validation so expressions and
+            // specialization-id layouts are handled consistently with glslang.
+            localSize.declared = true;
+            if (axis == 'x') {
+                localSize.x = clampedValue;
+            } else if (axis == 'y') {
+                localSize.y = clampedValue;
+            } else {
+                localSize.z = clampedValue;
+            }
+        }
+
+        return localSize;
+    }
+
+    static MobileGL::Uint GetComputeWorkGroupSizeLimit(MobileGL::Uint index) {
+        constexpr MobileGL::Uint kFrontendMinComputeWorkGroupSizes[] = {1024, 1024, 64};
+        MobileGL::Int backendValue = 0;
+        if (MobileGL::MG_Backend::gBackendFunctionsTable.GL.GetIntegeri_v) {
+            MobileGL::MG_Backend::gBackendFunctionsTable.GL.GetIntegeri_v(GL_MAX_COMPUTE_WORK_GROUP_SIZE, index,
+                                                                          &backendValue);
+        }
+
+        // TODO: Share these exposed compute limit helpers with GL_Getter.cpp instead of duplicating the frontend minima.
+        return std::max(static_cast<MobileGL::Uint>(std::max(backendValue, 0)),
+                        kFrontendMinComputeWorkGroupSizes[index]);
+    }
+
+    static unsigned long long GetComputeWorkGroupInvocationLimit() {
+        constexpr unsigned long long kFrontendMaxComputeWorkGroupInvocations = 1024;
+        if (!MobileGL::MG_Backend::pActiveBackendObject) return kFrontendMaxComputeWorkGroupInvocations;
+
+        return std::max(static_cast<unsigned long long>(std::max(
+                            MobileGL::MG_Backend::pActiveBackendObject->GetDynamicParameters()
+                                .MaxComputeWorkGroupInvocations,
+                            0)),
+                        kFrontendMaxComputeWorkGroupInvocations);
+    }
+
+    static std::optional<MobileGL::String> ValidateComputeLocalSizeLimits(const MobileGL::String& source) {
+        const ComputeLocalSize localSize = ParseComputeLocalSize(source);
+        if (!localSize.declared) return std::nullopt;
+
+        if (localSize.x > GetComputeWorkGroupSizeLimit(0) || localSize.y > GetComputeWorkGroupSizeLimit(1) ||
+            localSize.z > GetComputeWorkGroupSizeLimit(2)) {
+            return "Compute shader local_size exceeds GL_MAX_COMPUTE_WORK_GROUP_SIZE.";
+        }
+
+        const unsigned long long invocations = static_cast<unsigned long long>(localSize.x) * localSize.y * localSize.z;
+        if (invocations > GetComputeWorkGroupInvocationLimit()) {
+            return "Compute shader local_size product exceeds GL_MAX_COMPUTE_WORK_GROUP_INVOCATIONS.";
+        }
+
+        return std::nullopt;
+    }
+}
 
 namespace MobileGL::MG_State::GLState {
     void ShaderObject::SetShaderSource(const String& source) {
@@ -33,6 +159,16 @@ namespace MobileGL::MG_State::GLState {
         String compileSource = m_source;
         MG_Util::ShaderTranspiler::PreprocessShaderSource(m_stage, compileSource);
 
+        if (m_stage == ShaderStage::Compute) {
+            const std::optional<String> localSizeError = ValidateComputeLocalSizeLimits(compileSource);
+            if (localSizeError) {
+                m_compileStatus = false;
+                m_shader.reset();
+                m_infoLog = *localSizeError;
+                return;
+            }
+        }
+
         // Compile for OpenGL here, so that we can do validation and link
         // like a real OpenGL driver at linking stage
         // Will compile for other backends later.

MobileGL/MG_Util/ShaderTranspiler/ShaderCompiler.cpp

@@ -47,7 +47,8 @@ namespace MobileGL {
                 Resources.maxComputeWorkGroupCountZ = 65535;
                 Resources.maxComputeWorkGroupSizeX = 1024;
                 Resources.maxComputeWorkGroupSizeY = 1024;
-                Resources.maxComputeWorkGroupSizeZ = 64;
+                // TODO: Drive glslang compute resource limits from the active backend instead of this permissive cap.
+                Resources.maxComputeWorkGroupSizeZ = 1024;
                 Resources.maxComputeUniformComponents = 1024;
                 Resources.maxComputeTextureImageUnits = 16;
                 Resources.maxComputeImageUniforms = 8;