From 8cee05777192f75b8053d2bb8fb2f788532dc5ca Mon Sep 17 00:00:00 2001 From: Dan Zheng Date: Wed, 2 Oct 2019 00:04:01 -0700 Subject: [PATCH 1/2] [WIP] [AutoDiff] Fix AllocStack SILDebugVariable assertion failure. https://github.com/apple/swift/pull/27238 tightened debug info requirements, adding assertions to AllocStack and AllocBox constructors. This caused assertion failures for the differentiation transform, which does not provide SILDebugVariable when building AllocStack during differential/pullback generation: ``` Assertion failed: ((!dyn_cast_or_null(Loc.getAsASTNode()) || Var) && "location is a VarDecl, but SILDebugVariable is empty"), function createAllocStack, file /Users/danielzheng/swift-merge/swift/include/swift/SIL/SILBuilder.h, line 407. Stack dump: ... 1. Swift version 5.1-dev (LLVM c5340df2d1, Swift 7bc4a06c32) 2. While running pass #56204 SILModuleTransform "Differentiation". swift::SILBuilder::createAllocStack(...) (anonymous namespace)::PullbackEmitter::getAdjointBuffer(...) ``` There are a few potential fixes: - Use `RegularLocation::getAutoGeneratedLocation()` so that the assertion does not trigger. This needs to be done for all `createAllocStack` invocations in Differentiation.cpp. - Find a way to provide the correct SILDebugVariable. This seems hard because SILDebugVariable is constructed during SILGen and has a particular "ArgNo" field that seems hard to compute. This is an initial minimal test to see if tensorflow/swift-apis compilation will pass. More changes are needed for a robust final fix: - Update all `createAllocStack` calls, including those in differential generation. - Add apple/swift tests. --- lib/SILOptimizer/Mandatory/Differentiation.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/SILOptimizer/Mandatory/Differentiation.cpp b/lib/SILOptimizer/Mandatory/Differentiation.cpp index d039f10b382c8..2921cc4765896 100644 --- a/lib/SILOptimizer/Mandatory/Differentiation.cpp +++ b/lib/SILOptimizer/Mandatory/Differentiation.cpp @@ -6112,7 +6112,7 @@ class PullbackEmitter final : public SILInstructionVisitor { // Allocate local buffer and initialize to zero. auto bufObjectType = getRemappedTangentType(originalBuffer->getType()); auto *newBuf = localAllocBuilder.createAllocStack( - originalBuffer.getLoc(), bufObjectType); + RegularLocation::getAutoGeneratedLocation(), bufObjectType); // Temporarily change global builder insertion point and emit zero into the // local buffer. auto insertionPoint = builder.getInsertionBB(); From 2c6f07c5a7ec06bb180e6ebd8186f50d99c90b4e Mon Sep 17 00:00:00 2001 From: Dan Zheng Date: Wed, 2 Oct 2019 10:07:58 -0700 Subject: [PATCH 2/2] Fix differential generation for AllocStack. Use the VarInfo of the original AllocStack instruction. --- lib/SILOptimizer/Mandatory/Differentiation.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lib/SILOptimizer/Mandatory/Differentiation.cpp b/lib/SILOptimizer/Mandatory/Differentiation.cpp index 2921cc4765896..babb3e0aac2a9 100644 --- a/lib/SILOptimizer/Mandatory/Differentiation.cpp +++ b/lib/SILOptimizer/Mandatory/Differentiation.cpp @@ -4691,7 +4691,8 @@ class JVPEmitter final CLONE_AND_EMIT_TANGENT(AllocStack, asi) { auto &diffBuilder = getDifferentialBuilder(); auto *mappedAllocStackInst = diffBuilder.createAllocStack( - asi->getLoc(), getRemappedTangentType(asi->getElementType())); + asi->getLoc(), getRemappedTangentType(asi->getElementType()), + asi->getVarInfo()); bufferMap.try_emplace({asi->getParent(), asi}, mappedAllocStackInst); }