Add a deviceScopes stack to CompilerRuntime to be consulted by the eager op dispatch. (#20897)

Unfortunately, this info gets ignored when any ops get extracted out into a
graph, but hopefully static analysis of these calls could remove some of those differences.

Author: pschuh

include/swift/Runtime/RuntimeFunctions.def

@@ -1505,6 +1505,9 @@ FUNCTION(TFC_OpSetAttrOptionalTensorShapeArray,
          ARGS(Int8PtrTy, Int8PtrTy, Int8PtrTy, Int8PtrTy),
          ATTRS(NoUnwind))
 
+FUNCTION(TFC_OpSetDeviceFromScope, _swift_tfc_OpSetDeviceFromScope,
+         C_CC, RETURNS(VoidTy), ARGS(Int8PtrTy, Int8PtrTy), ATTRS(NoUnwind))
+
 // The swift function signature is (CTFEOp, UnsafePointer<Int8>, String) -> ().
 // In LLVM IR, the String gets exploded to (BridgeObjectPtrTy, Int64Ty), so
 // IRGen passes those two values in place of the Swift String.

lib/IRGen/IRGenSIL.cpp

@@ -2843,15 +2843,8 @@ void IRGenSILFunction::visitGraphOperationInst(GraphOperationInst *i) {
   }
 
   // Set the device.
-  opDevice = deviceInfo->handleDevicePlacement(opInfo.getOperationName(),
-                                               opDevice, i->getAttributes());
-  assert(!opDevice.empty());
-  if (opDevice != TF_ALL_DEVICES) {
-    auto *setDeviceFn = IGM.getTFE_OpSetDeviceFn();
-    auto device = IGM.getAddrOfGlobalString(opDevice);
-    Builder.CreateCall(setDeviceFn, {op, device, status});
-    checkOk(status);
-  }
+  Builder.CreateCall(IGM.getTFC_OpSetDeviceFromScopeFn(), {op, status});
+  checkOk(status);
 
   // If we have any opaque TensorGroup results, then we need to do extra runtime
   // work to determine the number of TensorFlow outputs and to allocate space

stdlib/private/TensorFlowUnittest/TensorFlowUnittest.swift

@@ -53,10 +53,13 @@ public func printT(_ message: String) {
 }
 
 extension TestSuite {
+  static let willTargetGPU: Bool =
+    CommandLine.arguments.contains("--target=gpu")
   // If the macro TPU is not specified, run the tests with CPU or GPU. Otherwise
   // use TPU.
   // For GPU execution, TensorFlow must have been built with --config=cuda,
   // and a gpu device must be available.
+  @inline(never)
   public func testAllBackends(_ name: String, _ body: @escaping () -> Void) {
 #if !TPU
     testCPUOrGPU(name, body)
@@ -66,12 +69,14 @@ extension TestSuite {
   }
   public func testCPUOrGPU(_ name: String, _ body: @escaping () -> Void) {
 #if !TPU
-    test(name + "_CPU_OR_GPU") {
+    test(name + (TestSuite.willTargetGPU ? "_GPU" : "_CPU")) {
       _RuntimeConfig.executionMode = .auto
       _RuntimeConfig.usesTFEagerAPI = true
       _RuntimeConfig.gpuMemoryAllowGrowth = true
       _RuntimeConfig.printsDebugLog = false
-      body()
+      withDevice(TestSuite.willTargetGPU ? .gpu : .cpu) {
+        body()
+      }
     }
 #endif // TPU
   }

stdlib/public/TensorFlow/CompilerRuntime.swift

@@ -261,6 +261,10 @@ public final class _ExecutionContext {
   /// The mutex for preventing potential concurrent access.
   private var mutex: pthread_mutex_t = pthread_mutex_t()
 
+  /// The stack for holding the current device scoping information.
+  @usableFromInline
+  var deviceScopes: [DeviceKind] = [.`default`]
+
   /// Initializes a new execution context by initializing available devices.
   @usableFromInline
   init() {
@@ -1507,6 +1511,67 @@ func _TFCOpSetAttrStringArray(_ op: CTFEOp,
   }
 }
 
+/// A TensorFlow device kind.
+public enum DeviceKind {
+  /// Default device.
+  case `default`
+  /// Central processing units.
+  case cpu
+  /// Graphics processing units.
+  case gpu
+  // TODO: TPU?
+}
+
+internal extension _ExecutionContext {
+  var currentDeviceKind: DeviceKind {
+    return sync {
+      deviceScopes.last ?? .default
+    }
+  }
+
+  @usableFromInline
+  func pushDeviceScope(_ kind: DeviceKind) {
+    sync {
+      deviceScopes.append(kind)
+    }
+  }
+
+  @usableFromInline
+  func popDeviceScope() {
+    sync {
+      internalConsistencyCheck(deviceScopes.popLast() != nil)
+    }
+  }
+}
+
+/// Executes a closure, making TensorFlow operations run on a specific kind of
+/// device unless specified otherwise.
+///
+/// - Parameters:
+///   - kind: A kind of device to run TensorFlow operations on.
+///   - body: A closure whose TensorFlow operations are to be executed on the
+///     specified kind of device.
+@inlinable
+public func withDevice<R>(_ kind: DeviceKind,
+                          perform body: () throws -> R) rethrows -> R {
+  _ExecutionContext.global.pushDeviceScope(kind)
+  let result = try body()
+  _ExecutionContext.global.popDeviceScope()
+  return result
+}
+
+@usableFromInline
+@_cdecl("_swift_tfc_OpSetDeviceFromScope")
+func _TFCOpSetDeviceFromScope(_ op: CTFEOp, _ status: CTFStatus) {
+  let context = _ExecutionContext.global
+  let device = context.currentDeviceKind
+  switch device {
+  case .cpu: TFE_OpSetDevice(op, context.cpuDeviceName, status)
+  case .gpu: TFE_OpSetDevice(op, context.gpuDeviceName!, status)
+  default: break
+  }
+}
+
 @usableFromInline
 @_cdecl("_swift_tfc_CheckOk")
 func _TFCCheckOk(_ s: CTFStatus) {

test/TensorFlowRuntime/accelerator_only.swift

@@ -1,4 +1,4 @@
-// RUN: %target-run-dynamic-compilation-swift
+// RUN: %target-run-dynamic-compilation-swift %swift-tensorflow-test-run-extra-options
 // RUN: %target-swift-frontend -emit-sil -O %s -verify | %FileCheck %s
 // REQUIRES: executable_test
 // REQUIRES: swift_test_mode_optimize

test/TensorFlowRuntime/codable.swift

@@ -1,4 +1,4 @@
-// RUN: %target-run-simple-swift
+// RUN: %target-run-simple-swift %swift-tensorflow-test-run-extra-options
 // REQUIRES: executable_test
 // REQUIRES: swift_test_mode_optimize
 // REQUIRES: objc_interop

test/TensorFlowRuntime/conflicting_module.swift

@@ -7,7 +7,7 @@
 // RUN: %target-build-swift %S/Inputs/ConflictingModule1.swift -emit-module -emit-library -module-name ConflictingModule1 -o %t/libConflictingModule1.%target-dylib-extension -emit-module-path %t/ConflictingModule1.swiftmodule
 // RUN: %target-build-swift %S/Inputs/ConflictingModule2.swift -emit-module -emit-library -module-name ConflictingModule2 -o %t/libConflictingModule2.%target-dylib-extension -emit-module-path %t/ConflictingModule2.swiftmodule
 // RUN: %target-build-swift %s -I %t -L %t -lConflictingModule1 -lConflictingModule2 -Xlinker -rpath -Xlinker %t -o %t/a.out
-// RUN: %target-run %t/a.out
+// RUN: %target-run %t/a.out %swift-tensorflow-test-run-extra-options
 // REQUIRES: executable_test
 
 import StdlibUnittest

test/TensorFlowRuntime/control_flow_1.swift

@@ -1,5 +1,5 @@
-// RUN: %target-run-simple-swift
-// RUN: %target-run-dynamic-compilation-swift
+// RUN: %target-run-simple-swift %swift-tensorflow-test-run-extra-options
+// RUN: %target-run-dynamic-compilation-swift %swift-tensorflow-test-run-extra-options
 // REQUIRES: executable_test
 // REQUIRES: swift_test_mode_optimize
 //

test/TensorFlowRuntime/control_flow_objc.swift

@@ -1,4 +1,4 @@
-// RUN: %target-run-simple-swift
+// RUN: %target-run-simple-swift %swift-tensorflow-test-run-extra-options
 // REQUIRES: executable_test
 // REQUIRES: swift_test_mode_optimize
 // REQUIRES: OS=macosx

test/TensorFlowRuntime/dataset_1.swift

@@ -1,5 +1,5 @@
-// RUN: %target-run-simple-swift
-// RUN: %target-run-dynamic-compilation-swift
+// RUN: %target-run-simple-swift %swift-tensorflow-test-run-extra-options
+// RUN: %target-run-dynamic-compilation-swift %swift-tensorflow-test-run-extra-options
 // REQUIRES: executable_test
 // REQUIRES: swift_test_mode_optimize
 //
@@ -100,7 +100,6 @@ DatasetTests.testAllBackends("Basic") {
   // OneShotIterator is not supported on GPU.
   model()
 }
-#endif
 
 DatasetTests.testAllBackends("MultiValue") {
   enableCPU()
@@ -137,5 +136,6 @@ DatasetTests.testAllBackends("MultiValue") {
   expectEqual(2, Tensor(handle: next.0).scalarized())
   expectEqual(12, Tensor(handle: next.1).scalarized())
 }
+#endif
 
 runAllTests()

test/TensorFlowRuntime/dataset_api.swift

@@ -1,5 +1,5 @@
-// RUN: %target-run-simple-swift
-// RUN: %target-run-dynamic-compilation-swift
+// RUN: %target-run-simple-swift %swift-tensorflow-test-run-extra-options
+// RUN: %target-run-dynamic-compilation-swift %swift-tensorflow-test-run-extra-options
 // REQUIRES: executable_test
 // REQUIRES: swift_test_mode_optimize
 //

test/TensorFlowRuntime/dataset_global.swift

@@ -1,6 +1,6 @@
 // TODO: Revert to %target-run-simple-swift once we complete send/recv support for resource/variant tensors.
-// RUN: %target-run-send-recv-handle-swift
-// RUN: %target-run-dynamic-compilation-swift
+// RUN: %target-run-send-recv-handle-swift %swift-tensorflow-test-run-extra-options
+// RUN: %target-run-dynamic-compilation-swift %swift-tensorflow-test-run-extra-options
 // REQUIRES: executable_test
 // REQUIRES: swift_test_mode_optimize
 
@@ -19,15 +19,14 @@ var DatasetGlobalTests = TestSuite("DatasetGlobal")
 // Fatal error: No unary variant device copy function found for direction: 1 and Variant type_name: tensorflow::DatasetVariantWrapper
 #if !CUDA
 
-let scalars = Tensor<Float>([0, 1, 2])
-let dataset = Dataset(elements: scalars)
-var iterator = dataset.makeIterator()
-
 DatasetGlobalTests.testCPUOrGPU("RuntimeConfigTest") {
   expectTrue(_RuntimeConfig.usesTFEagerAPI)
 }
 
 DatasetGlobalTests.testCPUOrGPU("DatasetAsGlobalVar") {
+  let scalars = Tensor<Float>([0, 1, 2])
+  let dataset = Dataset(elements: scalars)
+
   var expectedVal: Float = 0.0
   for item in dataset {
     _hostOp(item)
@@ -37,6 +36,10 @@ DatasetGlobalTests.testCPUOrGPU("DatasetAsGlobalVar") {
 }
 
 DatasetGlobalTests.testCPUOrGPU("IteratorAsGlobalVar") {
+  let scalars = Tensor<Float>([0, 1, 2])
+  let dataset = Dataset(elements: scalars)
+  var iterator = dataset.makeIterator()
+
   var expectedVal: Float = 0.0
 	while let item = iterator.next() {
     _hostOp(item)