Move benchmarking infra code to torchao (#766)

Summary:
Pull Request resolved: #766

Add a test to validate that benchmarking works on ios deive

Reviewed By: metascroy

Differential Revision: D61501686

torchao/experimental/benchmark_infra/ios/Entitlements-Dev.plist

@@ -0,0 +1,12 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>get-task-allow</key>
+	<true/>
+	<key>keychain-access-groups</key>
+	<array>
+		<string>T84QZS65DQ.platformFamily</string>
+	</array>
+</dict>
+</plist>

torchao/experimental/benchmark_infra/ios/TorchAOBenchmark-Info.plist

@@ -0,0 +1,34 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>CFBundleDevelopmentRegion</key>
+	<string>en</string>
+	<key>CFBundleDisplayName</key>
+	<string>${PRODUCT_NAME}</string>
+	<key>CFBundleExecutable</key>
+	<string>${EXECUTABLE_NAME}</string>
+	<key>CFBundleIdentifier</key>
+	<string>${FB_BUNDLE_ID}</string>
+	<key>CFBundleInfoDictionaryVersion</key>
+	<string>6.0</string>
+	<key>CFBundleName</key>
+	<string>${PRODUCT_NAME}</string>
+	<key>CFBundlePackageType</key>
+	<string>APPL</string>
+	<key>CFBundleShortVersionString</key>
+	<string>1.0</string>
+	<key>CFBundleSignature</key>
+	<string>????</string>
+	<key>CFBundleVersion</key>
+	<string>1.0</string>
+	<key>LSRequiresIPhoneOS</key>
+	<true/>
+	<key>UILaunchStoryboardName</key>
+	<string>LaunchScreen</string>
+	<key>UISupportedInterfaceOrientations</key>
+	<array>
+		<string>UIInterfaceOrientationPortrait</string>
+	</array>
+</dict>
+</plist>

torchao/experimental/benchmark_infra/ios/main_empty.mm

@@ -0,0 +1,8 @@
+// (c) Meta Platforms, Inc. and affiliates.
+
+#include <iostream>
+
+int main(int argc, char** argv) {
+  std::cout << "Default main when no benchmarking deps are specified\n";
+  return 0;
+}

torchao/experimental/benchmark_infra/ios/output_redirect.mm

@@ -0,0 +1,44 @@
+// (c) Meta Platforms, Inc. and affiliates.
+
+#import "output_redirect.h"
+
+#import <cstdio>
+#import <fstream>
+#import <iostream>
+#import <string>
+
+#import <Foundation/Foundation.h>
+
+class STDIORedirector {
+ public:
+  STDIORedirector() {
+    if (@available(iOS 17, *)) {
+      /* duplicate stdout */
+      std::string file_name =
+          std::string(std::getenv("HOME")) + "/tmp/BENCH_LOG";
+      redirect_out_ = fopen(file_name.c_str(), "w");
+      stdout_dupfd_ = dup(1);
+      /* replace stdout with our output fd */
+      dup2(fileno(redirect_out_), 1);
+      fflush(stdout);
+      setvbuf(stdout, nil, _IONBF, 0);
+      setvbuf(redirect_out_, nil, _IONBF, 0);
+    }
+  }
+
+  ~STDIORedirector() {
+    if (@available(iOS 17, *)) {
+      fflush(stdout);
+      /* restore stdout */
+      dup2(stdout_dupfd_, 1);
+      close(stdout_dupfd_);
+      fclose(redirect_out_);
+    }
+  }
+
+ private:
+  FILE* redirect_out_;
+  int stdout_dupfd_;
+};
+
+static STDIORedirector stdio_redirector_;

torchao/experimental/benchmark_infra/test/test_bench.cpp

@@ -0,0 +1,49 @@
+// (c) Meta Platforms, Inc. and affiliates.
+
+#include <benchmark/benchmark.h>
+
+#include <cstdint>
+#include <cstdlib>
+#include <memory>
+#include <random>
+
+namespace {
+std::random_device rd;
+std::mt19937 generator(rd());
+
+// This test is to validate that the benchmarking binary
+// can be run on any device. Right now, it is focusing
+// on benchmarking on laptop (x86 oe mac) and iOS
+static void TestBenchmark(benchmark::State& state) {
+  const int32_t K = state.range(0);
+  auto a = std::make_unique<float[]>(K);
+  auto b = std::make_unique<float[]>(K);
+  auto c = std::make_unique<float[]>(K);
+  static std::uniform_real_distribution<> real_distrib(-1.0, 1.0);
+  for (int ii = 0; ii < K; ++ii) {
+    a[ii] = real_distrib(generator);
+    b[ii] = real_distrib(generator);
+    c[ii] = 0;
+  }
+  for (auto _ : state) {
+    for (int ii = 0; ii < K; ++ii) {
+      c[ii] = a[ii] + b[ii];
+    }
+  }
+}
+
+BENCHMARK(TestBenchmark)->Args({4096 * 4})->UseRealTime();
+} // namespace
+
+int main(int argc, char** argv) {
+  char arg0_default[] = "benchmark";
+  char* args_default = arg0_default;
+  if (!argv) {
+    argc = 1;
+    argv = &args_default;
+  }
+  ::benchmark::Initialize(&argc, argv);
+  ::benchmark::RunSpecifiedBenchmarks();
+  ::benchmark::Shutdown();
+  return 0;
+}