[Fix] (DirectVulkan, ShaderTranspiler, MG_IntegrationTest, SelfTest, TraceReplay): use native subgroups and patch iterationRP's under-declared scratch

iterationRP's Program 203 declares shared vec2 prefixSumCache[32] for a 512-invocation workgroup indexed by gl_SubgroupID; any device narrower than 16 lanes partitions into more than 32 subgroups and the pack writes shared memory out of bounds (heap corruption on lavapipe's CPU rasterizer, ssim 0.028 on the CI retrace). Fix it where the fault lies - in the fixture - and keep the GL contract sound everywhere else: - FixIterationRPSubgroupScratchPass: fingerprint-gated SPIR-V pass that grows exactly that array to ceil(invocations/width) entries on sub-16-lane devices; every other module passes through byte-identical. - DeriveNumSubgroupsPass stays default-on for the Adreno topology bug and is made spec-sound: pipelines request REQUIRE_FULL_SUBGROUPS whenever the workgroup shape makes the flag legal (computeFullSubgroups enabled, local_size_x a multiple of the native width, subgroup count within maxComputeWorkgroupSubgroups). - EmulateSubgroupsPass: 32-lane virtual-subgroup lowering kept in-tree as a last resort, enabled only by MOBILEGL_MAGMA_EMULATE_SUBGROUP=1 on devices with no native subgroup support; fails closed on extended subgroup instructions and on modules whose added scratch would exceed maxComputeSharedMemorySize. - IterationRPFirstReductionScenario skips gracefully outside the pack's 16..256-lane source domain; the new IterationRPScratchFixScenario runs the fixture-shaped reduction on any width and asserts the exact width-independent total. DriverPost keeps reporting FAIL on out-of-domain devices. - Program203 -> IterationRP rename throughout; the per-trace num_subgroups_quirk plumbing is removed from the trace replayer, JNI chain, and CI workflows.
2026-09-13 14:48:32 +09:00 · 2026-08-19 09:48:11 -04:00
parent 2b6c2b561c
commit d8576a2ed3
38 changed files with 3354 additions and 396 deletions
@@ -8,7 +8,6 @@

 #include "ProgramFactory.h"

-#include "Config.h"
 #include "MG_Backend/DirectVulkan/DirectVulkanResourceState.h"
 #include "MG_Util/ShaderTranspiler/ShaderCompiler.h"
 #include "MG_Util/ShaderTranspiler/SpvcSession.h"
