[AutoDiff] Improve debugging utilities. (#32269)

- Show SIL type when printing `AdjointValue`.
- Add utilities to print `PullbackEmitter` adjoint value and buffer mappings.
- Print generated VJP before printing generated pullback.
  - This is useful because pullback generation may crash after VJP
    generation succeeds.

Author: dan-zheng

include/swift/SILOptimizer/Differentiation/AdjointValue.h

@@ -136,13 +136,12 @@ class AdjointValue final {
   void print(llvm::raw_ostream &s) const {
     switch (getKind()) {
     case AdjointValueKind::Zero:
-      s << "Zero";
+      s << "Zero[" << getType() << ']';
       break;
     case AdjointValueKind::Aggregate:
-      s << "Aggregate<";
+      s << "Aggregate[" << getType() << "](";
       if (auto *decl =
               getType().getASTType()->getStructOrBoundGenericStruct()) {
-        s << "Struct>(";
         interleave(
             llvm::zip(decl->getStoredProperties(), base->value.aggregate),
             [&s](std::tuple<VarDecl *, const AdjointValue &> elt) {
@@ -151,7 +150,6 @@ class AdjointValue final {
             },
             [&s] { s << ", "; });
       } else if (getType().is<TupleType>()) {
-        s << "Tuple>(";
         interleave(
             base->value.aggregate,
             [&s](const AdjointValue &elt) { elt.print(s); },
@@ -162,10 +160,11 @@ class AdjointValue final {
       s << ')';
       break;
     case AdjointValueKind::Concrete:
-      s << "Concrete(" << base->value.concrete << ')';
+      s << "Concrete[" << getType() << "](" << base->value.concrete << ')';
       break;
     }
   }
+
   SWIFT_DEBUG_DUMP { print(llvm::dbgs()); };
 };
 

include/swift/SILOptimizer/Differentiation/PullbackEmitter.h

@@ -307,6 +307,13 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
     return pullbackTrampolineBBMap.lookup({originalBlock, successorBlock});
   }
 
+  //--------------------------------------------------------------------------//
+  // Debugging utilities
+  //--------------------------------------------------------------------------//
+
+  void printAdjointValueMapping();
+  void printAdjointBufferMapping();
+
 public:
   //--------------------------------------------------------------------------//
   // Entry point

lib/SILOptimizer/Differentiation/PullbackEmitter.cpp

@@ -514,6 +514,66 @@ void PullbackEmitter::addToAdjointBuffer(SILBasicBlock *origBB,
   accumulateIndirect(adjointBuffer, rhsBufferAccess, loc);
 }
 
+//--------------------------------------------------------------------------//
+// Debugging utilities
+//--------------------------------------------------------------------------//
+
+void PullbackEmitter::printAdjointValueMapping() {
+  // Group original/adjoint values by basic block.
+  llvm::DenseMap<SILBasicBlock *, llvm::DenseMap<SILValue, AdjointValue>> tmp;
+  for (auto pair : valueMap) {
+    auto origPair = pair.first;
+    auto *origBB = origPair.first;
+    auto origValue = origPair.second;
+    auto adjValue = pair.second;
+    tmp[origBB].insert({origValue, adjValue});
+  }
+  // Print original/adjoint values per basic block.
+  auto &s = getADDebugStream() << "Adjoint value mapping:\n";
+  for (auto &origBB : getOriginal()) {
+    if (!pullbackBBMap.count(&origBB))
+      continue;
+    auto bbValueMap = tmp[&origBB];
+    s << "bb" << origBB.getDebugID();
+    s << " (size " << bbValueMap.size() << "):\n";
+    for (auto valuePair : bbValueMap) {
+      auto origValue = valuePair.first;
+      auto adjValue = valuePair.second;
+      s << "ORIG: " << origValue;
+      s << "ADJ: " << adjValue << '\n';
+    }
+    s << '\n';
+  }
+}
+
+void PullbackEmitter::printAdjointBufferMapping() {
+  // Group original/adjoint buffers by basic block.
+  llvm::DenseMap<SILBasicBlock *, llvm::DenseMap<SILValue, SILValue>> tmp;
+  for (auto pair : bufferMap) {
+    auto origPair = pair.first;
+    auto *origBB = origPair.first;
+    auto origBuf = origPair.second;
+    auto adjBuf = pair.second;
+    tmp[origBB][origBuf] = adjBuf;
+  }
+  // Print original/adjoint buffers per basic block.
+  auto &s = getADDebugStream() << "Adjoint buffer mapping:\n";
+  for (auto &origBB : getOriginal()) {
+    if (!pullbackBBMap.count(&origBB))
+      continue;
+    auto bbBufferMap = tmp[&origBB];
+    s << "bb" << origBB.getDebugID();
+    s << " (size " << bbBufferMap.size() << "):\n";
+    for (auto valuePair : bbBufferMap) {
+      auto origBuf = valuePair.first;
+      auto adjBuf = valuePair.second;
+      s << "ORIG: " << origBuf;
+      s << "ADJ: " << adjBuf << '\n';
+    }
+    s << '\n';
+  }
+}
+
 //--------------------------------------------------------------------------//
 // Member accessor pullback generation
 //--------------------------------------------------------------------------//

lib/SILOptimizer/Differentiation/VJPEmitter.cpp

@@ -842,15 +842,16 @@ bool VJPEmitter::run() {
   // `-enable-strip-ownership-after-serialization` is true.
   mergeBasicBlocks(vjp);
 
+  LLVM_DEBUG(getADDebugStream()
+             << "Generated VJP for " << original->getName() << ":\n"
+             << *vjp);
+
   // Generate pullback code.
   PullbackEmitter PullbackEmitter(*this);
   if (PullbackEmitter.run()) {
     errorOccurred = true;
     return true;
   }
-  LLVM_DEBUG(getADDebugStream()
-             << "Generated VJP for " << original->getName() << ":\n"
-             << *vjp);
   return errorOccurred;
 }