[DynamicCompilation] Add _AnyTensorHandle and Swift-C round trip entry points (#19529)

Author: rxwei

lib/SILOptimizer/Mandatory/TFPartition.cpp

@@ -94,10 +94,10 @@ static bool isUserIgnoredByPartitioning(SILInstruction *inst) {
   return isa<RefCountingInst>(inst);
 }
 
-/// Given a decl for a struct that has a single field (typically because it is
-/// known to be a standard library type like Int or Float), return the canonical
-/// type of the single member, asserting and aborting if we get something
-/// unexpected.
+/// Given a decl for a struct or class that has a single field (typically
+/// because it is known to be a standard library type like Int or Float), return
+/// the canonical type of the single member, asserting and aborting if we get
+/// something unexpected.
 static CanType getSingleElementDeclFieldType(NominalTypeDecl *decl) {
   auto *field = tf::getFieldIfContainsSingleField(decl);
   assert(field && "Struct should have one member");
@@ -3777,7 +3777,13 @@ void TFFunctionPartition::insertTensorComputationStartEndTerminate(
   auto tensorHandleDecl = ctx.getTensorHandleDecl();
   assert(getSingleElementDeclFieldType(tensorHandleDecl) &&
          "TensorHandle should have exactly one field");
-  auto tensorHandleMember = *tensorHandleDecl->getStoredProperties().begin();
+  auto *anyTensorHandleClass =
+      tensorHandleDecl->getSuperclass()->getAnyNominal();
+  auto anyTensorHandleSILTy = SILType::getPrimitiveObjectType(
+      anyTensorHandleClass->getDeclaredType()->getCanonicalType());
+  assert(anyTensorHandleClass);
+  auto *tensorHandleMember =
+      *anyTensorHandleClass->getStoredProperties().begin();
 
   // Ownership markers for CTensorHandle accesses.
   auto loadOwnership = hostFn.hasQualifiedOwnership()
@@ -3839,6 +3845,8 @@ void TFFunctionPartition::insertTensorComputationStartEndTerminate(
     // it.  If it is a scalar, then we need to box the scalar in a
     // CTensorHandle.
     if (isTensorHandle(tensorValue->getType().getASTType())) {
+      // Upcast to _AnyTensorHandle.
+      tensorValue = B.createUpcast(loc, tensorValue, anyTensorHandleSILTy);
       auto fieldAddress =
           B.createRefElementAddr(loc, tensorValue, tensorHandleMember);
       tensorValue = B.createLoad(loc, fieldAddress, loadOwnership);
@@ -3977,8 +3985,9 @@ void TFFunctionPartition::insertTensorComputationStartEndTerminate(
                                   /*objc*/ false, /*canAllocOnStack*/ false,
                                   /*elementTypes*/ {},
                                   /*elementCountOperands*/ {});
+    auto baseTH = B.createUpcast(loc, newTH, anyTensorHandleSILTy);
     auto fieldAddress =
-        B.createRefElementAddr(result.getLoc(), newTH, tensorHandleMember);
+        B.createRefElementAddr(result.getLoc(), baseTH, tensorHandleMember);
 
     B.createStore(result.getLoc(), newValue, fieldAddress, storeOwnership);
 

lib/SILOptimizer/Mandatory/TFUtilities.cpp

@@ -96,17 +96,26 @@ llvm::raw_ostream *tf::getTFDumpIntermediateStream() {
   return &fileStream;
 }
 
-/// If the specified decl has a single stored field, return it.  Otherwise
-/// return null.
+/// Given a nominal type decl, collect all fields. If it's a class decl, collect
+/// all fields along the inheritance hierarchy.
+static void getAllFields(NominalTypeDecl *decl,
+                         SmallVectorImpl<VarDecl *> &fields) {
+  for (auto *field : decl->getStoredProperties())
+    fields.push_back(field);
+  if (auto *classdecl = decl->getAsClassOrClassExtensionContext())
+    if (auto *superclass = classdecl->getSuperclassDecl())
+      getAllFields(superclass, fields);
+}
+
+/// If the specified decl has a single stored field, return it.  If it's a class
+/// type, return there's exactly one field in the entire inheritance hierarchy.
+/// Otherwise return null.
 VarDecl *tf::getFieldIfContainsSingleField(NominalTypeDecl *decl) {
-  // Check to see if there is a single stored field.
-  auto fieldIt = decl->getStoredProperties().begin();
-  if (fieldIt == decl->getStoredProperties().end())
-    return nullptr;
-  auto result = *fieldIt++;
-  if (fieldIt != decl->getStoredProperties().end())
-    return nullptr;
-  return result;
+  SmallVector<VarDecl *, 4> fields;
+  getAllFields(decl, fields);
+  if (fields.size() == 1)
+    return fields.front();
+  return nullptr;
 }
 
 bool tf::isTensorHandle(SILType ty) {

stdlib/public/TensorFlow/CompilerRuntime.swift

@@ -495,9 +495,8 @@ internal func dumpCTensorHandleContent(
   let dType: TF_DataType = TFE_TensorHandleDataType(inputTensorHandle)
   debugLog("Tensor \(idx) has TF data type \(dType).")
   switch dType {
-  case TF_INT8: dumpTensorContent(inputTensorHandle, Int8.self)
   case TF_UINT8: dumpTensorContent(inputTensorHandle, UInt8.self)
-  case TF_INT16: dumpTensorContent(inputTensorHandle, Int16.self)
+  case TF_INT8: dumpTensorContent(inputTensorHandle, Int8.self)
   case TF_UINT16: dumpTensorContent(inputTensorHandle, UInt16.self)
   case TF_INT16: dumpTensorContent(inputTensorHandle, Int16.self)
   case TF_UINT32: dumpTensorContent(inputTensorHandle, UInt32.self)
@@ -1132,7 +1131,7 @@ public func _TFCTerminateTensorComputation(_ computation: _TensorComputation) {
 ///   function.
 @inlinable
 @_silgen_name("_swift_tfc_CreateCTensorHandle")
-public func _TFCCreateCTensorHandle<T>(_ value : T,
+public func _TFCCreateCTensorHandle<T>(_ value: T,
                                        _ dtype: TF_DataType) -> CTensorHandle {
   // Create a new CTensor and initialize it to the scalar value.
   let tensor = TF_AllocateTensor(dtype, nil, 0, MemoryLayout<T>.stride)
@@ -1170,6 +1169,37 @@ public func _TFCExtractCTensorHandle(
   return handle.cTensorHandle
 }
 
+@inlinable
+@_silgen_name("_swift_tfc_GetCTensorHandleFromSwift")
+public func _TFCGetCTensorHandleFromSwift(
+  _ handle: _AnyTensorHandle
+) -> CTensorHandle {
+  return handle.cTensorHandle
+}
+
+@inlinable
+@_silgen_name("_swift_tfc_CreateTensorHandleFromC")
+public func _TFCCreateTensorHandleFromC(
+  _ cHandle: CTensorHandle
+) -> _AnyTensorHandle {
+  let dtype = TFE_TensorHandleDataType(cHandle)
+  switch dtype {
+  case TF_BFLOAT16: return TensorHandle<BFloat16>(owning: cHandle)
+  case TF_UINT8: return TensorHandle<UInt8>(owning: cHandle)
+  case TF_INT8: return TensorHandle<Int8>(owning: cHandle)
+  case TF_UINT16: return TensorHandle<UInt16>(owning: cHandle)
+  case TF_INT16: return TensorHandle<Int16>(owning: cHandle)
+  case TF_UINT32: return TensorHandle<UInt32>(owning: cHandle)
+  case TF_INT32: return TensorHandle<Int32>(owning: cHandle)
+  case TF_UINT64: return TensorHandle<UInt64>(owning: cHandle)
+  case TF_INT64: return TensorHandle<Int64>(owning: cHandle)
+  case TF_FLOAT: return TensorHandle<Float>(owning: cHandle)
+  case TF_DOUBLE: return TensorHandle<Double>(owning: cHandle)
+  case TF_BOOL: return TensorHandle<Bool>(owning: cHandle)
+  default: fatalError("Unsupported dtype \(dtype)")
+  }
+}
+
 @inlinable
 @_silgen_name("_swift_tfc_CreateFloatTensorHandleFromCTensorHandle")
 public func _TFCCreateTensorHandleFromCTensorHandle(

stdlib/public/TensorFlow/DataTypes.swift

@@ -21,6 +21,7 @@
 
 import CTensorFlow
 
+@_fixed_layout
 public struct _TensorDataType {
   internal var cDataType: TF_DataType
 

stdlib/public/TensorFlow/TensorHandle.swift

@@ -16,34 +16,52 @@
 
 import CTensorFlow
 
-/// `TensorHandle` is the type used by ops and the `#tfop()` syntax
-/// specifically. It includes a `Scalar` type, which compiler internals depend
-/// on to determine the datatypes of parameters when they are extracted
-/// into a tensor program.
+/// `_AnyTensorHandle` is the scalar-agnostic base type for `TensorHandle`, used
+/// specifically for low-level, type-erased passings of Swift-level tensor
+/// handles in the compiler.
 @_fixed_layout // required because the compiler accesses cTensorHandle directly.
-public final class TensorHandle<Scalar : AccelerableByTensorFlow> {
+public class _AnyTensorHandle {
   /// The underlying `TF_TensorHandle *`.
   ///
-  /// - Note: The compiler knows that `TensorHandle` has a single stored
+  /// - Note: The compiler knows that `_AnyTensorHandle` has a single stored
   /// property, and assumes that this is it. Changing the design of
   /// `TensorHandle` will require tweaking the compiler.
-  public let cTensorHandle: CTensorHandle
+  @usableFromInline let cTensorHandle: CTensorHandle
+  
+  /// Private initializer from a `CTensorHandle`. Should only be called from
+  /// `TensorHandle<Scalar>.init`.
+  fileprivate init(base: CTensorHandle) {
+    self.cTensorHandle = base
+  }
+}
 
+/// `TensorHandle` is the type used by ops and the `#tfop()` syntax
+/// specifically. It includes a `Scalar` type, which compiler internals depend
+/// on to determine the datatypes of parameters when they are extracted
+/// into a tensor program.
+@_fixed_layout // required because the compiler accesses cTensorHandle directly.
+public final class TensorHandle<Scalar> : _AnyTensorHandle
+  where Scalar : AccelerableByTensorFlow {
+  @usableFromInline
+  init(owning cTensorHandle: CTensorHandle) {
+    super.init(base: cTensorHandle)
+  }
+  
   @usableFromInline
-  init(copyingFromCTensor cTensor: CTensor) {
+  convenience init(copyingFromCTensor cTensor: CTensor) {
     let status = TF_NewStatus()
     let cTensorHandle = TFE_NewTensorHandle(cTensor, status)
     checkOk(status)
-    self.cTensorHandle = cTensorHandle!
-
+    self.init(owning: cTensorHandle!)
     TF_DeleteStatus(status)
   }
 
-  @usableFromInline
-  init(owning cTensorHandle: CTensorHandle) {
-    self.cTensorHandle = cTensorHandle
+  deinit {
+    debugLog("De-initializing TensorHandle.")
+    TFE_DeleteTensorHandle(cTensorHandle)
+    debugLog("Returning from deinit of TensorHandle.")
   }
-
+  
   /// Create a `TensorHandle` with a closure that initializes the underlying
   /// buffer.
   ///
@@ -71,12 +89,6 @@ public final class TensorHandle<Scalar : AccelerableByTensorFlow> {
     self.init(copyingFromCTensor: cTensor)
     TF_DeleteTensor(cTensor)
   }
-
-  deinit {
-    debugLog("De-initializing TensorHandle.")
-    TFE_DeleteTensorHandle(cTensorHandle)
-    debugLog("Returning from deinit of TensorHandle.")
-  }
 }
 
 internal extension TensorHandle {

test/TensorFlowRuntime/dynamic_compilation.swift

@@ -6,6 +6,7 @@
 // This file contains testing over a dataset as a global variable. This requires
 // sends/recvs support for variant handles.
 
+import CTensorFlow
 import TensorFlow
 import TensorFlowUnittest
 import StdlibUnittest

test/TensorFlowRuntime/runtime_entry_points.swift

@@ -0,0 +1,28 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+// REQUIRES: swift_test_mode_optimize
+
+import CTensorFlow
+import TensorFlow
+import TensorFlowUnittest
+import StdlibUnittest
+
+var RuntimeEntryPointTests = TestSuite("RuntimeEntryPoint")
+
+RuntimeEntryPointTests.testCPUOrGPU("RoundTrip_CTensorHandle_AnyTensorHandle") {
+  let zero: TensorHandle<Float> =
+    #tfop("Const", dtype: Float.self, value$tensor: 0.0)
+  var cHandle = _TFCGetCTensorHandleFromSwift(zero as _AnyTensorHandle)
+  let status = TF_NewStatus()
+  // We must do a copy, i.e. a retain on the tensor handle, to make sure it won't
+  // get double-free'd when both `zero` and `anyHandle` below go out of scope.
+  cHandle = TFE_TensorHandleCopySharingTensor(cHandle, status)
+  expectEqual(TF_GetCode(status), TF_OK)
+  TF_DeleteStatus(status)
+  let anyHandle = _TFCCreateTensorHandleFromC(cHandle)
+  let tensor = Tensor(handle: anyHandle as! TensorHandle<Float>)
+  print(tensor)
+  expectTrue(tensor == Tensor(0.0))
+}
+
+runAllTests()