[AutoDiff] Move dumpActivityInfo to DifferentiableActivityInfo class (#28597)

Author: bgogul

include/swift/SILOptimizer/Analysis/DifferentiableActivityAnalysis.h

@@ -131,7 +131,7 @@ class DifferentiableActivityInfo {
   SmallVector<SmallDenseSet<SILValue>, 4> usefulValueSets;
 
   /// The original function.
-  SILFunction &getFunction();
+  SILFunction &getFunction() const;
 
   /// Returns true if the given SILValue has a tangent space.
   bool hasTangentSpace(SILValue value) {
@@ -206,6 +206,14 @@ class DifferentiableActivityInfo {
   /// Returns the activity of the given value for the given `SILAutoDiffIndices`
   /// (parameter indices and result index).
   Activity getActivity(SILValue value, const SILAutoDiffIndices &indices) const;
+
+  /// Prints activity information for the `indices` of the given `value`.
+  void dump(SILValue value, const SILAutoDiffIndices &indices,
+            llvm::raw_ostream &s = llvm::dbgs()) const;
+
+  /// Prints activity information for the given `indices`.
+  void dump(SILAutoDiffIndices indices,
+            llvm::raw_ostream &s = llvm::dbgs()) const;
 };
 
 class DifferentiableActivityCollection {

lib/SILOptimizer/Analysis/DifferentiableActivityAnalysis.cpp

@@ -59,7 +59,7 @@ DifferentiableActivityInfo::DifferentiableActivityInfo(
   analyze(parent.domInfo, parent.postDomInfo);
 }
 
-SILFunction &DifferentiableActivityInfo::getFunction() {
+SILFunction &DifferentiableActivityInfo::getFunction() const {
   return parent.function;
 }
 
@@ -421,3 +421,32 @@ Activity DifferentiableActivityInfo::getActivity(
     activity |= ActivityFlags::Useful;
   return activity;
 }
+
+void DifferentiableActivityInfo::dump(SILValue value,
+                                      const SILAutoDiffIndices &indices,
+                                      llvm::raw_ostream &s) const {
+  s << '[';
+  auto activity = getActivity(value, indices);
+  switch (activity.toRaw()) {
+  case 0: s << "NONE"; break;
+  case (unsigned)ActivityFlags::Varied: s << "VARIED"; break;
+  case (unsigned)ActivityFlags::Useful: s << "USEFUL"; break;
+  case (unsigned)ActivityFlags::Active: s << "ACTIVE"; break;
+  }
+  s << "] " << value;
+}
+
+void DifferentiableActivityInfo::dump(SILAutoDiffIndices indices,
+                                      llvm::raw_ostream &s) const {
+  SILFunction &fn = getFunction();
+  s << "Activity info for " << fn.getName() << " at " << indices << '\n';
+  for (auto &bb : fn) {
+    s << "bb" << bb.getDebugID() << ":\n";
+    for (auto *arg : bb.getArguments())
+      dump(arg, indices, s);
+    for (auto &inst : bb)
+      for (auto res : inst.getResults())
+        dump(res, indices, s);
+    s << '\n';
+  }
+}

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -263,36 +263,6 @@ class DifferentiationTransformer {
 
 } // end anonymous namespace
 
-static void dumpActivityInfo(SILValue value,
-                             const SILAutoDiffIndices &indices,
-                             const DifferentiableActivityInfo &activityInfo,
-                             llvm::raw_ostream &s = llvm::dbgs()) {
-  s << '[';
-  auto activity = activityInfo.getActivity(value, indices);
-  switch (activity.toRaw()) {
-  case 0: s << "NONE"; break;
-  case (unsigned)ActivityFlags::Varied: s << "VARIED"; break;
-  case (unsigned)ActivityFlags::Useful: s << "USEFUL"; break;
-  case (unsigned)ActivityFlags::Active: s << "ACTIVE"; break;
-  }
-  s << "] " << value;
-}
-
-static void dumpActivityInfo(SILFunction &fn, SILAutoDiffIndices indices,
-                             const DifferentiableActivityInfo &activityInfo,
-                             llvm::raw_ostream &s = llvm::dbgs()) {
-  s << "Activity info for " << fn.getName() << " at " << indices << '\n';
-  for (auto &bb : fn) {
-    s << "bb" << bb.getDebugID() << ":\n";
-    for (auto *arg : bb.getArguments())
-      dumpActivityInfo(arg, indices, activityInfo, s);
-    for (auto &inst : bb)
-      for (auto res : inst.getResults())
-        dumpActivityInfo(res, indices, activityInfo, s);
-    s << '\n';
-  }
-}
-
 /// If the original function doesn't have a return, it cannot be differentiated.
 /// Returns true if error is emitted.
 static bool diagnoseNoReturn(ADContext &context, SILFunction *original,
@@ -1388,8 +1358,7 @@ class VJPEmitter final
     auto &activityInfo = activityCollection.getActivityInfo(
         vjp->getLoweredFunctionType()->getSubstGenericSignature(),
         AutoDiffDerivativeFunctionKind::VJP);
-    LLVM_DEBUG(
-        dumpActivityInfo(*original, indices, activityInfo, getADDebugStream()));
+    LLVM_DEBUG(activityInfo.dump(indices, getADDebugStream()));
     return activityInfo;
   }
 
@@ -2141,8 +2110,7 @@ class JVPEmitter final
     auto &activityInfo = activityCollection.getActivityInfo(
         jvp->getLoweredFunctionType()->getSubstGenericSignature(),
         AutoDiffDerivativeFunctionKind::JVP);
-    LLVM_DEBUG(
-        dumpActivityInfo(*original, indices, activityInfo, getADDebugStream()));
+    LLVM_DEBUG(activityInfo.dump(indices, getADDebugStream()));
     return activityInfo;
   }