From 3e83532c65aee42ec41efbc393bd6f4621883acf Mon Sep 17 00:00:00 2001
From: Michelle Casbon <texasmichelle@users.noreply.github.com>
Date: Tue, 11 Aug 2020 11:50:45 -0700
Subject: [PATCH] Add platform arguments to benchmarks (#654)

* Add platform argument to benchmarks

* Lint

* Make device selection more explicit

* Add default platform
---
 BenchmarksCore/BenchmarkArguments.swift | 37 ++++++++++++++++--
 BenchmarksCore/BenchmarkSettings.swift  | 52 ++++++++++++++++++++-----
 2 files changed, 76 insertions(+), 13 deletions(-)

diff --git a/BenchmarksCore/BenchmarkArguments.swift b/BenchmarksCore/BenchmarkArguments.swift
index 6248ea9242c..3ddad9963ee 100644
--- a/BenchmarksCore/BenchmarkArguments.swift
+++ b/BenchmarksCore/BenchmarkArguments.swift
@@ -28,23 +28,38 @@ public struct BenchmarkArguments: ParsableArguments {
   @Flag(help: "Use X10 backend.")
   var x10: Bool = false
 
+  @Flag(help: "Use CPU platform.")
+  var cpu: Bool = false
+
+  @Flag(help: "Use GPU platform.")
+  var gpu: Bool = false
+
+  @Flag(help: "Use TPU platform.")
+  var tpu: Bool = false
+
   @Flag(help: "Use synthetic data.")
-  var synthetic: Bool = false 
+  var synthetic: Bool = false
 
   @Flag(help: "Use real data.")
-  var real: Bool = false 
+  var real: Bool = false
 
   @Option(help: "File path for dataset loading.")
   var datasetFilePath: String?
 
   public init() {}
 
-  public init(arguments: Benchmark.BenchmarkArguments, batchSize: Int?, eager: Bool, x10: Bool,
-              synthetic: Bool, real: Bool, datasetFilePath: String?) {
+  public init(
+    arguments: Benchmark.BenchmarkArguments, batchSize: Int?, eager: Bool, x10: Bool,
+    cpu: Bool, gpu: Bool, tpu: Bool, synthetic: Bool, real: Bool,
+    datasetFilePath: String?
+  ) {
     self.arguments = arguments
     self.batchSize = batchSize
     self.eager = eager
     self.x10 = x10
+    self.cpu = cpu
+    self.gpu = gpu
+    self.tpu = tpu
     self.synthetic = synthetic
     self.real = real
     self.datasetFilePath = datasetFilePath
@@ -62,6 +77,11 @@ public struct BenchmarkArguments: ParsableArguments {
       throw ValidationError(
         "Can't specify both --eager and --x10 backends.")
     }
+
+    guard !(cpu && gpu) || !(cpu && tpu) || !(gpu && tpu) else {
+      throw ValidationError(
+        "Can't specify multiple platforms.")
+    }
   }
 
   public var settings: [BenchmarkSetting] {
@@ -76,6 +96,15 @@ public struct BenchmarkArguments: ParsableArguments {
     if eager {
       settings.append(Backend(.eager))
     }
+    if cpu {
+      settings.append(Platform(.cpu))
+    }
+    if gpu {
+      settings.append(Platform(.gpu))
+    }
+    if tpu {
+      settings.append(Platform(.tpu))
+    }
     if synthetic {
       settings.append(Synthetic(true))
     }
diff --git a/BenchmarksCore/BenchmarkSettings.swift b/BenchmarksCore/BenchmarkSettings.swift
index 027995efc66..8a631d5f548 100644
--- a/BenchmarksCore/BenchmarkSettings.swift
+++ b/BenchmarksCore/BenchmarkSettings.swift
@@ -47,6 +47,19 @@ public struct Backend: BenchmarkSetting {
   }
 }
 
+public struct Platform: BenchmarkSetting {
+  var value: Value
+  init(_ value: Value) {
+    self.value = value
+  }
+  public enum Value {
+    case `default`
+    case cpu
+    case gpu
+    case tpu
+  }
+}
+
 public struct DatasetFilePath: BenchmarkSetting {
   var value: String
   init(_ value: String) {
@@ -54,16 +67,16 @@ public struct DatasetFilePath: BenchmarkSetting {
   }
 }
 
-public extension BenchmarkSettings {
-  var batchSize: Int? {
+extension BenchmarkSettings {
+  public var batchSize: Int? {
     return self[BatchSize.self]?.value
   }
 
-  var length: Int? {
+  public var length: Int? {
     return self[Length.self]?.value
   }
 
-  var synthetic: Bool {
+  public var synthetic: Bool {
     if let value = self[Synthetic.self]?.value {
       return value
     } else {
@@ -71,7 +84,7 @@ public extension BenchmarkSettings {
     }
   }
 
-  var backend: Backend.Value {
+  public var backend: Backend.Value {
     if let value = self[Backend.self]?.value {
       return value
     } else {
@@ -79,19 +92,39 @@ public extension BenchmarkSettings {
     }
   }
 
-  var device: Device {
+  public var platform: Platform.Value {
+    if let value = self[Platform.self]?.value {
+      return value
+    } else {
+      fatalError("Platform setting must have a default.")
+    }
+  }
+
+  public var device: Device {
     // Note: The line is needed, or all GPU memory
     // will be exhausted on initial allocation of the model.
     // TODO: Remove the following tensor workaround when above is fixed.
     let _ = _ExecutionContext.global
 
     switch backend {
-    case .eager: return Device.defaultTFEager
-    case .x10: return Device.defaultXLA
+    case .eager:
+      switch platform {
+      case .default: return Device.defaultTFEager
+      case .cpu: return Device(kind: .CPU, ordinal: 0, backend: .TF_EAGER)
+      case .gpu: return Device(kind: .GPU, ordinal: 0, backend: .TF_EAGER)
+      case .tpu: fatalError("TFEager is unsupported on TPU.")
+      }
+    case .x10:
+      switch platform {
+      case .default: return Device.defaultXLA
+      case .cpu: return Device(kind: .CPU, ordinal: 0, backend: .XLA)
+      case .gpu: return Device(kind: .GPU, ordinal: 0, backend: .XLA)
+      case .tpu: return (Device.allDevices.filter { $0.kind == .TPU }).first!
+      }
     }
   }
 
-  var datasetFilePath: String? {
+  public var datasetFilePath: String? {
     return self[DatasetFilePath.self]?.value
   }
 }
@@ -100,6 +133,7 @@ public let defaultSettings: [BenchmarkSetting] = [
   TimeUnit(.s),
   InverseTimeUnit(.s),
   Backend(.eager),
+  Platform(.default),
   Synthetic(false),
   Columns([
     "name",