[llvm][automemcpy] Allow distribution filtering in analysis

Differential Revision: https://reviews.llvm.org/D120037

2 files changed, 17 insertions, 2 deletions

diff --git a/libc/benchmarks/automemcpy/lib/CodeGen.cpp b/libc/benchmarks/automemcpy/lib/CodeGen.cpp
index c150ab5..2ada57a 100644
--- a/libc/benchmarks/automemcpy/lib/CodeGen.cpp
+++ b/libc/benchmarks/automemcpy/lib/CodeGen.cpp
@@ -613,7 +613,7 @@ llvm::ArrayRef<MemmoveConfiguration> getMemmoveConfigurations() {
 // Stores `VolatileStr` into a cache and returns a StringRef of the cached
 // version.
 StringRef getInternalizedString(std::string VolatileStr) {
-  static llvm::StringSet<> StringCache;
+  static llvm::StringSet StringCache;
   return StringCache.insert(std::move(VolatileStr)).first->getKey();
 }
 
diff --git a/libc/benchmarks/automemcpy/lib/ResultAnalyzerMain.cpp b/libc/benchmarks/automemcpy/lib/ResultAnalyzerMain.cpp
index 6a657e4..4a6caec 100644
--- a/libc/benchmarks/automemcpy/lib/ResultAnalyzerMain.cpp
+++ b/libc/benchmarks/automemcpy/lib/ResultAnalyzerMain.cpp
@@ -20,6 +20,12 @@ namespace llvm {
 static cl::list<std::string> InputFilenames(cl::Positional, cl::OneOrMore,
                                             cl::desc("<input json files>"));
 
+// User can filter the distributions to be taken into account.
+static cl::list<std::string>
+    KeepOnlyDistributions("keep-only-distributions", cl::ZeroOrMore,
+                          cl::desc("<comma separated list of distribution "
+                                   "names, keeps all if unspecified>"));
+
 namespace automemcpy {
 
 // This is defined in the autogenerated 'Implementations.cpp' file.
@@ -48,7 +54,7 @@ static const FunctionDescriptor &getFunctionDescriptor(StringRef FunctionName) {
 // Functions and distributions names are stored quite a few times so it's more
 // efficient to internalize these strings and refer to them through 'StringRef'.
 static StringRef getInternalizedString(StringRef VolatileStr) {
-  static llvm::StringSet<> StringCache;
+  static llvm::StringSet StringCache;
   return StringCache.insert(VolatileStr).first->getKey();
 }
 
@@ -121,6 +127,15 @@ int Main(int argc, char **argv) {
     llvm::append_range(Samples, Result.Samples);
   }
 
+  if (!KeepOnlyDistributions.empty()) {
+    llvm::StringSet ValidDistributions;
+    ValidDistributions.insert(KeepOnlyDistributions.begin(),
+                              KeepOnlyDistributions.end());
+    llvm::erase_if(Samples, [&ValidDistributions](const Sample &S) {
+      return !ValidDistributions.contains(S.Id.Distribution.Name);
+    });
+  }
+
   // Extracts median of throughputs.
   std::vector<FunctionData> Functions = getThroughputs(Samples);
   fillScores(Functions);