From 303e67012ad10bcf21356810643cc2b3c38bfa29 Mon Sep 17 00:00:00 2001 From: Jay Kwak <82421531+jkwak-work@users.noreply.github.com> Date: Thu, 13 Feb 2025 22:24:51 -0800 Subject: [PATCH 1/5] Deduplicate witness table based on type and operand-0 This commit is to experiment the idea of what would happen if we deduplicate the witness table based only on its type and its first operand. It turned out that it causes a crash, because it doesn't carry the decorations and childs properly. It means we need to also consider them as well as its type and its first operand. But if we do, we cannot deduplicate anything. It seems that this is not the right place to apply the deduplication. --- source/slang/slang-ir-autodiff.cpp | 11 +++++++ source/slang/slang-ir-specialize.cpp | 48 ++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+) diff --git a/source/slang/slang-ir-autodiff.cpp b/source/slang/slang-ir-autodiff.cpp index 9075002e0d..a22699c8f1 100644 --- a/source/slang/slang-ir-autodiff.cpp +++ b/source/slang/slang-ir-autodiff.cpp @@ -1634,6 +1634,17 @@ void DifferentiableTypeConformanceContext::buildGlobalWitnessDictionary() { if (auto pairType = as(globalInst)) { +#if 0 + if (auto vecType = as(pairType->getValueType())) + { + auto op0 = vecType->getOperand(0); + if (op0->getOp() == kIROp_FloatType) + { + int a = 0; + ++a; + } + } +#endif addTypeToDictionary(pairType->getValueType(), pairType->getWitness()); } diff --git a/source/slang/slang-ir-specialize.cpp b/source/slang/slang-ir-specialize.cpp index a9b0d44121..e6fb8048da 100644 --- a/source/slang/slang-ir-specialize.cpp +++ b/source/slang/slang-ir-specialize.cpp @@ -54,6 +54,36 @@ struct SpecializationContext SpecializationOptions options; bool changed = false; + // Keep the record of cloned IRs and avoid duplication. + struct WitnessTableKey + { + IRWitnessTableType* witnessTableType; + IRInst* concreteType; + IRInstListBase decorationsAndChildren; + + WitnessTableKey(IRWitnessTable* wt) + : witnessTableType(as(wt->getFullType())) + , concreteType(wt->getOperand(0)) + , decorationsAndChildren(wt->getDecorationsAndChildren()) + {} + + bool operator==(const WitnessTableKey &other) const + { + if (witnessTableType != other.witnessTableType) + return false; + if (concreteType != other.concreteType) + return false; + if (decorationsAndChildren.first != other.decorationsAndChildren.first) + return false; + return true; + } + HashCode getHashCode() const + { + return combineHash(HashCode(witnessTableType), HashCode(concreteType)); + } + }; + Dictionary mapClonedWitnessTable; + SpecializationContext(IRModule* inModule, TargetProgram* target, SpecializationOptions options) : workList(*inModule->getContainerPool().getList()) @@ -3098,6 +3128,24 @@ IRInst* specializeGenericImpl( // IRInst* clonedInst = cloneInst(&env, builder, ii); + // Avoid duplicated witness tables. + // + if (auto clonedWitness = as(clonedInst)) + { + IRInst* cachedInst; + if (context->mapClonedWitnessTable.tryGetValue(clonedWitness, cachedInst)) + { + builder->setInsertBefore(clonedInst); + clonedInst->replaceUsesWith(cachedInst); + clonedInst->removeAndDeallocate(); + clonedInst = cachedInst; + } + else + { + context->mapClonedWitnessTable.add(clonedWitness, clonedInst); + } + } + // Any new instructions we create during cloning were // not present when we initially built our work list, // so we need to make sure to consider them now. From 4aa5a1427cdc42c352806b8a358214ecb757d8e1 Mon Sep 17 00:00:00 2001 From: Jay Kwak <82421531+jkwak-work@users.noreply.github.com> Date: Fri, 14 Feb 2025 14:09:53 -0800 Subject: [PATCH 2/5] Swap the cloned IR properly --- source/slang/slang-ir-autodiff.cpp | 11 ----------- source/slang/slang-ir-specialize.cpp | 21 +++++++-------------- 2 files changed, 7 insertions(+), 25 deletions(-) diff --git a/source/slang/slang-ir-autodiff.cpp b/source/slang/slang-ir-autodiff.cpp index a22699c8f1..9075002e0d 100644 --- a/source/slang/slang-ir-autodiff.cpp +++ b/source/slang/slang-ir-autodiff.cpp @@ -1634,17 +1634,6 @@ void DifferentiableTypeConformanceContext::buildGlobalWitnessDictionary() { if (auto pairType = as(globalInst)) { -#if 0 - if (auto vecType = as(pairType->getValueType())) - { - auto op0 = vecType->getOperand(0); - if (op0->getOp() == kIROp_FloatType) - { - int a = 0; - ++a; - } - } -#endif addTypeToDictionary(pairType->getValueType(), pairType->getWitness()); } diff --git a/source/slang/slang-ir-specialize.cpp b/source/slang/slang-ir-specialize.cpp index e6fb8048da..31bafdd129 100644 --- a/source/slang/slang-ir-specialize.cpp +++ b/source/slang/slang-ir-specialize.cpp @@ -59,24 +59,18 @@ struct SpecializationContext { IRWitnessTableType* witnessTableType; IRInst* concreteType; - IRInstListBase decorationsAndChildren; WitnessTableKey(IRWitnessTable* wt) : witnessTableType(as(wt->getFullType())) , concreteType(wt->getOperand(0)) - , decorationsAndChildren(wt->getDecorationsAndChildren()) - {} + { + } - bool operator==(const WitnessTableKey &other) const + bool operator==(const WitnessTableKey& other) const { - if (witnessTableType != other.witnessTableType) - return false; - if (concreteType != other.concreteType) - return false; - if (decorationsAndChildren.first != other.decorationsAndChildren.first) - return false; - return true; - } + return witnessTableType == other.witnessTableType && concreteType == other.concreteType; + } + HashCode getHashCode() const { return combineHash(HashCode(witnessTableType), HashCode(concreteType)); @@ -3135,8 +3129,7 @@ IRInst* specializeGenericImpl( IRInst* cachedInst; if (context->mapClonedWitnessTable.tryGetValue(clonedWitness, cachedInst)) { - builder->setInsertBefore(clonedInst); - clonedInst->replaceUsesWith(cachedInst); + env.mapOldValToNew[ii] = cachedInst; clonedInst->removeAndDeallocate(); clonedInst = cachedInst; } From 4734d570060d108ebe4b27de95de0db7ad2791c6 Mon Sep 17 00:00:00 2001 From: Jay Kwak <82421531+jkwak-work@users.noreply.github.com> Date: Fri, 14 Feb 2025 17:18:50 -0800 Subject: [PATCH 3/5] Avoid clonning instances that will be deallocated --- source/slang/slang-ir-specialize.cpp | 70 ++++++++++++++++++++-------- 1 file changed, 50 insertions(+), 20 deletions(-) diff --git a/source/slang/slang-ir-specialize.cpp b/source/slang/slang-ir-specialize.cpp index 31bafdd129..b6340ee7f5 100644 --- a/source/slang/slang-ir-specialize.cpp +++ b/source/slang/slang-ir-specialize.cpp @@ -57,12 +57,11 @@ struct SpecializationContext // Keep the record of cloned IRs and avoid duplication. struct WitnessTableKey { - IRWitnessTableType* witnessTableType; - IRInst* concreteType; + IRType* witnessTableType; + IRType* concreteType; - WitnessTableKey(IRWitnessTable* wt) - : witnessTableType(as(wt->getFullType())) - , concreteType(wt->getOperand(0)) + WitnessTableKey(IRType* wtt, IRType* ct) + : witnessTableType(wtt), concreteType(ct) { } @@ -3117,27 +3116,58 @@ IRInst* specializeGenericImpl( return specializedVal; } - // For any instruction other than a `return`, we will - // simply clone it completely into the global scope. - // - IRInst* clonedInst = cloneInst(&env, builder, ii); + IRInst* clonedInst; - // Avoid duplicated witness tables. + // Deduplicate witness tables. + // As long as its type and the concrete type are same, they shouldn't be cloned. // - if (auto clonedWitness = as(clonedInst)) + if (auto witnessTable = as(ii)) { - IRInst* cachedInst; - if (context->mapClonedWitnessTable.tryGetValue(clonedWitness, cachedInst)) - { - env.mapOldValToNew[ii] = cachedInst; - clonedInst->removeAndDeallocate(); - clonedInst = cachedInst; - } - else + auto witnessTableType = witnessTable->getFullType(); + + // We need to apply specialization from `env` to `concreteType`. + auto concreteType = witnessTable->getConcreteType(); + IRType* clonedConcreteType = as(cloneInst(&env, builder, concreteType)); + + SpecializationContext::WitnessTableKey cacheKey( + witnessTableType, + clonedConcreteType); + + if (!context->mapClonedWitnessTable.tryGetValue(cacheKey, clonedInst)) { - context->mapClonedWitnessTable.add(clonedWitness, clonedInst); + // Not found from cache + // Need to create a new one + IRInst* newOperands[] = {clonedConcreteType}; + + clonedInst = builder->emitIntrinsicInst( + witnessTableType, + kIROp_WitnessTable, + 1, + newOperands); + + clonedInst->sourceLoc = ii->sourceLoc; + + if (clonedInst != ii) + { + // TODO: Decoration shouldn't matter but it casues a trouble if we don't + // clone them here. It is unclear why. + cloneInstDecorationsAndChildren(&env, builder->getModule(), ii, clonedInst); + } + + context->mapClonedWitnessTable.add(cacheKey, clonedInst); } + + env.mapOldValToNew[ii] = clonedInst; } + else + { + // For any instruction other than a `return`, we will + // simply clone it completely into the global scope. + // + clonedInst = cloneInst(&env, builder, ii); + } + + SLANG_ASSERT(clonedInst); // Any new instructions we create during cloning were // not present when we initially built our work list, From 5e68cb39d2f57b6caf212c73c13260f6190a21db Mon Sep 17 00:00:00 2001 From: Jay Kwak <82421531+jkwak-work@users.noreply.github.com> Date: Fri, 14 Feb 2025 17:30:08 -0800 Subject: [PATCH 4/5] Clone the WitenessTableType as well to be safe --- source/slang/slang-ir-specialize.cpp | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/source/slang/slang-ir-specialize.cpp b/source/slang/slang-ir-specialize.cpp index b6340ee7f5..22811ef875 100644 --- a/source/slang/slang-ir-specialize.cpp +++ b/source/slang/slang-ir-specialize.cpp @@ -3123,27 +3123,26 @@ IRInst* specializeGenericImpl( // if (auto witnessTable = as(ii)) { + // CLone type and concreteType to use as cache-key auto witnessTableType = witnessTable->getFullType(); + IRType* clonedWitnessTableType = + as(cloneInst(&env, builder, witnessTableType)); - // We need to apply specialization from `env` to `concreteType`. auto concreteType = witnessTable->getConcreteType(); IRType* clonedConcreteType = as(cloneInst(&env, builder, concreteType)); SpecializationContext::WitnessTableKey cacheKey( - witnessTableType, + clonedWitnessTableType, clonedConcreteType); if (!context->mapClonedWitnessTable.tryGetValue(cacheKey, clonedInst)) { - // Not found from cache - // Need to create a new one - IRInst* newOperands[] = {clonedConcreteType}; - + // Not found from cache and need to create a new one clonedInst = builder->emitIntrinsicInst( - witnessTableType, + clonedWitnessTableType, kIROp_WitnessTable, 1, - newOperands); + (IRInst* const*)&clonedConcreteType); clonedInst->sourceLoc = ii->sourceLoc; From 9dbbcdf9278147e16ecd42575b5dab70ea5b79d4 Mon Sep 17 00:00:00 2001 From: Jay Kwak <82421531+jkwak-work@users.noreply.github.com> Date: Tue, 18 Feb 2025 12:50:47 -0800 Subject: [PATCH 5/5] Changed the comments about cloning the children --- source/slang/slang-ir-specialize.cpp | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/source/slang/slang-ir-specialize.cpp b/source/slang/slang-ir-specialize.cpp index 22811ef875..c28c6010ec 100644 --- a/source/slang/slang-ir-specialize.cpp +++ b/source/slang/slang-ir-specialize.cpp @@ -3123,7 +3123,7 @@ IRInst* specializeGenericImpl( // if (auto witnessTable = as(ii)) { - // CLone type and concreteType to use as cache-key + // Clone type and concreteType to use as cache-key auto witnessTableType = witnessTable->getFullType(); IRType* clonedWitnessTableType = as(cloneInst(&env, builder, witnessTableType)); @@ -3137,7 +3137,7 @@ IRInst* specializeGenericImpl( if (!context->mapClonedWitnessTable.tryGetValue(cacheKey, clonedInst)) { - // Not found from cache and need to create a new one + // It is not found from cache and we need to create a new one. clonedInst = builder->emitIntrinsicInst( clonedWitnessTableType, kIROp_WitnessTable, @@ -3146,12 +3146,11 @@ IRInst* specializeGenericImpl( clonedInst->sourceLoc = ii->sourceLoc; - if (clonedInst != ii) - { - // TODO: Decoration shouldn't matter but it casues a trouble if we don't - // clone them here. It is unclear why. - cloneInstDecorationsAndChildren(&env, builder->getModule(), ii, clonedInst); - } + // It seems that the children are same when the witnessTableType and the + // concrete type are same. + // If it changes in the future, we will need to clone the witness-table-entries + // from the children and use them as the cacheKey as well. + cloneInstDecorationsAndChildren(&env, builder->getModule(), ii, clonedInst); context->mapClonedWitnessTable.add(cacheKey, clonedInst); }