Wrapped TF C APIs as runtime functions, so that IRGen can generate calls to them (#19555)

* Wrapped TF C APIs as runtime functions, so that IRGen can generate calls to
them.

This patch "disolved" the call to the packaged experimental C API
TF_RunConstOp() into a set of finer-grained C API calls.

An alternative to the runtime function approach is to use Clang importer to call
the C APIs. Clang importer allows the C APIs to become available as SIL
functions with the same names, and thus can be found and called in IRGen via
silModule.findFunction() (see
https://github.com/apple/swift/blob/ad7def2c6bffd95d62e5e665c9faed0f8dac49f5/lib/IRGen/IRGenSIL.cpp#L1956-L1957). The
concerns with this approach are:

1. There are no precedents in IRGen that calls C APIs via this clang importer
route. Having IRGen depend on clang importer may not be desirable.

2. Some C APIs cannot be found via clang importer (e.g. TF_NewStatus() can be
found, but TFE_NewOp() cannot). As a work-around, these C APIs must be called in
the user module code first, before IRGen over the user module can locate them as
SIL functions. This feels brittle.

3. When passing objects between compiler runtime entry points (swift functions)
such as @_silgen_name("_swift_tfc_GetGlobalEagerContext") and the C APIs, IRGen
have to issue bitcast to convert between a real struct type (e.g. TFE_Context*)
and void* (i8* in LLVM type) as used in the entry points.

In contrast, the wrapped runtime functions consistently use void* in their
interfaces, so that IRGen need not do such bitcasts.

One downside with the runtime function approach is that we need to wrap ~20 TF C
APIs in the new files TensorFlow.{h,cpp}. This is however mostly a one-time eng cost.

* Reverted some compiler rt changes.

* Addressed feedback:
1. Removed runtime func impls, when we can us the C API ones.
2. Changed some compiler RT entry points to using @_cdecl, so that they can be
called in IRGen directly (instead of going through clang importer and silgen names).

* Addressed comments from @rxwei

* Added "find_package(TensorFlow REQUIRED)" to
swift/stdlib/public/runtime/CMakeLists.txt, because CI
run (e.g. https://ci-external.swift.org/view/Pull%20Request/job/swift-PR-TensorFlow-Linux/859/console)
reported:
```
CMake Error at stdlib/public/runtime/CMakeLists.txt:30 (include_directories):
  include_directories given empty-string as include directory.
```

* Moved code in TensorFlow.{h,cpp} to the CTensorFLow
 module (ctensorflow_init.{h,cpp}), so that the stdlib/public/runtime/ need not
 depend on the TF library.

Also addresed more feedback.

Author: Mingsheng Hong

include/swift/Runtime/RuntimeFunctions.def

@@ -1306,6 +1306,83 @@ FUNCTION(AutoDiffPopFromTape, swift_autodiffPopFromTape, C_CC,
          ARGS(OpaquePtrTy),
          ATTRS(NoUnwind))
 
+// SWIFT_ENABLE_TENSORFLOW
+//===----------------------------------------------------------------------===//
+// - MARK: Runtime functions issued by IRGen in dynamic compilation mode.
+// - Naming convention:
+//   - TF_XXX functions correspond to TF C APIs
+//   - TFE_XXX functions correspond to TF eager C APIs
+//   - TFC_XXX functions are C functions defined in the compiler runtime
+//===----------------------------------------------------------------------===//
+
+// Function declarations, with the definitions are in the C API impl.
+FUNCTION(TF_NewStatus, TF_NewStatus, C_CC,
+         RETURNS(Int8PtrTy),
+         ARGS(),
+         ATTRS(NoUnwind))
+
+FUNCTION(TF_DeleteStatus, TF_DeleteStatus, C_CC,
+         RETURNS(),
+         ARGS(Int8PtrTy),
+         ATTRS(NoUnwind))
+
+FUNCTION(TFE_NewOp, TFE_NewOp, C_CC,
+         RETURNS(Int8PtrTy),
+         ARGS(Int8PtrTy, Int8PtrTy, Int8PtrTy),
+         ATTRS(NoUnwind))
+
+FUNCTION(TFE_DeleteOp, TFE_DeleteOp, C_CC,
+         RETURNS(),
+         ARGS(Int8PtrTy),
+         ATTRS(NoUnwind))
+
+FUNCTION(TFE_OpSetAttrType, TFE_OpSetAttrType, C_CC,
+         RETURNS(),
+         ARGS(Int8PtrTy, Int8PtrTy, Int32Ty),
+         ATTRS(NoUnwind))
+
+FUNCTION(TFE_OpSetAttrTensor, TFE_OpSetAttrTensor, C_CC,
+         RETURNS(),
+         ARGS(Int8PtrTy, Int8PtrTy, Int8PtrTy, Int8PtrTy),
+         ATTRS(NoUnwind))
+
+FUNCTION(TF_DeleteTensor, TF_DeleteTensor, C_CC,
+         RETURNS(),
+         ARGS(Int8PtrTy),
+         ATTRS(NoUnwind))
+
+// Functions with definitions in runtime library.
+FUNCTION(TFC_CreateScalarFloatTensor, swift_tfc_CreateScalarFloatTensor, C_CC,
+         RETURNS(Int8PtrTy),
+         ARGS(Int32Ty),
+         ATTRS(NoUnwind))
+
+FUNCTION(TFE_Execute, swift_tfc_TFE_Execute, C_CC,
+         RETURNS(),
+         ARGS(Int8PtrTy, Int8PtrPtrTy, Int32PtrTy, Int8PtrTy),
+         ATTRS(NoUnwind))
+
+FUNCTION(TFC_GetGlobalEagerContext, _swift_tfc_GetGlobalEagerContext, C_CC,
+         RETURNS(Int8PtrTy),
+         ARGS(),
+         ATTRS(NoUnwind))
+
+// TODO: enable these decls once we have AnyTensorHandle.
+// FUNCTION(TFC_ExtractFloatCTensorHandle, _swift_tfc_ExtractFloatCTensorHandle, C_CC,
+//          RETURNS(Int8PtrTy),
+//          ARGS(Int8PtrTy),
+//          ATTRS(NoUnwind))
+
+// FUNCTION(TFC_CreateFloatTensorHandleFromCTensorHandle, _swift_tfc_CreateFloatTensorHandleFromCTensorHandle, C_CC,
+//          RETURNS(Int8PtrTy),
+//          ARGS(Int8PtrTy),
+//          ATTRS(NoUnwind))
+
+FUNCTION(TFC_CheckOk, _swift_tfc_CheckOk, C_CC,
+         RETURNS(),
+         ARGS(Int8PtrTy),
+         ATTRS(NoUnwind))
+
 #undef RETURNS
 #undef ARGS
 #undef ATTRS

lib/IRGen/IRGenSIL.cpp

@@ -928,7 +928,23 @@ class IRGenSILFunction :
       }
     }
   }