@@ -34,6 +33,32 @@ namespace MobileGL::MG_Backend::DirectVulkan {
        using SpvcSession = MG_Util::ShaderTranspiler::SpvcSession;
        using SessionUsageBit = MG_Util::ShaderTranspiler::SessionUsageBit;

+        // Local size of a compute module, read from OpExecutionMode LocalSize; all-zero
+        // when absent. The compile chain pins SPIR-V 1.3, where a literal local size
+        // always reaches the module as this execution mode (LocalSizeId does not exist
+        // yet).
+        struct ComputeLocalSize {
+            Uint32 x = 0;
+            Uint32 y = 0;
+            Uint32 z = 0;
+            Uint64 Total() const { return static_cast<Uint64>(x) * y * z; }
+        };
+        ComputeLocalSize TryGetComputeLocalSize(const Vector<Uint>& spirv) {
+            constexpr SizeT kHeaderWords = 5;
+            constexpr Uint32 kOpExecutionMode = 16;
+            constexpr Uint32 kModeLocalSize = 17;
+            for (SizeT offset = kHeaderWords; offset < spirv.size();) {
+                const Uint32 wordCount = spirv[offset] >> 16u;
+                const Uint32 opcode = spirv[offset] & 0xffffu;
+                if (wordCount == 0 || offset + wordCount > spirv.size()) break;
+                if (opcode == kOpExecutionMode && wordCount >= 6 && spirv[offset + 2] == kModeLocalSize) {
+                    return {spirv[offset + 3], spirv[offset + 4], spirv[offset + 5]};
+                }
+                offset += wordCount;
+            }
+            return {};
+        }
+
        struct DescriptorKey {
            ProgramFactory::DescriptorBindingKind kind = ProgramFactory::DescriptorBindingKind::None;
            String name;
@@ -3164,20 +3189,56 @@ namespace MobileGL::MG_Backend::DirectVulkan {
                }
            }

-            // NumSubgroups is defined by the local workgroup dimensions and SubgroupSize. Derive
-            // it in SPIR-V instead of trusting a driver builtin that can disagree with the
-            // SubgroupId topology produced by the same compute dispatch (Adreno reports 1 while
-            // emitting IDs 0..7 for a 512-invocation, 64-wide workgroup).
-            if (MG_Config::Features.NumSubgroupsQuirk && shaders[i] &&
-                shaders[i]->GetShaderStage() == ShaderStage::Compute) {
-                Vector<Uint> derivedNumSubgroupsSpirv;
-                if (MG_Util::ShaderTranspiler::ShaderCompiler::DeriveNumSubgroupsForVulkan(
-                        moduleSpirvs[i], derivedNumSubgroupsSpirv, enableSpirvValidation)) {
-                    moduleSpirvs[i] = std::move(derivedNumSubgroupsSpirv);
+            // GL_KHR_shader_subgroup handling (SubgroupSupportPolicy.h). Native subgroup
+            // operations execute natively; two module repairs keep the GL contract intact
+            // around them. The opt-in emulation path replaces them only on devices with no
+            // subgroup support at all (MOBILEGL_MAGMA_EMULATE_SUBGROUP).
+            if (shaders[i] && shaders[i]->GetShaderStage() == ShaderStage::Compute) {
+                if (m_subgroupPolicy.emulateSubgroups) {
+                    Vector<Uint> emulatedSpirv;
+                    if (MG_Util::ShaderTranspiler::ShaderCompiler::EmulateSubgroupsForVulkan(
+                            moduleSpirvs[i], emulatedSpirv,
+                            m_subgroupPolicy.maxComputeSharedMemoryBytes, enableSpirvValidation)) {
+                        moduleSpirvs[i] = std::move(emulatedSpirv);
+                    } else {
+                        MGLOG_E("ProgramFactory: subgroup emulation failed for program %u; the "
+                                "module keeps subgroup operations the device cannot execute",
+                                program.GetExternalIndex());
+                    }
                } else {
-                    MGLOG_E("ProgramFactory: failed to derive gl_NumSubgroups for program %u; "
-                            "compute shaders may observe a driver-inconsistent subgroup count",
-                            program.GetExternalIndex());
+                    // iterationRP under-declares its cross-subgroup scratch
+                    // (prefixSumCache[32] for 512 invocations); on a sub-16-lane device
+                    // grow that one fingerprinted array to what the topology needs.
+                    if (m_subgroupPolicy.fixIterationRPSubgroupScratch) {
+                        Vector<Uint> patchedSpirv;
+                        if (MG_Util::ShaderTranspiler::ShaderCompiler::FixIterationRPSubgroupScratchForVulkan(
+                                moduleSpirvs[i], patchedSpirv, m_subgroupPolicy.nativeSubgroupSize,
+                                enableSpirvValidation)) {
+                            moduleSpirvs[i] = std::move(patchedSpirv);
+                        } else {
+                            MGLOG_E("ProgramFactory: iterationRP subgroup scratch patch failed for "
+                                    "program %u; the pack's declared array sizes stay in effect",
+                                    program.GetExternalIndex());
+                        }
+                    }
+                    // gl_NumSubgroups must agree with the gl_SubgroupID range GL promises;
+                    // derive it from the workgroup dimensions and gl_SubgroupSize instead of
+                    // trusting a driver builtin that can disagree with the topology the same
+                    // dispatch emits (Adreno reports 1 while emitting IDs 0..7 for a
+                    // 512-invocation, 64-wide workgroup). The ceil() partition this derives
+                    // is pinned by REQUIRE_FULL_SUBGROUPS at pipeline creation whenever the
+                    // workgroup shape makes that flag legal (see the stage setup below).
+                    if (m_subgroupPolicy.deriveNumSubgroups) {
+                        Vector<Uint> derivedNumSubgroupsSpirv;
+                        if (MG_Util::ShaderTranspiler::ShaderCompiler::DeriveNumSubgroupsForVulkan(
+                                moduleSpirvs[i], derivedNumSubgroupsSpirv, enableSpirvValidation)) {
+                            moduleSpirvs[i] = std::move(derivedNumSubgroupsSpirv);
+                        } else {
+                            MGLOG_E("ProgramFactory: failed to derive gl_NumSubgroups for program %u; "
+                                    "compute shaders may observe a driver-inconsistent subgroup count",
+                                    program.GetExternalIndex());
+                        }
+                    }
                }
            }

@@ -3326,6 +3387,27 @@ namespace MobileGL::MG_Backend::DirectVulkan {
            stage.stage = ToVkStage(shaderStage);
            stage.module = module;
            stage.pName = "main";
+            // Pin the full-subgroup launch the derived gl_NumSubgroups assumes. Legal
+            // exactly when the computeFullSubgroups feature is enabled and local_size_x is
+            // a multiple of the subgroup size (VUID-VkPipelineShaderStageCreateInfo-
+            // flags-02759/-02785), and only worth requesting while the resulting subgroup
+            // count fits the device's maxComputeWorkgroupSubgroups (lavapipe caps it at
+            // 32, below a 512-invocation dispatch's 64). With the bit set, "Full
+            // Subgroups" guarantees every subgroup launches with all invocations active,
+            // making the subgroup count exactly invocations / size. Shapes the flag
+            // cannot cover (e.g. 32x16 on a 64-wide device) fall back to the driver's
+            // own - spec-encouraged - tight partitioning, which the DriverPost witness
+            // verifies per device.
+            if (shaderStage == ShaderStage::Compute && m_subgroupPolicy.requireFullSubgroups &&
+                !m_subgroupPolicy.emulateSubgroups && m_subgroupPolicy.nativeSubgroupSize != 0) {
+                const ComputeLocalSize localSize = TryGetComputeLocalSize(moduleSpv);
+                const Uint64 fullSubgroupCount =
+                    localSize.Total() / m_subgroupPolicy.nativeSubgroupSize;
+                if (localSize.x != 0 && localSize.x % m_subgroupPolicy.nativeSubgroupSize == 0 &&
+                    fullSubgroupCount <= m_subgroupPolicy.maxComputeWorkgroupSubgroups) {
+                    stage.flags |= VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT;
+                }
+            }

            entry.modules.push_back(module);
            entry.stages.push_back(stage);
@@ -372,16 +372,38 @@ namespace MobileGL::MG_Backend::DirectVulkan {
            virtual void OnProgramEvicted(HashType programHash, VkDescriptorSetLayout descriptorSetLayout) = 0;
        };

+        // How this factory's compute modules implement GL_KHR_shader_subgroup. Computed
+        // once at renderer initialization (SubgroupSupportPolicy.h + the device's
+        // subgroup properties) so lowering can never disagree with the advertised
+        // capabilities. Native subgroup operations always execute natively; the two
+        // repair passes patch modules AROUND them, and the emulation only replaces them
+        // on opted-in devices with no subgroup support at all.
+        struct SubgroupLoweringPolicy {
+            Bool emulateSubgroups = false;      // MOBILEGL_MAGMA_EMULATE_SUBGROUP, no-native-support devices
+            Bool fixIterationRPSubgroupScratch = false; // patch iterationRP's under-declared scratch
+            Bool deriveNumSubgroups = false;    // repair the NumSubgroups builtin
+            Bool requireFullSubgroups = false;  // computeFullSubgroups enabled on the device
+            Uint32 nativeSubgroupSize = 0;
+            // Full-subgroup launches are bounded by this device limit; a dispatch whose
+            // workgroup needs more subgroups than this cannot request the flag.
+            Uint32 maxComputeWorkgroupSubgroups = 0;
+            // VkPhysicalDeviceLimits::maxComputeSharedMemorySize; bounds the scratch the
+            // emulation pass may add (0 falls back to the Vulkan minimum, 16384).
+            Uint32 maxComputeSharedMemoryBytes = 0;
+        };
+
        explicit ProgramFactory(VkDevice device, const VulkanRendererConfig& config, Uint32 maxBindings,
                                Bool shaderDrawParametersEnabled,
                                Bool unformattedFloatStorageImagesEnabled,
                                Bool enableSpirvValidation,
-                                UpdateAfterBindLimits updateAfterBindLimits)
+                                UpdateAfterBindLimits updateAfterBindLimits,
+                                SubgroupLoweringPolicy subgroupPolicy)
            : m_device(device), m_maxBindings(maxBindings), m_config(config),
              m_shaderDrawParametersEnabled(shaderDrawParametersEnabled),
              m_unformattedFloatStorageImagesEnabled(unformattedFloatStorageImagesEnabled),
              m_enableSpirvValidation(enableSpirvValidation),
-              m_updateAfterBindLimits(updateAfterBindLimits) {
+              m_updateAfterBindLimits(updateAfterBindLimits),
+              m_subgroupPolicy(subgroupPolicy) {
            VkProgramObject::s_device = device;
        }
        // Destroys the pass-through tessellation control modules. Runs while the device is
@@ -511,6 +533,7 @@ namespace MobileGL::MG_Backend::DirectVulkan {
        // the factory lets each reflected layout choose ordinary descriptors when its
        // own counts would exceed the update-after-bind budget.
        UpdateAfterBindLimits m_updateAfterBindLimits{};
+        SubgroupLoweringPolicy m_subgroupPolicy{};
        // See SetDefaultFramebufferHeight. 0 means "not known yet"; the FragCoordYFlip bit is
        // never set before the swapchain exists, so no variant can be compiled against it.
        Uint32 m_defaultFramebufferHeight = 0;
@@ -8,6 +8,7 @@

 #include "VulkanRenderer.h"

+#include "MG_Backend/DirectVulkan/SubgroupSupportPolicy.h"
 #include "MG_Backend/DirectGLES/Utils.h"
 #include "VertexInputStateFactory.h"
 #include "VertexInputStateBuilder.h"
@@ -3058,11 +3059,22 @@ void main() {
            }
            PipelineFactory::SetSuppressBlendedDepthWrite(suppressBlendedDepthWrite);
        }
+        ProgramFactory::SubgroupLoweringPolicy subgroupPolicy{};
+        subgroupPolicy.emulateSubgroups = ShouldEmulateSubgroups(m_nativeSubgroupSupported);
+        subgroupPolicy.fixIterationRPSubgroupScratch =
+            m_nativeSubgroupSupported && ShouldFixIterationRPSubgroupScratch();
+        subgroupPolicy.deriveNumSubgroups =
+            m_nativeSubgroupSupported && ShouldDeriveNumSubgroups();
+        subgroupPolicy.requireFullSubgroups = m_computeFullSubgroupsFeatureEnabled;
+        subgroupPolicy.nativeSubgroupSize = m_nativeSubgroupSize;
+        subgroupPolicy.maxComputeWorkgroupSubgroups = m_maxComputeWorkgroupSubgroups;
+        subgroupPolicy.maxComputeSharedMemoryBytes =
+            m_physicalDevice.properties.limits.maxComputeSharedMemorySize;
        m_programFactory = MakeUnique<ProgramFactory>(m_device, m_config, maxProgramBindings,
                                                      m_shaderDrawParametersFeatureEnabled,
                                                      m_unformattedFloatStorageImagesEnabled,
                                                      MG_Config::Features.EnableSpirvValidation,
-                                                      m_updateAfterBindLimits);
+                                                      m_updateAfterBindLimits, subgroupPolicy);
        MOBILEGL_ASSERT(m_programFactory != nullptr, "ProgramFactory creation failed.");
        // The swapchain already exists at this point (Initialize creates it first), so seed the
        // height the factory could not be told about from CreateSwapchain.
@@ -12740,6 +12752,73 @@ void main() {
            }
        }

+        // Native subgroup topology, and VK_EXT_subgroup_size_control's
+        // computeFullSubgroups feature. REQUIRE_FULL_SUBGROUPS on a compute stage is what
+        // turns the derived gl_NumSubgroups (DeriveNumSubgroupsPass) from
+        // encouraged-but-unspecified driver behaviour into a spec guarantee: with the bit
+        // set and local_size_x a multiple of the subgroup size, every subgroup launches
+        // full, so the subgroup count is exactly invocations / size ("Full Subgroups",
+        // VUID-VkPipelineShaderStageCreateInfo-flags-02759/-02785).
+        m_nativeSubgroupSize = 0;
+        m_nativeSubgroupSupported = false;
+        m_computeFullSubgroupsFeatureEnabled = false;
+        if (getPhysicalDeviceProperties2 != nullptr) {
+            VkPhysicalDeviceSubgroupProperties subgroupProperties{};
+            subgroupProperties.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_PROPERTIES;
+            VkPhysicalDeviceProperties2 subgroupPropertyQuery{};
+            subgroupPropertyQuery.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2;
+            subgroupPropertyQuery.pNext = &subgroupProperties;
+            getPhysicalDeviceProperties2(m_physicalDevice.handle, &subgroupPropertyQuery);
+            // Mirrors the loader's HasUsableShaderSubgroupSupport gate, including the
+            // MOBILEGL_DISABLE_SUBGROUP escape hatch, so the module lowerings can never
+            // disagree with the advertised capabilities.
+            const Bool usableSubgroups =
+                subgroupProperties.subgroupSize > 0 &&
+                (subgroupProperties.supportedStages & VK_SHADER_STAGE_COMPUTE_BIT) != 0 &&
+                (subgroupProperties.supportedOperations & VK_SUBGROUP_FEATURE_BASIC_BIT) != 0;
+            if (usableSubgroups && !MG_Config::Features.DisableSubgroup) {
+                m_nativeSubgroupSize = subgroupProperties.subgroupSize;
+                m_nativeSubgroupSupported = true;
+            }
+        }
+        VkPhysicalDeviceSubgroupSizeControlFeaturesEXT subgroupSizeControlFeatures{};
+        subgroupSizeControlFeatures.sType =
+            VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_FEATURES_EXT;
+        m_maxComputeWorkgroupSubgroups = 0;
+        if (m_nativeSubgroupSupported &&
+            IsExtensionSupported(availableExtensions, VK_EXT_SUBGROUP_SIZE_CONTROL_EXTENSION_NAME) &&
+            getPhysicalDeviceFeatures2 != nullptr) {
+            VkPhysicalDeviceFeatures2 featureQuery{};
+            featureQuery.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2;
+            featureQuery.pNext = &subgroupSizeControlFeatures;
+            getPhysicalDeviceFeatures2(m_physicalDevice.handle, &featureQuery);
+            if (getPhysicalDeviceProperties2 != nullptr) {
+                VkPhysicalDeviceSubgroupSizeControlPropertiesEXT subgroupSizeControlProperties{};
+                subgroupSizeControlProperties.sType =
+                    VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SUBGROUP_SIZE_CONTROL_PROPERTIES_EXT;
+                VkPhysicalDeviceProperties2 propertyQuery{};
+                propertyQuery.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2;
+                propertyQuery.pNext = &subgroupSizeControlProperties;
+                getPhysicalDeviceProperties2(m_physicalDevice.handle, &propertyQuery);
+                m_maxComputeWorkgroupSubgroups =
+                    subgroupSizeControlProperties.maxComputeWorkgroupSubgroups;
+            }
+            if (subgroupSizeControlFeatures.computeFullSubgroups == VK_TRUE) {
+                if (!IsExtensionAlreadyEnabled(enabledDeviceExtensions,
+                                               VK_EXT_SUBGROUP_SIZE_CONTROL_EXTENSION_NAME)) {
+                    enabledDeviceExtensions.push_back(VK_EXT_SUBGROUP_SIZE_CONTROL_EXTENSION_NAME);
+                }
+                // Only the full-subgroups guarantee is wanted; required/varying subgroup
+                // sizes stay unrequested.
+                subgroupSizeControlFeatures.subgroupSizeControl = VK_FALSE;
+                subgroupSizeControlFeatures.pNext = const_cast<void*>(deviceCreateInfo.pNext);
+                deviceCreateInfo.pNext = &subgroupSizeControlFeatures;
+                m_computeFullSubgroupsFeatureEnabled = true;
+                MGLOG_I("Enabled optional device extension: %s (computeFullSubgroups)",
+                        VK_EXT_SUBGROUP_SIZE_CONTROL_EXTENSION_NAME);
+            }
+        }
+
        // VK_EXT_transform_feedback backs GL transform feedback capture.
        m_transformFeedbackFeatureEnabled = false;
        VkPhysicalDeviceTransformFeedbackFeaturesEXT transformFeedbackFeatures{};
@@ -554,6 +554,16 @@ namespace MobileGL::MG_Backend::DirectVulkan {
        Bool m_samplerAnisotropyFeatureEnabled = false;
        Bool m_shaderDrawParametersExtensionEnabled = false;
        Bool m_shaderDrawParametersFeatureEnabled = false;
+        // Native subgroup topology, queried at device creation for the compute-module
+        // subgroup repairs (SubgroupSupportPolicy.h) and the REQUIRE_FULL_SUBGROUPS
+        // stage flag; 0 / false when the device has no usable compute subgroups or
+        // MOBILEGL_DISABLE_SUBGROUP forced them off.
+        Uint32 m_nativeSubgroupSize = 0;
+        Bool m_nativeSubgroupSupported = false;
+        Bool m_computeFullSubgroupsFeatureEnabled = false;
+        // VkPhysicalDeviceSubgroupSizeControlProperties::maxComputeWorkgroupSubgroups;
+        // 0 when the extension (and therefore the full-subgroups flag) is unavailable.
+        Uint32 m_maxComputeWorkgroupSubgroups = 0;
        Bool m_unformattedFloatStorageImagesEnabled = false;
        // Set only after descriptor-indexing feature AND property queries prove that
        // update-after-bind is legal for every descriptor category this renderer emits.