Factored IRGenSILFunction::visitGraphOperationInst() to handle ops other than "Const". (#19575)

* Factored IRGenSILFunction::visitGraphOperationInst() to handle ops other than
"Const".

Added a new unit test based on "Add".

* Polished the code a bit.

* Reverted #include reordering made by clang-format, to avoid unnecessary
potential merge conflicts when we downstream code from the master branch.

Author: Mingsheng Hong

include/swift/AST/TensorFlow.h

@@ -30,6 +30,10 @@ namespace tf {
   /// Return true if the given type represents a TensorFlow dtype.
   bool isTensorFlowDType(Type ty);
 
+  /// This function maps a Swift type (either a language type like Float or an
+  /// LLVM Builtin type like Builtin.f32) into the TensorFlow TF_DataType value.
+  unsigned convertSwiftTypeToTF(Type ty);
+
   /// If the specified type is the well-known TensorHandle<T> type, then return
   /// "T".  If not, return a null type.
   Type getTensorHandleElementType(Type ty);

include/swift/Runtime/RuntimeFunctions.def

@@ -1336,6 +1336,11 @@ FUNCTION(TFE_DeleteOp, TFE_DeleteOp, C_CC,
          ARGS(Int8PtrTy),
          ATTRS(NoUnwind))
 
+FUNCTION(TFE_OpAddInput, TFE_OpAddInput, C_CC,
+         RETURNS(),
+         ARGS(Int8PtrTy, Int8PtrTy, Int8PtrTy),
+         ATTRS(NoUnwind))
+
 FUNCTION(TFE_OpSetAttrType, TFE_OpSetAttrType, C_CC,
          RETURNS(),
          ARGS(Int8PtrTy, Int8PtrTy, Int32Ty),

include/swift/SIL/GraphOperationInfo.h

@@ -19,6 +19,8 @@
 #define SWIFT_SIL_GRAPH_OPERATION_INFO_H
 
 #include "swift/AST/Identifier.h"
+#include "swift/AST/TensorFlow.h"
+#include "swift/SIL/SILType.h"
 #include "swift/SIL/SILValue.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
@@ -185,6 +187,11 @@ struct GraphOperationInfo {
   static std::pair<llvm::StringRef, ArgumentLowering>
   decodeArgumentName(StringRef Name);
 };
+
+/// Determine whether the specified type is one of our well-known types, and
+/// if so, which one it is.
+TFValueKind classifyTensorFlowValue(SILType ty);
+
 } // end namespace tf
 } // end namespace swift
 #endif // SWIFT_SIL_GRAPH_OPERATION_INFO_H

lib/AST/CMakeLists.txt

@@ -3,6 +3,11 @@ if (SWIFT_FORCE_OPTIMIZED_TYPECHECKER)
   set(EXTRA_AST_FLAGS "FORCE_BUILD_OPTIMIZED")
 endif()
 
+if(SWIFT_ENABLE_TENSORFLOW)
+  find_package(TensorFlow REQUIRED)
+  include_directories(BEFORE "${TF_INCLUDE_DIR}")
+endif()
+
 add_swift_library(swiftAST STATIC
   AccessScopeChecker.cpp
   AccessRequests.cpp

lib/AST/GraphOperationInfo.cpp

@@ -149,3 +149,9 @@ std::pair<StringRef, GraphOperationInfo::ArgumentLowering>
 GraphOperationInfo::StructuredArgument::getArgumentNameAndLowering() const {
   return decodeArgumentName(Name);
 }
+
+/// Determine whether the specified type is one of our well-known types, and
+/// if so, which one it is.
+TFValueKind tf::classifyTensorFlowValue(SILType ty) {
+  return classifyTensorFlowValue(ty.getASTType());
+}

lib/AST/TensorFlow.cpp

@@ -17,7 +17,11 @@
 
 #include "swift/AST/TensorFlow.h"
 #include "swift/AST/Decl.h"
+#include "swift/AST/Module.h"
 #include "swift/AST/Types.h"
+#ifdef SWIFT_ENABLE_TENSORFLOW
+#include "tensorflow/c/c_api.h"
+#endif
 using namespace swift;
 using namespace tf;
 
@@ -36,6 +40,88 @@ bool tf::isTensorFlowDType(Type ty) {
   return !conformances.empty();
 }
 
+static bool is64(Type ty) {
+  return ty->getASTContext().LangOpts.Target.isArch64Bit();
+}
+
+/// This function maps a Swift type (either a language type like Float or an
+/// LLVM Builtin type like Builtin.f32) into the TensorFlow TF_DataType value.
+///
+/// This returns 0 (which is an invalid tensorflow type ID) on error.
+///
+unsigned tf::convertSwiftTypeToTF(Type ty) {
+#ifdef SWIFT_ENABLE_TENSORFLOW
+  // Handle wrappers like Float, which come up in TensorHandle<Float>
+  if (auto *s = ty->getAs<StructType>()) {
+    // Make sure the type is defined inside the Swift module.
+    auto context = s->getDecl()->getDeclContext()->getParentModule();
+    if (!context || context->getName().str() != "Swift")
+      return 0;
+
+    return llvm::StringSwitch<unsigned>(s->getDecl()->getNameStr())
+        .Case("Bool", TF_BOOL)
+        .Case("Int8", TF_INT8)
+        .Case("UInt8", TF_UINT8)
+        .Case("Int16", TF_INT16)
+        .Case("UInt16", TF_UINT16)
+        .Case("Int32", TF_INT32)
+        .Case("UInt32", TF_UINT32)
+        .Case("Int64", TF_INT64)
+        .Case("UInt64", TF_UINT64)
+        .Case("Int8", TF_INT8)
+        .Case("UInt8", TF_UINT8)
+        .Case("BFloat16", TF_BFLOAT16)
+        .Case("Float", TF_FLOAT)
+        .Case("Double", TF_DOUBLE)
+        .Case("Int", is64(s) ? TF_INT64 : TF_INT32)
+        .Case("UInt", is64(s) ? TF_UINT64 : TF_UINT32)
+        .Case("String", TF_STRING)
+        .Default(0);
+  }
+
+  // BuiltinIntegerType doesn't carry sign information, which TensorFlow needs,
+  // so we can't rely on getting type information from the builtin types
+  // themselves.  For now we'll just use signed types.
+  if (auto *BII = ty->getAs<BuiltinIntegerType>()) {
+    if (BII->getWidth().isPointerWidth())
+      return is64(ty) ? TF_INT64 : TF_INT32;
+
+    switch (BII->getFixedWidth()) {
+    case 1:
+      return TF_BOOL;
+    case 8:
+      return TF_INT8;
+    case 16:
+      return TF_INT16;
+    case 32:
+      return TF_INT32;
+    case 64:
+      return TF_INT64;
+    }
+  }
+
+  if (auto *BIF = ty->getAs<BuiltinFloatType>()) {
+    switch (BIF->getFPKind()) {
+    case BuiltinFloatType::IEEE16:
+      return TF_HALF;
+    case BuiltinFloatType::IEEE32:
+      return TF_FLOAT;
+    case BuiltinFloatType::IEEE64:
+      return TF_DOUBLE;
+    case BuiltinFloatType::IEEE80:
+    case BuiltinFloatType::IEEE128:
+    case BuiltinFloatType::PPC128:
+      return 0;
+    }
+  }
+
+  if (auto *BRPT = ty->getAs<BuiltinRawPointerType>()) {
+    return TF_STRING;
+  }
+#endif
+  return 0;
+}
+
 /// If the specified type is the well-known TensorHandle<T> type, then return
 /// "T".  If not, return a null type.
 Type tf::getTensorHandleElementType(Type ty) {