-  
+
+  // SWIFT_ENABLE_TENSORFLOW
+  // Returns the LLVM function with `funcName`. It must exist.
+  llvm::Function *findFunction(StringRef funcName, SILModule &silModule) {
+    LLVM_DEBUG(llvm::dbgs() << "IRGen for calling " << funcName << "().\n");
+    auto silFn = silModule.findFunction(funcName, SILLinkage::PublicExternal);
+    assert(silFn);
+    llvm::Function *fn = IGM.getAddrOfSILFunction(silFn, NotForDefinition);
+    assert(fn);
+    return fn;
+  }
+
+  void checkOk(llvm::Value *status) {
+    auto *checkOkFn = IGM.getTFC_CheckOkFn();
+    Builder.CreateCall(checkOkFn, {status});
+  }
+
   //===--------------------------------------------------------------------===//
   // SIL instruction lowering
   //===--------------------------------------------------------------------===//
@@ -1897,6 +1913,23 @@ static void abortOnGraphOp(IRGenFunction &IGF, llvm::StringRef errMessage) {
   IGF.Builder.CreateCall(abortFunc, {});
 }
 
+// Create and return a (i8*-typed) address value to a constant string.
+static llvm::Value *createStringValAddr(IRGenModule &IGM, StringRef strVal) {
+  auto &llvmModule = IGM.Module;
+  auto &llvmContext = llvmModule.getContext();
+  auto opNameVal = llvm::ConstantDataArray::getString(llvmContext, strVal);
+  auto global =
+      new llvm::GlobalVariable(IGM.Module, opNameVal->getType(), true,
+                               llvm::GlobalValue::PrivateLinkage, opNameVal);
+
+  // Make an i8*.
+  auto zero = llvm::ConstantInt::get(IGM.Int32Ty, 0);
+  llvm::Constant *indices[] = {zero, zero};
+  auto opNameValAddr = llvm::ConstantExpr::getInBoundsGetElementPtr(
+      global->getValueType(), global, indices);
+  return opNameValAddr;
+}
+
 /// For now we lower any graph_op inst into a const TF node. (super fast tensor
 /// computation, but those who don't care about correctness. :-) )
 /// TODO: Fix this mis-compilation.
