From d35e452368b0364bea5d6c6e33377d6e6750c37c Mon Sep 17 00:00:00 2001 From: Swung0x48 Date: Mon, 6 Jul 2026 01:25:44 +0800 Subject: [PATCH] [Fix] (MG_Util/ShaderTranspiler): rewrite workgroup vec3 composite loads --- MobileGL/MG_Test/Program/ProgramUtilTest.cpp | 9 +- .../DecomposeWorkgroupVec3Pass.cpp | 354 ++++++++++++------ 2 files changed, 246 insertions(+), 117 deletions(-) diff --git a/MobileGL/MG_Test/Program/ProgramUtilTest.cpp b/MobileGL/MG_Test/Program/ProgramUtilTest.cpp index c2e2e473..d9e875ba 100644 --- a/MobileGL/MG_Test/Program/ProgramUtilTest.cpp +++ b/MobileGL/MG_Test/Program/ProgramUtilTest.cpp @@ -832,10 +832,7 @@ void main() { vec3 loaded = shared_memory[row][col]; float x = shared_memory[row][col].x; - vec3 rowCopy[9]; - for (uint i = 0u; i < 9u; ++i) { - rowCopy[i] = shared_memory[row][i]; - } + vec3 rowCopy[9] = shared_memory[0]; memoryBarrierShared(); barrier(); @@ -893,6 +890,10 @@ TEST_F(ProgramUtilTest, DecomposeWorkgroupVec3InSpirvPass) { EXPECT_NE(source.find("shared float"), std::string::npos) << "Expected `shared float` in decomposed output:\n" << source; + + EXPECT_EQ(source.find("= shared_memory[0]"), std::string::npos) + << "Decomposed output kept an invalid whole-row shared-memory load:\n" + << source; } TEST_F(ProgramUtilTest, DecomposeWorkgroupVec3IgnoresNonWorkgroupVec3) { diff --git a/MobileGL/MG_Util/ShaderTranspiler/SpirvPasses/DecomposeWorkgroupVec3Pass.cpp b/MobileGL/MG_Util/ShaderTranspiler/SpirvPasses/DecomposeWorkgroupVec3Pass.cpp index 08a73e00..dc4b5d3e 100644 --- a/MobileGL/MG_Util/ShaderTranspiler/SpirvPasses/DecomposeWorkgroupVec3Pass.cpp +++ b/MobileGL/MG_Util/ShaderTranspiler/SpirvPasses/DecomposeWorkgroupVec3Pass.cpp @@ -100,6 +100,17 @@ namespace MobileGL { return typeId; } + uint32_t GetArrayLength(IRContext* context, const Instruction* arrayTypeInst) { + assert(arrayTypeInst != nullptr && arrayTypeInst->opcode() == spv::Op::OpTypeArray); + const uint32_t lengthId = arrayTypeInst->GetSingleWordInOperand(1); + const Instruction* lengthInst = context->get_def_use_mgr()->GetDef(lengthId); + if (lengthInst == nullptr || lengthInst->opcode() != spv::Op::OpConstant || + lengthInst->NumInOperands() == 0) { + return 0; + } + return lengthInst->GetSingleWordInOperand(0); + } + // Inserts an instruction before |where|, sets its debug info, and updates // def-use and block mapping. Returns a pointer to the inserted instruction. Instruction* InsertBefore(BasicBlock* block, Instruction* where, @@ -112,6 +123,191 @@ namespace MobileGL { context->set_instr_block(&*iter, block); return &*iter; } + + uint32_t BuildCompositeLoad(IRContext* context, BasicBlock* block, Instruction* where, + uint32_t originalTypeId, uint32_t decomposedPtrId, + const std::unordered_map& vec3ToArr3, + const Instruction* sourceLoad) { + auto* defUseMgr = context->get_def_use_mgr(); + auto* typeMgr = context->get_type_mgr(); + auto* constMgr = context->get_constant_mgr(); + Instruction* originalTypeInst = defUseMgr->GetDef(originalTypeId); + assert(originalTypeInst != nullptr); + + if (originalTypeInst->opcode() == spv::Op::OpTypeVector) { + const auto vecIt = vec3ToArr3.find(originalTypeId); + assert(vecIt != vec3ToArr3.end()); + const uint32_t scalarTypeId = originalTypeInst->GetSingleWordInOperand(0); + const uint32_t ptrScalarTypeId = + typeMgr->FindPointerToType(scalarTypeId, spv::StorageClass::Workgroup); + std::vector componentIds; + componentIds.reserve(3); + for (uint32_t c = 0; c < 3; ++c) { + const uint32_t constCId = constMgr->GetUIntConstId(c); + const uint32_t acId = context->TakeNextId(); + auto acInst = MakeUnique( + context, spv::Op::OpAccessChain, ptrScalarTypeId, acId, + std::initializer_list{ + {SPV_OPERAND_TYPE_ID, {decomposedPtrId}}, + {SPV_OPERAND_TYPE_ID, {constCId}}}); + InsertBefore(block, where, std::move(acInst), context); + + const uint32_t loadId = context->TakeNextId(); + auto loadInst = MakeUnique( + context, spv::Op::OpLoad, scalarTypeId, loadId, + std::initializer_list{ + {SPV_OPERAND_TYPE_ID, {acId}}}); + for (uint32_t opIdx = 1; opIdx < sourceLoad->NumInOperands(); ++opIdx) { + loadInst->AddOperand(Operand(sourceLoad->GetInOperand(opIdx))); + } + Instruction* loadPtr = InsertBefore(block, where, std::move(loadInst), context); + componentIds.push_back(loadPtr->result_id()); + } + + const uint32_t compositeId = context->TakeNextId(); + auto composite = MakeUnique( + context, spv::Op::OpCompositeConstruct, originalTypeId, compositeId, + std::initializer_list{}); + for (uint32_t componentId : componentIds) { + composite->AddOperand({SPV_OPERAND_TYPE_ID, {componentId}}); + } + InsertBefore(block, where, std::move(composite), context); + return compositeId; + } + + if (originalTypeInst->opcode() == spv::Op::OpTypeArray) { + const uint32_t originalElemTypeId = originalTypeInst->GetSingleWordInOperand(0); + const uint32_t vec3LeafId = FindVec3LeafTypeId(context, originalElemTypeId); + assert(vec3LeafId != 0); + const auto arr3It = vec3ToArr3.find(vec3LeafId); + assert(arr3It != vec3ToArr3.end()); + const uint32_t decomposedElemTypeId = + originalElemTypeId == vec3LeafId + ? arr3It->second + : RebuildPointeeType(context, originalElemTypeId, arr3It->second, vec3ToArr3); + const uint32_t ptrDecomposedElemTypeId = + typeMgr->FindPointerToType(decomposedElemTypeId, spv::StorageClass::Workgroup); + const uint32_t length = GetArrayLength(context, originalTypeInst); + assert(length != 0); + + std::vector elementIds; + elementIds.reserve(length); + for (uint32_t i = 0; i < length; ++i) { + const uint32_t constIId = constMgr->GetUIntConstId(i); + const uint32_t elemPtrId = context->TakeNextId(); + auto elemAc = MakeUnique( + context, spv::Op::OpAccessChain, ptrDecomposedElemTypeId, elemPtrId, + std::initializer_list{ + {SPV_OPERAND_TYPE_ID, {decomposedPtrId}}, + {SPV_OPERAND_TYPE_ID, {constIId}}}); + InsertBefore(block, where, std::move(elemAc), context); + elementIds.push_back(BuildCompositeLoad(context, block, where, originalElemTypeId, + elemPtrId, vec3ToArr3, sourceLoad)); + } + + const uint32_t compositeId = context->TakeNextId(); + auto composite = MakeUnique( + context, spv::Op::OpCompositeConstruct, originalTypeId, compositeId, + std::initializer_list{}); + for (uint32_t elementId : elementIds) { + composite->AddOperand({SPV_OPERAND_TYPE_ID, {elementId}}); + } + InsertBefore(block, where, std::move(composite), context); + return compositeId; + } + + assert(false && "DecomposeWorkgroupVec3Pass: unsupported composite load type"); + return 0; + } + + void BuildCompositeStore(IRContext* context, BasicBlock* block, Instruction* where, + uint32_t originalTypeId, uint32_t decomposedPtrId, uint32_t valueId, + const std::unordered_map& vec3ToArr3, + const Instruction* sourceStore) { + auto* defUseMgr = context->get_def_use_mgr(); + auto* typeMgr = context->get_type_mgr(); + auto* constMgr = context->get_constant_mgr(); + Instruction* originalTypeInst = defUseMgr->GetDef(originalTypeId); + assert(originalTypeInst != nullptr); + + if (originalTypeInst->opcode() == spv::Op::OpTypeVector) { + const auto vecIt = vec3ToArr3.find(originalTypeId); + assert(vecIt != vec3ToArr3.end()); + const uint32_t scalarTypeId = originalTypeInst->GetSingleWordInOperand(0); + const uint32_t ptrScalarTypeId = + typeMgr->FindPointerToType(scalarTypeId, spv::StorageClass::Workgroup); + for (uint32_t c = 0; c < 3; ++c) { + const uint32_t extractId = context->TakeNextId(); + auto extract = MakeUnique( + context, spv::Op::OpCompositeExtract, scalarTypeId, extractId, + std::initializer_list{ + {SPV_OPERAND_TYPE_ID, {valueId}}, + {SPV_OPERAND_TYPE_LITERAL_INTEGER, {c}}}); + InsertBefore(block, where, std::move(extract), context); + + const uint32_t constCId = constMgr->GetUIntConstId(c); + const uint32_t acId = context->TakeNextId(); + auto acInst = MakeUnique( + context, spv::Op::OpAccessChain, ptrScalarTypeId, acId, + std::initializer_list{ + {SPV_OPERAND_TYPE_ID, {decomposedPtrId}}, + {SPV_OPERAND_TYPE_ID, {constCId}}}); + InsertBefore(block, where, std::move(acInst), context); + + auto compStore = MakeUnique( + context, spv::Op::OpStore, 0, 0, + std::initializer_list{ + {SPV_OPERAND_TYPE_ID, {acId}}, + {SPV_OPERAND_TYPE_ID, {extractId}}}); + for (uint32_t opIdx = 2; opIdx < sourceStore->NumInOperands(); ++opIdx) { + compStore->AddOperand(Operand(sourceStore->GetInOperand(opIdx))); + } + InsertBefore(block, where, std::move(compStore), context); + } + return; + } + + if (originalTypeInst->opcode() == spv::Op::OpTypeArray) { + const uint32_t originalElemTypeId = originalTypeInst->GetSingleWordInOperand(0); + const uint32_t vec3LeafId = FindVec3LeafTypeId(context, originalElemTypeId); + assert(vec3LeafId != 0); + const auto arr3It = vec3ToArr3.find(vec3LeafId); + assert(arr3It != vec3ToArr3.end()); + const uint32_t decomposedElemTypeId = + originalElemTypeId == vec3LeafId + ? arr3It->second + : RebuildPointeeType(context, originalElemTypeId, arr3It->second, vec3ToArr3); + const uint32_t ptrDecomposedElemTypeId = + typeMgr->FindPointerToType(decomposedElemTypeId, spv::StorageClass::Workgroup); + const uint32_t length = GetArrayLength(context, originalTypeInst); + assert(length != 0); + + for (uint32_t i = 0; i < length; ++i) { + const uint32_t extractId = context->TakeNextId(); + auto extract = MakeUnique( + context, spv::Op::OpCompositeExtract, originalElemTypeId, extractId, + std::initializer_list{ + {SPV_OPERAND_TYPE_ID, {valueId}}, + {SPV_OPERAND_TYPE_LITERAL_INTEGER, {i}}}); + InsertBefore(block, where, std::move(extract), context); + + const uint32_t constIId = constMgr->GetUIntConstId(i); + const uint32_t elemPtrId = context->TakeNextId(); + auto elemAc = MakeUnique( + context, spv::Op::OpAccessChain, ptrDecomposedElemTypeId, elemPtrId, + std::initializer_list{ + {SPV_OPERAND_TYPE_ID, {decomposedPtrId}}, + {SPV_OPERAND_TYPE_ID, {constIId}}}); + InsertBefore(block, where, std::move(elemAc), context); + + BuildCompositeStore(context, block, where, originalElemTypeId, elemPtrId, + extractId, vec3ToArr3, sourceStore); + } + return; + } + + assert(false && "DecomposeWorkgroupVec3Pass: unsupported composite store type"); + } } // namespace spvtools::opt::Pass::Status DecomposeWorkgroupVec3Pass::Process() { @@ -133,7 +329,7 @@ namespace MobileGL { // ===================================================================== std::unordered_set vec3TypeIds; std::unordered_map vec3ToArr3; - std::unordered_map ptrVec3ToPtrArr3; + std::unordered_map ptrCompositeToDecomposed; for (Instruction& typeInst : ctx->types_values()) { if (typeInst.opcode() == spv::Op::OpTypeVector && @@ -230,21 +426,29 @@ namespace MobileGL { continue; } const uint32_t pointeeId = typeInst.GetSingleWordInOperand(1); - auto it = vec3ToArr3.find(pointeeId); + const uint32_t vec3LeafId = FindVec3LeafTypeId(ctx, pointeeId); + if (vec3LeafId == 0) { + continue; + } + auto it = vec3ToArr3.find(vec3LeafId); if (it == vec3ToArr3.end()) { continue; } - const uint32_t ptrArrId = - typeMgr->FindPointerToType(it->second, spv::StorageClass::Workgroup); - ptrVec3ToPtrArr3[typeInst.result_id()] = ptrArrId; + const uint32_t newPointeeId = + pointeeId == vec3LeafId + ? it->second + : RebuildPointeeType(ctx, pointeeId, it->second, vec3ToArr3); + const uint32_t newPtrId = + typeMgr->FindPointerToType(newPointeeId, spv::StorageClass::Workgroup); + ptrCompositeToDecomposed[typeInst.result_id()] = newPtrId; } // ===================================================================== // Phase 3: Rewrite access chain result types. // Any OpAccessChain/OpInBoundsAccessChain whose result type is a - // ptr_Workgroup_vec3 must now produce ptr_Workgroup_arr3. - // (Access chains that go one level deeper to a component already have - // result type ptr_scalar, which is unaffected.) + // Workgroup pointer whose pointee contains vec3 must now produce the + // decomposed pointer type. Access chains that go one level deeper to a + // scalar component already have result type ptr_scalar and are unaffected. // ===================================================================== // We also need to handle chained access chains: an access chain whose // base is itself an access chain that was rewritten. Since the result @@ -258,31 +462,30 @@ namespace MobileGL { inst.opcode() != spv::Op::OpInBoundsAccessChain) { continue; } - auto it = ptrVec3ToPtrArr3.find(inst.type_id()); - if (it != ptrVec3ToPtrArr3.end()) { + auto it = ptrCompositeToDecomposed.find(inst.type_id()); + if (it != ptrCompositeToDecomposed.end()) { accessChainsToFix.push_back(&inst); } } } } for (Instruction* ac : accessChainsToFix) { - const uint32_t newTypeId = ptrVec3ToPtrArr3[ac->type_id()]; + const uint32_t newTypeId = ptrCompositeToDecomposed[ac->type_id()]; ac->SetResultType(newTypeId); defUseMgr->AnalyzeInstUse(ac); } - // Build the set of decomposed pointer type ids (ptr_Workgroup_arr3) for - // fast lookup when matching loads/stores. - std::unordered_set ptrArr3TypeIds; - for (const auto& [oldPtr, newPtr] : ptrVec3ToPtrArr3) { - ptrArr3TypeIds.insert(newPtr); + // Build the set of decomposed pointer type ids for fast lookup when matching + // loads/stores. + std::unordered_set decomposedPtrTypeIds; + for (const auto& [oldPtr, newPtr] : ptrCompositeToDecomposed) { + decomposedPtrTypeIds.insert(newPtr); } // ===================================================================== - // Phase 4: Rewrite whole-vec3 OpLoad. - // OpLoad %v3float %ptr (ptr is now ptr_Workgroup_arr3) - // -> 3x OpAccessChain %ptr_float %ptr %c + OpLoad %float - // -> OpCompositeConstruct %v3float %f0 %f1 %f2 + // Phase 4: Rewrite whole-composite OpLoad. + // OpLoad %v3float %ptr -> vec3 component loads + // OpLoad %arr_v3float %ptr_to_row -> per-element vec3 loads + array construct // ===================================================================== std::vector loadsToRewrite; for (auto& func : *ctx->module()) { @@ -291,7 +494,7 @@ namespace MobileGL { if (inst.opcode() != spv::Op::OpLoad) { continue; } - if (vec3ToArr3.find(inst.type_id()) == vec3ToArr3.end()) { + if (FindVec3LeafTypeId(ctx, inst.type_id()) == 0) { continue; } // Check the pointer operand's type. @@ -300,7 +503,7 @@ namespace MobileGL { if (ptrDef == nullptr) { continue; } - if (ptrArr3TypeIds.find(ptrDef->type_id()) == ptrArr3TypeIds.end()) { + if (decomposedPtrTypeIds.find(ptrDef->type_id()) == decomposedPtrTypeIds.end()) { continue; } loadsToRewrite.push_back(&inst); @@ -310,47 +513,9 @@ namespace MobileGL { for (Instruction* load : loadsToRewrite) { BasicBlock* block = ctx->get_instr_block(load); - const uint32_t vec3TypeId = load->type_id(); - const uint32_t floatTypeId = - defUseMgr->GetDef(vec3TypeId)->GetSingleWordInOperand(0); - const uint32_t ptrFloatTypeId = typeMgr->FindPointerToType( - floatTypeId, spv::StorageClass::Workgroup); const uint32_t ptrId = load->GetSingleWordInOperand(0); - - std::vector compLoadIds; - compLoadIds.reserve(3); - for (uint32_t c = 0; c < 3; ++c) { - const uint32_t constCId = constMgr->GetUIntConstId(c); - const uint32_t acId = ctx->TakeNextId(); - auto acInst = MakeUnique( - ctx, spv::Op::OpAccessChain, ptrFloatTypeId, acId, - std::initializer_list{ - {SPV_OPERAND_TYPE_ID, {ptrId}}, - {SPV_OPERAND_TYPE_ID, {constCId}}}); - Instruction* acPtr = InsertBefore(block, load, std::move(acInst), ctx); - - const uint32_t loadId = ctx->TakeNextId(); - auto loadInst = MakeUnique( - ctx, spv::Op::OpLoad, floatTypeId, loadId, - std::initializer_list{ - {SPV_OPERAND_TYPE_ID, {acId}}}); - // Copy memory operands (alignment etc.) from original load. - for (uint32_t opIdx = 1; opIdx < load->NumInOperands(); ++opIdx) { - loadInst->AddOperand(Operand(load->GetInOperand(opIdx))); - } - Instruction* loadPtr = - InsertBefore(block, load, std::move(loadInst), ctx); - compLoadIds.push_back(loadPtr->result_id()); - } - - const uint32_t compositeId = ctx->TakeNextId(); - auto composite = MakeUnique( - ctx, spv::Op::OpCompositeConstruct, vec3TypeId, compositeId, - std::initializer_list{}); - for (uint32_t cId : compLoadIds) { - composite->AddOperand({SPV_OPERAND_TYPE_ID, {cId}}); - } - InsertBefore(block, load, std::move(composite), ctx); + const uint32_t compositeId = + BuildCompositeLoad(ctx, block, load, load->type_id(), ptrId, vec3ToArr3, load); ctx->ReplaceAllUsesWith(load->result_id(), compositeId); ctx->KillNamesAndDecorates(load->result_id()); @@ -358,10 +523,9 @@ namespace MobileGL { } // ===================================================================== - // Phase 5: Rewrite whole-vec3 OpStore. - // OpStore %ptr %vec3val (ptr is now ptr_Workgroup_arr3) - // -> 3x OpCompositeExtract %float %val %c - // -> 3x OpAccessChain %ptr_float %ptr %c + OpStore %ptr_c %comp_c + // Phase 5: Rewrite whole-composite OpStore. + // OpStore %ptr %vec3val -> per-component scalar stores + // OpStore %ptr %arr_v3float -> per-element vec3 stores // ===================================================================== std::vector storesToRewrite; for (auto& func : *ctx->module()) { @@ -375,7 +539,12 @@ namespace MobileGL { if (ptrDef == nullptr) { continue; } - if (ptrArr3TypeIds.find(ptrDef->type_id()) == ptrArr3TypeIds.end()) { + if (decomposedPtrTypeIds.find(ptrDef->type_id()) == decomposedPtrTypeIds.end()) { + continue; + } + const uint32_t valId = inst.GetSingleWordInOperand(1); + Instruction* valDef = defUseMgr->GetDef(valId); + if (valDef == nullptr || FindVec3LeafTypeId(ctx, valDef->type_id()) == 0) { continue; } storesToRewrite.push_back(&inst); @@ -387,56 +556,15 @@ namespace MobileGL { BasicBlock* block = ctx->get_instr_block(store); const uint32_t ptrId = store->GetSingleWordInOperand(0); const uint32_t valId = store->GetSingleWordInOperand(1); - Instruction* ptrDef = defUseMgr->GetDef(ptrId); - const uint32_t arr3TypeId = ptrDef->type_id() == 0 - ? 0 - : defUseMgr->GetDef(ptrDef->type_id())->GetSingleWordInOperand(1); - // The pointee is the scalar array [3]; element type is the scalar. - Instruction* arr3TypeInst = defUseMgr->GetDef(arr3TypeId); - const uint32_t scalarTypeId = arr3TypeInst->GetSingleWordInOperand(0); - const uint32_t ptrScalarTypeId = typeMgr->FindPointerToType( - scalarTypeId, spv::StorageClass::Workgroup); - - for (uint32_t c = 0; c < 3; ++c) { - const uint32_t constCId = constMgr->GetUIntConstId(c); - // Extract component c from the stored value. - const uint32_t extractId = ctx->TakeNextId(); - auto extract = MakeUnique( - ctx, spv::Op::OpCompositeExtract, scalarTypeId, extractId, - std::initializer_list{ - {SPV_OPERAND_TYPE_ID, {valId}}, - {SPV_OPERAND_TYPE_LITERAL_INTEGER, {c}}}); - InsertBefore(block, store, std::move(extract), ctx); - - // Access chain into the array at index c. - const uint32_t acId = ctx->TakeNextId(); - auto acInst = MakeUnique( - ctx, spv::Op::OpAccessChain, ptrScalarTypeId, acId, - std::initializer_list{ - {SPV_OPERAND_TYPE_ID, {ptrId}}, - {SPV_OPERAND_TYPE_ID, {constCId}}}); - InsertBefore(block, store, std::move(acInst), ctx); - - // Store the component. - auto compStore = MakeUnique( - ctx, spv::Op::OpStore, 0, 0, - std::initializer_list{ - {SPV_OPERAND_TYPE_ID, {acId}}, - {SPV_OPERAND_TYPE_ID, {extractId}}}); - // Copy memory operands (alignment etc.) from original store. - for (uint32_t opIdx = 2; opIdx < store->NumInOperands(); ++opIdx) { - compStore->AddOperand(Operand(store->GetInOperand(opIdx))); - } - InsertBefore(block, store, std::move(compStore), ctx); - } - + Instruction* valDef = defUseMgr->GetDef(valId); + BuildCompositeStore(ctx, block, store, valDef->type_id(), ptrId, valId, vec3ToArr3, store); ctx->KillInst(store); } // ===================================================================== // Phase 6: Assert on unsupported atomic/CopyMemory on vec3 pointers. // After phases 3-5, any remaining instruction whose pointer operand's - // type is ptr_Workgroup_vec3 indicates an unsupported pattern. + // type is a decomposed Workgroup pointer indicates an unsupported pattern. // ===================================================================== for (auto& func : *ctx->module()) { for (auto& bb : func) { @@ -466,8 +594,8 @@ namespace MobileGL { if (ptrDef == nullptr) { continue; } - if (ptrVec3ToPtrArr3.find(ptrDef->type_id()) != - ptrVec3ToPtrArr3.end()) { + if (decomposedPtrTypeIds.find(ptrDef->type_id()) != + decomposedPtrTypeIds.end()) { MOBILEGL_ASSERT(false, "DecomposeWorkgroupVec3Pass: unsupported atomic/CopyMemory " "on workgroup vec3 pointer");