[AutoDiff] NFC: Reimplement VJPCloner using pimpl pattern.

`VJPCloner.h` is now tiny: `VJPCloner` exposes only a `bool run()` entry point.

All of the implementation is moved to `VJPCloner::Implementation` in
`VJPCloner.cpp`. Methods can be defined directly in `VJPCloner.cpp` without
separate declarations.

Author: dan-zheng

include/swift/SILOptimizer/Differentiation/PullbackCloner.h

@@ -18,29 +18,14 @@
 #ifndef SWIFT_SILOPTIMIZER_UTILS_DIFFERENTIATION_PULLBACKCLONER_H
 #define SWIFT_SILOPTIMIZER_UTILS_DIFFERENTIATION_PULLBACKCLONER_H
 
-#include "swift/SILOptimizer/Analysis/DifferentiableActivityAnalysis.h"
-#include "swift/SILOptimizer/Differentiation/AdjointValue.h"
-#include "swift/SILOptimizer/Differentiation/DifferentiationInvoker.h"
-#include "swift/SILOptimizer/Differentiation/LinearMapInfo.h"
-
-#include "swift/SIL/TypeSubstCloner.h"
-#include "llvm/ADT/DenseMap.h"
-
 namespace swift {
-
-class SILDifferentiabilityWitness;
-class SILBasicBlock;
-class SILFunction;
-class SILInstruction;
-
 namespace autodiff {
 
-class ADContext;
 class VJPCloner;
 
 /// A helper class for generating pullback functions.
 class PullbackCloner final {
-  struct Implementation;
+  class Implementation;
   Implementation &impl;
 
 public:

include/swift/SILOptimizer/Differentiation/VJPCloner.h

@@ -22,149 +22,41 @@
 #include "swift/SILOptimizer/Differentiation/DifferentiationInvoker.h"
 #include "swift/SILOptimizer/Differentiation/LinearMapInfo.h"
 
-#include "swift/SIL/TypeSubstCloner.h"
-#include "llvm/ADT/DenseMap.h"
-
 namespace swift {
-
-class SILDifferentiabilityWitness;
-class SILBasicBlock;
-class SILFunction;
-class SILInstruction;
-
 namespace autodiff {
 
 class ADContext;
 class PullbackCloner;
 
-class VJPCloner final
-    : public TypeSubstCloner<VJPCloner, SILOptFunctionBuilder> {
-  friend class PullbackCloner;
-
-private:
-  /// The global context.
-  ADContext &context;
-
-  /// The original function.
-  SILFunction *const original;
-
-  /// The differentiability witness.
-  SILDifferentiabilityWitness *const witness;
-
-  /// The VJP function.
-  SILFunction *const vjp;
-
-  /// The pullback function.
-  SILFunction *pullback;
-
-  /// The differentiation invoker.
-  DifferentiationInvoker invoker;
-
-  /// Info from activity analysis on the original function.
-  const DifferentiableActivityInfo &activityInfo;
-
-  /// The linear map info.
-  LinearMapInfo pullbackInfo;
-
-  /// Caches basic blocks whose phi arguments have been remapped (adding a
-  /// predecessor enum argument).
-  SmallPtrSet<SILBasicBlock *, 4> remappedBasicBlocks;
-
-  bool errorOccurred = false;
-
-  /// Mapping from original blocks to pullback values. Used to build pullback
-  /// struct instances.
-  llvm::DenseMap<SILBasicBlock *, SmallVector<SILValue, 8>> pullbackValues;
-
-  ASTContext &getASTContext() const { return vjp->getASTContext(); }
-  SILModule &getModule() const { return vjp->getModule(); }
-  const SILAutoDiffIndices getIndices() const {
-    return witness->getSILAutoDiffIndices();
-  }
-
-  static SubstitutionMap getSubstitutionMap(SILFunction *original,
-                                            SILFunction *vjp);
-
-  static const DifferentiableActivityInfo &
-  getActivityInfo(ADContext &context, SILFunction *original,
-                  SILAutoDiffIndices indices, SILFunction *vjp);
+/// A helper class for generating VJP functions.
+class VJPCloner final {
+  class Implementation;
+  Implementation &impl;
 
 public:
+  /// Creates a VJP cloner.
+  ///
+  /// The parent VJP cloner stores the original function and an empty
+  /// to-be-generated pullback function.
   explicit VJPCloner(ADContext &context, SILFunction *original,
                      SILDifferentiabilityWitness *witness, SILFunction *vjp,
                      DifferentiationInvoker invoker);
-
-  SILFunction *createEmptyPullback();
-
-  /// Run VJP generation. Returns true on error.
+  ~VJPCloner();
+
+  ADContext &getContext() const;
+  SILModule &getModule() const;
+  SILFunction &getOriginal() const;
+  SILFunction &getVJP() const;
+  SILFunction &getPullback() const;
+  SILDifferentiabilityWitness *getWitness() const;
+  const SILAutoDiffIndices getIndices() const;
+  DifferentiationInvoker getInvoker() const;
+  LinearMapInfo &getPullbackInfo() const;
+  const DifferentiableActivityInfo &getActivityInfo() const;
+
+  /// Performs VJP generation on the empty VJP function. Returns true if any
+  /// error occurs.
   bool run();
-
-  void postProcess(SILInstruction *orig, SILInstruction *cloned);
-
-  /// Remap original basic blocks, adding predecessor enum arguments.
-  SILBasicBlock *remapBasicBlock(SILBasicBlock *bb);
-
-  /// General visitor for all instructions. If any error is emitted by previous
-  /// visits, bail out.
-  void visit(SILInstruction *inst);
-
-  void visitSILInstruction(SILInstruction *inst);
-
-private:
-  /// Get the lowered SIL type of the given AST type.
-  SILType getLoweredType(Type type);
-
-  /// Get the lowered SIL type of the given nominal type declaration.
-  SILType getNominalDeclLoweredType(NominalTypeDecl *nominal);
-
-  // Creates a trampoline block for given original terminator instruction, the
-  // pullback struct value for its parent block, and a successor basic block.
-  //
-  // The trampoline block has the same arguments as and branches to the remapped
-  // successor block, but drops the last predecessor enum argument.
-  //
-  // Used for cloning branching terminator instructions with specific
-  // requirements on successor block arguments, where an additional predecessor
-  // enum argument is not acceptable.
-  SILBasicBlock *createTrampolineBasicBlock(TermInst *termInst,
-                                            StructInst *pbStructVal,
-                                            SILBasicBlock *succBB);
-
-  /// Build a pullback struct value for the given original terminator
-  /// instruction.
-  StructInst *buildPullbackValueStructValue(TermInst *termInst);
-
-  /// Build a predecessor enum instance using the given builder for the given
-  /// original predecessor/successor blocks and pullback struct value.
-  EnumInst *buildPredecessorEnumValue(SILBuilder &builder,
-                                      SILBasicBlock *predBB,
-                                      SILBasicBlock *succBB,
-                                      SILValue pbStructVal);
-
-public:
-  void visitReturnInst(ReturnInst *ri);
-
-  void visitBranchInst(BranchInst *bi);
-
-  void visitCondBranchInst(CondBranchInst *cbi);
-
-  void visitSwitchEnumInstBase(SwitchEnumInstBase *inst);
-
-  void visitSwitchEnumInst(SwitchEnumInst *sei);
-
-  void visitSwitchEnumAddrInst(SwitchEnumAddrInst *seai);
-
-  void visitCheckedCastBranchInst(CheckedCastBranchInst *ccbi);
-
-  void visitCheckedCastValueBranchInst(CheckedCastValueBranchInst *ccvbi);
-
-  void visitCheckedCastAddrBranchInst(CheckedCastAddrBranchInst *ccabi);
-
-  // If an `apply` has active results or active inout arguments, replace it
-  // with an `apply` of its VJP.
-  void visitApplyInst(ApplyInst *ai);
-
-  void visitDifferentiableFunctionInst(DifferentiableFunctionInst *dfi);
 };
 
 } // end namespace autodiff

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -18,7 +18,11 @@
 #define DEBUG_TYPE "differentiation"
 
 #include "swift/SILOptimizer/Differentiation/PullbackCloner.h"
+#include "swift/SILOptimizer/Analysis/DifferentiableActivityAnalysis.h"
 #include "swift/SILOptimizer/Differentiation/ADContext.h"
+#include "swift/SILOptimizer/Differentiation/AdjointValue.h"
+#include "swift/SILOptimizer/Differentiation/DifferentiationInvoker.h"
+#include "swift/SILOptimizer/Differentiation/LinearMapInfo.h"
 #include "swift/SILOptimizer/Differentiation/Thunk.h"
 #include "swift/SILOptimizer/Differentiation/VJPCloner.h"
 
@@ -27,8 +31,10 @@
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/SIL/InstructionUtils.h"
 #include "swift/SIL/Projection.h"
+#include "swift/SIL/TypeSubstCloner.h"
 #include "swift/SILOptimizer/PassManager/PrettyStackTrace.h"
 #include "swift/SILOptimizer/Utils/SILOptFunctionBuilder.h"
+#include "llvm/ADT/DenseMap.h"
 
 namespace swift {
 
@@ -49,8 +55,12 @@ class VJPCloner;
 /// instructions per basic block in reverse order. This visitation order is
 /// necessary for generating pullback functions, whose control flow graph is
 /// ~a transposed version of the original function's control flow graph.
-struct PullbackCloner::Implementation final
+class PullbackCloner::Implementation final
     : public SILInstructionVisitor<PullbackCloner::Implementation> {
+
+public:
+  explicit Implementation(VJPCloner &vjpCloner);
+
 private:
   /// The parent VJP cloner.
   VJPCloner &vjpCloner;
@@ -124,23 +134,21 @@ struct PullbackCloner::Implementation final
 
   bool errorOccurred = false;
 
-  ADContext &getContext() const { return vjpCloner.context; }
+  ADContext &getContext() const { return vjpCloner.getContext(); }
   SILModule &getModule() const { return getContext().getModule(); }
   ASTContext &getASTContext() const { return getPullback().getASTContext(); }
-  SILFunction &getOriginal() const { return *vjpCloner.original; }
-  SILFunction &getPullback() const { return *vjpCloner.pullback; }
-  SILDifferentiabilityWitness *getWitness() const { return vjpCloner.witness; }
-  DifferentiationInvoker getInvoker() const { return vjpCloner.invoker; }
-  LinearMapInfo &getPullbackInfo() { return vjpCloner.pullbackInfo; }
+  SILFunction &getOriginal() const { return vjpCloner.getOriginal(); }
+  SILFunction &getPullback() const { return vjpCloner.getPullback(); }
+  SILDifferentiabilityWitness *getWitness() const {
+    return vjpCloner.getWitness();
+  }
+  DifferentiationInvoker getInvoker() const { return vjpCloner.getInvoker(); }
+  LinearMapInfo &getPullbackInfo() const { return vjpCloner.getPullbackInfo(); }
   const SILAutoDiffIndices getIndices() const { return vjpCloner.getIndices(); }
   const DifferentiableActivityInfo &getActivityInfo() const {
-    return vjpCloner.activityInfo;
+    return vjpCloner.getActivityInfo();
   }
 
-public:
-  explicit Implementation(VJPCloner &vjpCloner);
-
-private:
   //--------------------------------------------------------------------------//
   // Pullback struct mapping
   //--------------------------------------------------------------------------//
@@ -1520,9 +1528,10 @@ PullbackCloner::Implementation::Implementation(VJPCloner &vjpCloner)
   auto *domAnalysis = passManager.getAnalysis<DominanceAnalysis>();
   auto *postDomAnalysis = passManager.getAnalysis<PostDominanceAnalysis>();
   auto *postOrderAnalysis = passManager.getAnalysis<PostOrderAnalysis>();
-  domInfo = domAnalysis->get(vjpCloner.original);
-  postDomInfo = postDomAnalysis->get(vjpCloner.original);
-  postOrderInfo = postOrderAnalysis->get(vjpCloner.original);
+  auto *original = &vjpCloner.getOriginal();
+  domInfo = domAnalysis->get(original);
+  postDomInfo = postDomAnalysis->get(original);
+  postOrderInfo = postOrderAnalysis->get(original);
 }
 
 PullbackCloner::PullbackCloner(VJPCloner &vjpCloner)