@@ -1938,59 +1971,127 @@ void IRGenSILFunction::visitGraphOperationInst(GraphOperationInst *i) {
 
   auto &silModule = CurSILFn->getModule();
 
+  auto *TFNewStatusFn = IGM.getTF_NewStatusFn();
+  auto status = Builder.CreateCall(TFNewStatusFn, {});
+
+  if (opInfo.getOperationName() != "Const") {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Done with IRGen for dummy graph_op; setting explosion.\n");
+    // The added explosion is incorrect, but is good enough for unit testing.
+    Explosion e;
+    e.add(TFNewStatusFn);
+
+    SILValue result = i->getResults()[0];
+    setLoweredExplosion(result, e);
+    return;
+  }
+
   // The true return type is TFE_Context*, which is an opaque pointer, so it
   // 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 *getContextFn = IGM.getTFC_GetGlobalEagerContextFn();
   auto eagerContext = Builder.CreateCall(getContextFn, {});
 
-  // 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});
+  // Create a TFE op as in:
+  //   auto* op = TFE_NewOp(ctx, "Const", status);
+  auto *TFENewOpFn = IGM.getTFE_NewOpFn();
+  // TODO: remove the hard-coded "Const"
+  auto opNameValAddr = createStringValAddr(IGM, "Const");
+  auto op =
+      Builder.CreateCall(TFENewOpFn, {eagerContext, opNameValAddr, status});
+  checkOk(status);
+
+  // Set up dtype attr as in:
+  //   TFE_OpSetAttrType(op, "dtype", TF_FLOAT);
+  auto *setAttrTypeFn = IGM.getTFE_OpSetAttrTypeFn();
+  auto dtypeValAddr = createStringValAddr(IGM, "dtype");
+  // TODO: do not hard-code 1 as the TF_FLOAT enum value.
+  Builder.CreateCall(
+      setAttrTypeFn,
+      {op, dtypeValAddr, llvm::ConstantInt::get(IGM.Int32Ty, 1 /*TF_FLOAT*/)});
+
+  auto typeAttr = i->getAttribute(0);
+  auto typeAttrInfo =
+      GraphOperationInfo::decodeArgumentName(typeAttr.name.str());
+  assert(typeAttrInfo.first == "dtype");
+  assert(typeAttr.value.getKind() == SymbolicValue::Metatype);
+
+  auto valueAttr = i->getAttribute(1);
+  auto valueAttrInfo =
+      GraphOperationInfo::decodeArgumentName(valueAttr.name.str());
+  assert(valueAttrInfo.first == "value");
+  assert(valueAttr.value.getKind() == SymbolicValue::Float);
+  assert(valueAttrInfo.second ==
+         GraphOperationInfo::ArgumentLowering::TensorAttribute);
+  auto apfloat = valueAttr.value.getFloatValue();
+  // CreateScalarFloatTensor() takes an int instead of float, as runtime
+  // functions that take/return float values do not yet exist.
+  auto constVal = llvm::ConstantInt::get(IGM.Int32Ty, apfloat.convertToFloat());
+  LLVM_DEBUG(llvm::dbgs() << "The const value is " << *constVal << ".\n");
+
+  auto *createTensorFn = IGM.getTFC_CreateScalarFloatTensorFn();
+  auto tensor = Builder.CreateCall(createTensorFn, {constVal});
+
+  // Set up the tensor-typed value attr as in:
+  //   TFE_OpSetAttrTensor(op, "value", tensor, status);
+  auto *setTensorAttrFn = IGM.getTFE_OpSetAttrTensorFn();
+  auto valueAttrAddr = createStringValAddr(IGM, "value");
+  Builder.CreateCall(setTensorAttrFn, {op, valueAttrAddr, tensor, status});
+  checkOk(status);
+
+  auto *deleteTensorFn = IGM.getTF_DeleteTensorFn();
+  Builder.CreateCall(deleteTensorFn, {tensor});
+
+  // Now we execute the TFE op as in:
+  //   TFE_TensorHandle* retval;
+  //   int num_retvals = 1;
+  //   TFE_Execute(op, &retval, &num_retvals, status);
+  //
+  // The LLVM IR code looks like:
+  //   %returnValues = alloca %struct.TFE_TensorHandle*, align 8
+  //   %returnValueCount = alloca i32, align 8
+  //   store i32 1, i32* %returnValueCount, align 8
+  //   %134 = bitcast i8** %returnValues to %struct.TFE_TensorHandle**
+  //   call void @TFE_Execute(%struct.TFE_Op* %130,
+  //                          %struct.TFE_TensorHandle** %134,
+  //                          i32* %returnValueCount,
+  //                          %struct.TF_Status* %128)
+  //   %135 = load %struct.TFE_TensorHandle*, %struct.TFE_TensorHandle** %134,
+  //
+  // FIXME: getPointerAlignment is likely excessive. "align 4" might be
+  // sufficient.
+  auto returnValueCount =
+      createAlloca(IGM.Int32Ty, IGM.getPointerAlignment(), "returnValueCount");
+  auto expectedReturnValueCount =
+      llvm::ConstantInt::get(IGM.Int32Ty, i->getNumResults());
+  Builder.CreateStore(expectedReturnValueCount, returnValueCount);
+  auto returnValues = createAlloca(IGM.Int8PtrTy, expectedReturnValueCount,
+                                   IGM.getPointerAlignment(), "returnValues");
+  auto *tfeExecuteFn = IGM.getTFE_ExecuteFn();
+  Builder.CreateCall(tfeExecuteFn, {op, returnValues.getAddress(),
+                                    returnValueCount.getAddress(), status});
+  checkOk(status);
+
+  // TODO: add sanity check that the returned returnValueCount has value equal
+  // to expectedReturnValueCount.
+
+  // Clean up env as in:
+  //   TFE_DeleteOp(op);
+  //   TF_DeleteStatus(status);
+  auto *deleteOpFn = IGM.getTFE_DeleteOpFn();
+  Builder.CreateCall(deleteOpFn, {op});
+  auto *deleteStatusFn = IGM.getTF_DeleteStatusFn();
+  Builder.CreateCall(deleteStatusFn, {status});
+
+  auto cTensorHandle =
+      Builder.CreateLoad(returnValues.getAddress(), IGM.getPointerAlignment());
+  LLVM_DEBUG(llvm::dbgs() << "The returned tensor handle is " << *cTensorHandle
+                          << ".\n");
 
   // 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::Function *createHandleFn = findFunction(
+      "_swift_tfc_CreateFloatTensorHandleFromCTensorHandle", silModule);
+  auto tensorHandle = Builder.CreateCall(createHandleFn, {cTensorHandle});
 
   LLVM_DEBUG(
       llvm::dbgs() << "Done with IRGen for graph_op; setting explosion.\n");

