Extended graph_op IRGen to support calling C APIs directly.

In this PR, we call the TF C API TFE_RunConstOp() instead of the previous
compiler runtime entry point @_silgen_name("_swift_tfc_RunEagerConstTest"), and
then call compiler runtime entry point
@_silgen_name("_swift_tfc_CreateFloatTensorHandleFromCTensorHandle")to wrap it
into a TensorHandle<Float>.

The main changes are:

1. To create an llvm::Function object based on a function name such as
"_swift_tfc_GetGlobalEagerContext", we first call silModule.findFunction() to
get the SILFunction, and then call IGM.getAddrOfSILFunction() to get the
llvm::Function.

To make this work, we fixed a bug in SILDeserializer::readSILFunctionChecked(),
which pervious a function decl (not body) from being deserialized in the IRGen stage.

2. We obtain from llvm::Function objects the LLVM type objects for C data types
TFE_Context* and TFE_TensorHandle*, and then generate bitcast instructions to
put the function params into the right types, before issuing the relevant
function call.

Next steps:
1. Replace the TFE_RunConstOp() call with a sequence of TF eager C API calls.
2. Generalize the graph_op decoding logic to handle graph_op's other than
Const.
3. Support generic tf datatype T instead of the hard-coded Float.
4. Figure out a way to call do scalar promotion even in the case of -Onone,
since Tensor<Float>(1.0) becomes a pseudo graph_op "tfc.scalarToTensor", which
gets should be turn into a "Const" graph_op.

Author: Mingsheng Hong

lib/IRGen/IRGenSIL.cpp

@@ -34,7 +34,11 @@
 #include "swift/Basic/STLExtras.h"
 #include "swift/SIL/Dominance.h"
 #include "swift/SIL/InstructionUtils.h"
+// SWIFT_ENABLE_TENSORFLOW
+#include "swift/SIL/GraphOperationInfo.h"
 #include "swift/SIL/PrettyStackTrace.h"
+// SWIFT_ENABLE_TENSORFLOW
+#include "swift/SIL/SILConstants.h"
 #include "swift/SIL/SILDebugScope.h"
 #include "swift/SIL/SILDeclRef.h"
 #include "swift/SIL/SILLinkage.h"
@@ -81,6 +85,7 @@
 
 using namespace swift;
 using namespace irgen;
