From f65d756bff8d4c5cbc15bd0322a2ae8e6b896a21 Mon Sep 17 00:00:00 2001 From: Ellie Hermaszewska Date: Tue, 29 Oct 2024 14:49:26 +0800 Subject: format * format * Minor test fixes * enable checking cpp format in ci --- source/slang/slang-ir-specialize-address-space.cpp | 561 +++++++++++---------- 1 file changed, 286 insertions(+), 275 deletions(-) (limited to 'source/slang/slang-ir-specialize-address-space.cpp') diff --git a/source/slang/slang-ir-specialize-address-space.cpp b/source/slang/slang-ir-specialize-address-space.cpp index 836601a00..a3c45617d 100644 --- a/source/slang/slang-ir-specialize-address-space.cpp +++ b/source/slang/slang-ir-specialize-address-space.cpp @@ -1,172 +1,172 @@ #include "slang-ir-specialize-address-space.h" -#include "slang-ir.h" +#include "slang-ir-clone.h" #include "slang-ir-insts.h" #include "slang-ir-util.h" -#include "slang-ir-clone.h" +#include "slang-ir.h" namespace Slang { - struct AddressSpaceContext : public AddressSpaceSpecializationContext +struct AddressSpaceContext : public AddressSpaceSpecializationContext +{ + IRModule* module; + + Dictionary mapInstToAddrSpace; + InitialAddressSpaceAssigner* addrSpaceAssigner; + + AddressSpaceContext(IRModule* inModule, InitialAddressSpaceAssigner* inAddrSpaceAssigner) + : module(inModule), addrSpaceAssigner(inAddrSpaceAssigner) { - IRModule* module; + } - Dictionary mapInstToAddrSpace; - InitialAddressSpaceAssigner* addrSpaceAssigner; + AddressSpace getAddressSpaceFromVarType(IRInst* type) + { + return addrSpaceAssigner->getAddressSpaceFromVarType(type); + } - AddressSpaceContext(IRModule* inModule, InitialAddressSpaceAssigner* inAddrSpaceAssigner) - : module(inModule) - , addrSpaceAssigner(inAddrSpaceAssigner) - { - } + AddressSpace getLeafInstAddressSpace(IRInst* inst) + { + return addrSpaceAssigner->getLeafInstAddressSpace(inst); + } - AddressSpace getAddressSpaceFromVarType(IRInst* type) - { - return addrSpaceAssigner->getAddressSpaceFromVarType(type); - } + AddressSpace getAddrSpace(IRInst* inst) override + { + auto addrSpace = mapInstToAddrSpace.tryGetValue(inst); + if (addrSpace) + return *addrSpace; + return AddressSpace::Generic; + } - AddressSpace getLeafInstAddressSpace(IRInst* inst) - { - return addrSpaceAssigner->getLeafInstAddressSpace(inst); - } + List workList; - AddressSpace getAddrSpace(IRInst* inst) override - { - auto addrSpace = mapInstToAddrSpace.tryGetValue(inst); - if (addrSpace) - return *addrSpace; - return AddressSpace::Generic; - } + struct FuncSpecializationKey + { + private: + IRFunc* func; + List argAddrSpaces; + HashCode hashCode; - List workList; - - struct FuncSpecializationKey + public: + IRFunc* getFunc() const { return func; } + ArrayView getArgAddrSpaces() const { return argAddrSpaces.getArrayView(); } + + FuncSpecializationKey() = default; + + FuncSpecializationKey(IRFunc* func, List argAddrSpaces) + : func(func), argAddrSpaces(argAddrSpaces) { - private: - IRFunc* func; - List argAddrSpaces; - HashCode hashCode; - public: - IRFunc* getFunc() const { return func; } - ArrayView getArgAddrSpaces() const { return argAddrSpaces.getArrayView(); } - - FuncSpecializationKey() = default; - - FuncSpecializationKey(IRFunc* func, List argAddrSpaces) - : func(func) - , argAddrSpaces(argAddrSpaces) + Hasher hasher; + hasher.addHash(Slang::getHashCode(func)); + for (auto addrSpace : argAddrSpaces) { - Hasher hasher; - hasher.addHash(Slang::getHashCode(func)); - for (auto addrSpace : argAddrSpaces) - { - hasher.addHash((HashCode)addrSpace); - } - hashCode = hasher.getResult(); + hasher.addHash((HashCode)addrSpace); } + hashCode = hasher.getResult(); + } - bool operator==(const FuncSpecializationKey& key) const + bool operator==(const FuncSpecializationKey& key) const + { + if (func != key.func) + return false; + if (argAddrSpaces.getCount() != key.argAddrSpaces.getCount()) + return false; + for (Index i = 0; i < argAddrSpaces.getCount(); i++) { - if (func != key.func) + if (argAddrSpaces[i] != key.argAddrSpaces[i]) return false; - if (argAddrSpaces.getCount() != key.argAddrSpaces.getCount()) - return false; - for (Index i = 0; i < argAddrSpaces.getCount(); i++) - { - if (argAddrSpaces[i] != key.argAddrSpaces[i]) - return false; - } - return true; } + return true; + } - HashCode getHashCode() const - { - return hashCode; - } - }; + HashCode getHashCode() const { return hashCode; } + }; - Dictionary functionSpecializations; + Dictionary functionSpecializations; - IRFunc* specializeFunc(const FuncSpecializationKey& key) - { - auto func = key.getFunc(); - IRCloneEnv cloneEnv; - IRBuilder builder(module); + IRFunc* specializeFunc(const FuncSpecializationKey& key) + { + auto func = key.getFunc(); + IRCloneEnv cloneEnv; + IRBuilder builder(module); - // First, clone the function body. - builder.setInsertBefore(func); - auto specializedFunc = as(cloneInst(&cloneEnv, &builder, func)); + // First, clone the function body. + builder.setInsertBefore(func); + auto specializedFunc = as(cloneInst(&cloneEnv, &builder, func)); - // Update the parameter types with new address spaces in the specialized function. - Index paramIndex = 0; - for (auto param : specializedFunc->getParams()) + // Update the parameter types with new address spaces in the specialized function. + Index paramIndex = 0; + for (auto param : specializedFunc->getParams()) + { + auto paramType = param->getFullType(); + auto ptrType = as(paramType); + if (ptrType) { - auto paramType = param->getFullType(); - auto ptrType = as(paramType); - if (ptrType) - { - auto paramAddrSpace = key.getArgAddrSpaces()[paramIndex]; - auto newParamType = builder.getPtrType(ptrType->getOp(), ptrType->getValueType(), paramAddrSpace); - param->setFullType(newParamType); - mapInstToAddrSpace[param] = paramAddrSpace; - } - paramIndex++; + auto paramAddrSpace = key.getArgAddrSpaces()[paramIndex]; + auto newParamType = + builder.getPtrType(ptrType->getOp(), ptrType->getValueType(), paramAddrSpace); + param->setFullType(newParamType); + mapInstToAddrSpace[param] = paramAddrSpace; } + paramIndex++; + } - // Update the function type. - fixUpFuncType(specializedFunc); + // Update the function type. + fixUpFuncType(specializedFunc); - functionSpecializations[key] = specializedFunc; - return specializedFunc; - } + functionSpecializations[key] = specializedFunc; + return specializedFunc; + } - AddressSpace getFuncResultAddrSpace(IRFunc* callee) - { - auto funcType = as(callee->getDataType()); - return getAddressSpaceFromVarType(funcType->getResultType()); - } + AddressSpace getFuncResultAddrSpace(IRFunc* callee) + { + auto funcType = as(callee->getDataType()); + return getAddressSpaceFromVarType(funcType->getResultType()); + } - // Return true if the address space of the function return type is changed. - bool processFunction(IRFunc* func) + // Return true if the address space of the function return type is changed. + bool processFunction(IRFunc* func) + { + bool retValAddrSpaceChanged = false; + Dictionary mapVarValueToAddrSpace; + bool changed = true; + while (changed) { - bool retValAddrSpaceChanged = false; - Dictionary mapVarValueToAddrSpace; - bool changed = true; - while (changed) + changed = false; + for (auto block : func->getBlocks()) { - changed = false; - for (auto block : func->getBlocks()) + bool isFirstBlock = block == func->getFirstBlock(); + + for (auto inst : block->getChildren()) { - bool isFirstBlock = block == func->getFirstBlock(); + // If we have already assigned an address space to this instruction, then skip + // it. + if (mapInstToAddrSpace.containsKey(inst)) + { + // TODO: if the inst is a phi node, we need to check if the address space of + // the phi arguments is consistent. If not, then we need to report an error. + // For now, we just skip the checks. + continue; + } - for (auto inst : block->getChildren()) + // If the inst already has a pointer type with explicit address space, then use + // it. + if (auto ptrType = as(inst->getDataType())) { - // If we have already assigned an address space to this instruction, then skip it. - if (mapInstToAddrSpace.containsKey(inst)) + if (ptrType->hasAddressSpace()) { - // TODO: if the inst is a phi node, we need to check if the address space of the phi arguments - // is consistent. If not, then we need to report an error. - // For now, we just skip the checks. + mapInstToAddrSpace[inst] = ptrType->getAddressSpace(); continue; } + } - // If the inst already has a pointer type with explicit address space, then use it. - if (auto ptrType = as(inst->getDataType())) - { - if (ptrType->hasAddressSpace()) - { - mapInstToAddrSpace[inst] = ptrType->getAddressSpace(); - continue; - } - } - - // Otherwise, try to assign an address space based on the instruction type. - switch (inst->getOp()) - { - case kIROp_Var: - case kIROp_RWStructuredBufferGetElementPtr: + // Otherwise, try to assign an address space based on the instruction type. + switch (inst->getOp()) + { + case kIROp_Var: + case kIROp_RWStructuredBufferGetElementPtr: { - // The address space of these insts should be assigned by the initial address space assigner. + // The address space of these insts should be assigned by the initial + // address space assigner. AddressSpace addrSpace = AddressSpace::Generic; if (addrSpaceAssigner->tryAssignAddressSpace(inst, addrSpace)) { @@ -175,21 +175,21 @@ namespace Slang } break; } - case kIROp_GetElementPtr: - case kIROp_FieldAddress: - case kIROp_GetOffsetPtr: - case kIROp_BitCast: - if (!mapInstToAddrSpace.containsKey(inst)) + case kIROp_GetElementPtr: + case kIROp_FieldAddress: + case kIROp_GetOffsetPtr: + case kIROp_BitCast: + if (!mapInstToAddrSpace.containsKey(inst)) + { + auto addrSpace = getAddrSpace(inst->getOperand(0)); + if (addrSpace != AddressSpace::Generic) { - auto addrSpace = getAddrSpace(inst->getOperand(0)); - if (addrSpace != AddressSpace::Generic) - { - mapInstToAddrSpace[inst] = addrSpace; - changed = true; - } + mapInstToAddrSpace[inst] = addrSpace; + changed = true; } - break; - case kIROp_Store: + } + break; + case kIROp_Store: { auto addrSpace = getAddrSpace(inst->getOperand(1)); if (addrSpace != AddressSpace::Generic) @@ -200,195 +200,206 @@ namespace Slang } } break; - case kIROp_Load: + case kIROp_Load: { - if (auto addrSpace = mapVarValueToAddrSpace.tryGetValue(inst->getOperand(0))) + if (auto addrSpace = + mapVarValueToAddrSpace.tryGetValue(inst->getOperand(0))) { mapInstToAddrSpace[inst] = *addrSpace; changed = true; } } break; - case kIROp_Param: - if (!isFirstBlock) + case kIROp_Param: + if (!isFirstBlock) + { + auto phiArgs = getPhiArgs(inst); + AddressSpace addrSpace = AddressSpace::Generic; + for (auto arg : phiArgs) { - auto phiArgs = getPhiArgs(inst); - AddressSpace addrSpace = AddressSpace::Generic; - for (auto arg : phiArgs) + auto argAddrSpace = getAddrSpace(arg); + if (argAddrSpace != AddressSpace::Generic) { - auto argAddrSpace = getAddrSpace(arg); - if (argAddrSpace != AddressSpace::Generic) + if (addrSpace != AddressSpace::Generic && + addrSpace != argAddrSpace) { - if (addrSpace != AddressSpace::Generic && addrSpace != argAddrSpace) - { - // TODO: this is an error in user code, because the address spaces of the - // phi arguments don't match. - } - addrSpace = argAddrSpace; + // TODO: this is an error in user code, because the + // address spaces of the phi arguments don't match. } + addrSpace = argAddrSpace; } - if (addrSpace != AddressSpace::Generic) - { - mapInstToAddrSpace[inst] = addrSpace; - changed = true; - } - break; + } + if (addrSpace != AddressSpace::Generic) + { + mapInstToAddrSpace[inst] = addrSpace; + changed = true; } break; - case kIROp_Call: + } + break; + case kIROp_Call: + { + auto callInst = as(inst); + auto callee = as(inst->getOperand(0)); + if (callee) { - auto callInst = as(inst); - auto callee = as(inst->getOperand(0)); - if (callee) + List argAddrSpaces; + bool fullySpecialized = true; + for (UInt i = 0; i < callInst->getArgCount(); i++) { - List argAddrSpaces; - bool fullySpecialized = true; - for (UInt i = 0; i < callInst->getArgCount(); i++) + auto arg = callInst->getArg(i); + auto argAddrSpace = getAddrSpace(arg); + argAddrSpaces.add(getAddrSpace(arg)); + if (argAddrSpace == AddressSpace::Generic && + as(arg->getDataType())) { - auto arg = callInst->getArg(i); - auto argAddrSpace = getAddrSpace(arg); - argAddrSpaces.add(getAddrSpace(arg)); - if (argAddrSpace == AddressSpace::Generic && - as(arg->getDataType())) - { - fullySpecialized = false; - break; - } - } - if (!fullySpecialized) + fullySpecialized = false; break; - - FuncSpecializationKey key(callee, argAddrSpaces); - IRFunc* specializedCallee = nullptr; - if (IRFunc** specializedFunc = functionSpecializations.tryGetValue(key)) - { - specializedCallee = *specializedFunc; - } - else - { - specializedCallee = specializeFunc(key); - workList.add(specializedCallee); - } - IRBuilder builder(callInst); - builder.setInsertBefore(callInst); - if (specializedCallee != callInst->getCallee()) - { - callInst = as(builder.replaceOperand(callInst->getOperands(), specializedCallee)); - } - auto callResultAddrSpace = getFuncResultAddrSpace(specializedCallee); - if (callResultAddrSpace != AddressSpace::Generic) - { - mapInstToAddrSpace[callInst] = callResultAddrSpace; - changed = true; } } + if (!fullySpecialized) + break; + + FuncSpecializationKey key(callee, argAddrSpaces); + IRFunc* specializedCallee = nullptr; + if (IRFunc** specializedFunc = + functionSpecializations.tryGetValue(key)) + { + specializedCallee = *specializedFunc; + } + else + { + specializedCallee = specializeFunc(key); + workList.add(specializedCallee); + } + IRBuilder builder(callInst); + builder.setInsertBefore(callInst); + if (specializedCallee != callInst->getCallee()) + { + callInst = as(builder.replaceOperand( + callInst->getOperands(), + specializedCallee)); + } + auto callResultAddrSpace = + getFuncResultAddrSpace(specializedCallee); + if (callResultAddrSpace != AddressSpace::Generic) + { + mapInstToAddrSpace[callInst] = callResultAddrSpace; + changed = true; + } } - break; - case kIROp_Return: + } + break; + case kIROp_Return: + { + auto retVal = inst->getOperand(0); + auto addrSpace = getAddrSpace(retVal); + if (addrSpace != AddressSpace::Generic) { - auto retVal = inst->getOperand(0); - auto addrSpace = getAddrSpace(retVal); - if (addrSpace != AddressSpace::Generic) + auto funcType = as(func->getDataType()); + AddressSpace resultAddrSpace = getFuncResultAddrSpace(func); + if (resultAddrSpace != addrSpace) { - auto funcType = as(func->getDataType()); - AddressSpace resultAddrSpace = getFuncResultAddrSpace(func); - if (resultAddrSpace != addrSpace) - { - auto ptrResultType = as(funcType->getResultType()); - SLANG_ASSERT(ptrResultType); - IRBuilder builder(func); - auto newResultType = builder.getPtrType(ptrResultType->getOp(), ptrResultType->getValueType(), addrSpace); - fixUpFuncType(func, newResultType); - retValAddrSpaceChanged = true; - } + auto ptrResultType = + as(funcType->getResultType()); + SLANG_ASSERT(ptrResultType); + IRBuilder builder(func); + auto newResultType = builder.getPtrType( + ptrResultType->getOp(), + ptrResultType->getValueType(), + addrSpace); + fixUpFuncType(func, newResultType); + retValAddrSpaceChanged = true; } } - break; } + break; } } } - return retValAddrSpaceChanged; } + return retValAddrSpaceChanged; + } - static void setDataType(IRInst* inst, IRType* dataType) + static void setDataType(IRInst* inst, IRType* dataType) + { + auto rate = inst->getRate(); + if (!rate) { - auto rate = inst->getRate(); - if (!rate) - { - inst->setFullType(dataType); - return; - } - - IRBuilder builder(inst); - builder.setInsertBefore(inst); - auto newType = builder.getRateQualifiedType(rate, dataType); - inst->setFullType(newType); + inst->setFullType(dataType); + return; } - void applyAddressSpaceToInstType() + IRBuilder builder(inst); + builder.setInsertBefore(inst); + auto newType = builder.getRateQualifiedType(rate, dataType); + inst->setFullType(newType); + } + + void applyAddressSpaceToInstType() + { + for (auto [inst, addrSpace] : mapInstToAddrSpace) { - for (auto [inst, addrSpace] : mapInstToAddrSpace) + auto ptrType = as(inst->getDataType()); + if (ptrType) { - auto ptrType = as(inst->getDataType()); - if (ptrType) + if (ptrType->getAddressSpace() != addrSpace) { - if (ptrType->getAddressSpace() != addrSpace) - { - IRBuilder builder(inst); - auto newType = builder.getPtrType(ptrType->getOp(), ptrType->getValueType(), addrSpace); - setDataType(inst, newType); - } + IRBuilder builder(inst); + auto newType = + builder.getPtrType(ptrType->getOp(), ptrType->getValueType(), addrSpace); + setDataType(inst, newType); } } } + } - void processModule() + void processModule() + { + for (auto globalInst : module->getGlobalInsts()) { - for (auto globalInst : module->getGlobalInsts()) + auto addrSpace = getLeafInstAddressSpace(globalInst); + if (addrSpace != AddressSpace::Generic) { - auto addrSpace = getLeafInstAddressSpace(globalInst); - if (addrSpace != AddressSpace::Generic) - { - mapInstToAddrSpace[globalInst] = addrSpace; - } - if (auto func = as(globalInst)) - { - if (func->findDecoration()) - workList.add(func); - } + mapInstToAddrSpace[globalInst] = addrSpace; + } + if (auto func = as(globalInst)) + { + if (func->findDecoration()) + workList.add(func); } + } - HashSet newWorkList; - while (workList.getCount()) + HashSet newWorkList; + while (workList.getCount()) + { + for (Index i = 0; i < workList.getCount(); i++) { - for (Index i = 0; i < workList.getCount(); i++) + auto func = workList[i]; + bool resultTypeChanged = processFunction(func); + if (resultTypeChanged) { - auto func = workList[i]; - bool resultTypeChanged = processFunction(func); - if (resultTypeChanged) + for (auto use = func->firstUse; use; use = use->nextUse) { - for (auto use = func->firstUse; use; use = use->nextUse) + if (auto callInst = as(use->getUser())) { - if (auto callInst = as(use->getUser())) - { - newWorkList.add(getParentFunc(callInst)); - } + newWorkList.add(getParentFunc(callInst)); } } } - workList.clear(); - for (auto f : newWorkList) - workList.add(f); } - - applyAddressSpaceToInstType(); + workList.clear(); + for (auto f : newWorkList) + workList.add(f); } - }; - void specializeAddressSpace(IRModule* module, InitialAddressSpaceAssigner* addrSpaceAssigner) - { - AddressSpaceContext context(module, addrSpaceAssigner); - context.processModule(); + applyAddressSpaceToInstType(); } +}; + +void specializeAddressSpace(IRModule* module, InitialAddressSpaceAssigner* addrSpaceAssigner) +{ + AddressSpaceContext context(module, addrSpaceAssigner); + context.processModule(); } +} // namespace Slang -- cgit v1.2.3