stdlib/public/CTensorFlow/ctensorflow_init.cpp

@@ -1,5 +1,8 @@
 #include "ctensorflow_init.h"
 
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/c_api_experimental.h"
+#include "tensorflow/c/eager/c_api.h"
 #include "tensorflow/core/platform/init_main.h"
 
 #include <assert.h>
@@ -31,4 +34,22 @@ void InitTensorFlowRuntime(unsigned char enable_debug_logging,
   tensorflow::port::InitMain(/*usage=*/nullptr, &my_argc, &tmpArgv);
 }
 
+void *swift_tfc_CreateScalarFloatTensor(int32_t val) {
+  auto *tensor =
+      TF_AllocateTensor(TF_FLOAT, /*shape.data()*/ nullptr, /*shape.size()*/ 0,
+                        TF_DataTypeSize(TF_FLOAT) * 1);
+  auto *ptr = reinterpret_cast<char *>(TF_TensorData(tensor));
+  *reinterpret_cast<float *>(ptr) = static_cast<float>(val);
+  return tensor;
+}
+
+void swift_tfc_TFE_Execute(void *op, void **retvals, int32_t *num_retvals,
+                           void *status) {
+  int int_num_retvals = *num_retvals;
+  TFE_Execute(reinterpret_cast<TFE_Op *>(op),
+              reinterpret_cast<TFE_TensorHandle **>(retvals), &int_num_retvals,
+              reinterpret_cast<TF_Status *>(status));
+  *num_retvals = int_num_retvals;
+}
+
 }  // extern "C"

stdlib/public/CTensorFlow/ctensorflow_init.h

@@ -1,6 +1,8 @@
 #ifndef SWIFT_SRC_SWIFT_STDLIB_CTENSORFLOW_INIT_H_
 #define SWIFT_SRC_SWIFT_STDLIB_CTENSORFLOW_INIT_H_
 
+#include <stdint.h>
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -14,6 +16,17 @@ extern "C" {
 extern void InitTensorFlowRuntime(unsigned char enable_debug_logging,
                                   int verbose_level);
 
+//===----------------------------------------------------------------------===//
+// - MARK: Runtime functions to be called via IRGen.
+//===----------------------------------------------------------------------===//
+
+// Caller owns the returned tensor.
+// TODO: Generalize to create tensors from other shapes and dtypes.
+void *swift_tfc_CreateScalarFloatTensor(int32_t val);
+
+void swift_tfc_TFE_Execute(void *op, void **retvals, int32_t *num_retvals,
+                           void *status);
+
 #ifdef __cplusplus
 } /* end extern "C" */
 #endif