+using swift::tf::GraphOperationInfo;
 
 // FIXME: Remove this option entirely and turn this on by default.
 llvm::cl::opt<bool> DebugInfoInlinedGenerics(
@@ -1875,9 +1880,6 @@ void IRGenSILFunction::visitGraphOperationInst(GraphOperationInst *i) {
   if (!llvm::TFDynamicCompilation)
     llvm_unreachable("graph_op is not valid in canonical SIL");
 
-  auto &llvmModule = IGM.Module;
-  auto &llvmContext = llvmModule.getContext();
-
   tf::GraphOperationInfo opInfo(i);
   SmallVector<tf::GraphOperationInfo::StructuredOperand, 4> structuredOperands;
   auto opName = opInfo.decodeName(structuredOperands);
@@ -1895,20 +1897,64 @@ void IRGenSILFunction::visitGraphOperationInst(GraphOperationInst *i) {
   // 2. Run the graph_op
   // 3. Set the output tensor handles via setLoweredExplosion()
 
+  auto &silModule = CurSILFn->getModule();
+
   // The true return type is TFE_Context*, which is an opaque pointer, so it
-  // maps to void* in the Swift-C calling convention.
-  auto getContextFn = llvmModule.getOrInsertFunction(
-      "_swift_tfc_GetGlobalEagerContext",
-      llvm::TypeBuilder<void *(), /*cross_compilable=*/false>::get(
-          llvmContext));
+  // maps to void* in the Swift-C calling convention. `eagerContext` has type
+  // void*, or i8* in LLVM type system.
+  auto getContextSilFn = silModule.findFunction(
+      "_swift_tfc_GetGlobalEagerContext", SILLinkage::PublicExternal);
+  assert(getContextSilFn);
+  llvm::Constant *getContextFn =
+      IGM.getAddrOfSILFunction(getContextSilFn, NotForDefinition);
+  assert(getContextFn);
   auto eagerContext = Builder.CreateCall(getContextFn, {});
 
-  // The true function type is TFE_Context* -> TFE_TensorHandle*.
-  auto testFunc = llvmModule.getOrInsertFunction(
-      "_swift_tfc_RunEagerConstTest",
-      llvm::TypeBuilder<void *(void *), false>::get(llvmContext));
-  auto tensorHandle = Builder.CreateCall(testFunc, {eagerContext});
-
+  // For now we call a hard-coded C API to run a const op:
+  //   TFE_TensorHandle* TFE_RunConstOp(TFE_Context* ctx)
+  // TODO: Remove this hard-coded C API call.
+  LLVM_DEBUG(llvm::dbgs() << "IRGen for TFE_RunConstOp().\n");
+  auto TFERunConstSilFn =
+      silModule.findFunction("TFE_RunConstOp", SILLinkage::PublicExternal);
+  assert(TFERunConstSilFn);
+  llvm::Function *TFERunConstFn =
+      IGM.getAddrOfSILFunction(TFERunConstSilFn, NotForDefinition);
+  assert(TFERunConstFn);
+
+  // We need to cast `eagerContext` of type i8* to %struct.TFE_Context*
+  auto *funcTy = TFERunConstFn->getFunctionType();
+  assert(funcTy->getNumParams() == 1);
+  auto *tfeContextTy = funcTy->getParamType(0);
+  LLVM_DEBUG(llvm::dbgs() << "  Param 0 of TFE_RunConstOp() has type "
+                          << *tfeContextTy << ".\n");
+  auto eagerContextTyped = Builder.CreateBitCast(eagerContext, tfeContextTy);
+
+  LLVM_DEBUG(llvm::dbgs() << "  Creating call over TFE_RunConstOp().\n");
+  auto cTensorHandle = Builder.CreateCall(TFERunConstFn, {eagerContextTyped});
+
+  // Wrap `cTensorHandle` into a TensorHandle<T> object.
+  // This requires casting `cTensorHandle` of i8* type to
+  // %struct.TFE_TensorHandle*.
+  LLVM_DEBUG(llvm::dbgs() << "IRGen for creating result TensorHandle.\n");
+  auto createHandleSilFn = silModule.findFunction(
+      "_swift_tfc_CreateFloatTensorHandleFromCTensorHandle",
+      SILLinkage::PublicExternal);
+  assert(createHandleSilFn);
+  llvm::Function *createHandleFn =
+      IGM.getAddrOfSILFunction(createHandleSilFn, NotForDefinition);
+  assert(createHandleFn);
+  auto *createHandleFnTy = createHandleFn->getFunctionType();
+  assert(createHandleFnTy->getNumParams() == 1);
+  auto *cTensorHandleTy = createHandleFnTy->getParamType(0);
+  LLVM_DEBUG(llvm::dbgs() << "  Param 0 of tensor handle creation fn has type "
+                          << *cTensorHandleTy << ".\n");
+  auto cTensorHandleTyped =
+      Builder.CreateBitCast(cTensorHandle, cTensorHandleTy);
+  LLVM_DEBUG(llvm::dbgs() << "  Creating call over tensor handle creation.\n");
+  auto tensorHandle = Builder.CreateCall(createHandleFn, {cTensorHandleTyped});
+
+  LLVM_DEBUG(
+      llvm::dbgs() << "Done with IRGen for graph_op; setting explosion.\n");
   Explosion e;
   e.add(tensorHandle);
 

lib/Serialization/DeserializeSIL.cpp

@@ -427,8 +427,9 @@ SILDeserializer::readSILFunctionChecked(DeclID FID, SILFunction *existingFn,
     break;
 
   case SILStage::Lowered:
-    llvm_unreachable("cannot deserialize into a module that has entered "
-                     "Lowered stage");
+    if (!declarationOnly) // SWIFT_ENABLE_TENSORFLOW
+      llvm_unreachable("cannot deserialize into a module that has entered "
+                       "Lowered stage");
   }
 
   if (FID == 0)

stdlib/public/TensorFlow/CompilerRuntime.swift

@@ -1166,3 +1166,20 @@ public func _GetGlobalEagerContext() -> CTFEContext {
   return _ExecutionContext.global.eagerContext
 }
 
+// TODO: replace these functions with the generic versions commented out below.
+@inlinable
+@_silgen_name("_swift_tfc_ExtractFloatCTensorHandle")
+public func _ExtractCTensorHandle(
+  _ handle: TensorHandle<Float>
+) -> CTensorHandle {
+  return handle.cTensorHandle
+}
+
+@inlinable
+@_silgen_name("_swift_tfc_CreateFloatTensorHandleFromCTensorHandle")
+public func _CreateTensorHandleFromCTensorHandle(
+  _ ownedCHandle: CTensorHandle
+) -> TensorHandle<Float> {
+  return TensorHandle<Float>(owning: ownedCHandle)
+}
+