[SIL] Add GraphOperationInst. (#17077)

* [SIL] Add `GraphOperationInst`.

Add `GraphOperationInst` to represent graph operations.

A graph operation may have:
- Zero or more SILValue operands
- Zero or more attributes (consisting of a name and a `SymbolicValue` constant value)
- One or more result types

`GraphOperationInst` will replace `BuiltinInst` as the representation for
TensorFlow operations in the TFDeabstraction/TFPartition passes.

Todos:
 - @lattner: Edit deabstraction/partitioning to use `GraphOperationInst`.
 - Low priority: Implement SIL serialization/deserialization for `GraphOperationInst`.

Author: dan-zheng

docs/SIL.rst

@@ -5341,7 +5341,7 @@ gradient
   sil-autodiff-preserving-result ::= '[' 'preserving_result' ']'
 
   %original = function_ref @original : $(Float, Float) -> Float
-  %original_grad = gradient [wrt 0, 1] [preserving_result]
+  %original_grad = gradient [wrt 0, 1] [preserving_result] \
     %original : $(Float, Float) -> Float
 
 Computes the gradient function of a value ``%original`` using reverse-mode
@@ -5352,6 +5352,37 @@ automatic differentiation.
 This instruction is only valid in raw SIL and is rewritten by the automatic
 differentiation pass.
 
+.. SWIFT_ENABLE_TENSORFLOW
+
+Graph Program Extraction
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+graph_op
+````````
+::
+
+  sil-instruction ::= 'graph_op' string-literal
+                      '(' (sil-operand (',' sil-operand)*)? ')'
+                      ('{' (sil-graph-op-attr (',' sil-graph-op-attr)*)? '}')?
+                      ':' sil-type (',' sil-type)*
+  sil-graph-op-attr ::= sil-identifier ':' sil-symbolic-value
+  sil-symbolic-value ::= i[0-9]+ int-literal |
+                         f(32|64) float-literal |
+                         sil-type |
+                         '[' (sil-symbolic-value (',' sil-symbolic-value)*)? ']'
+
+  %add = graph_op "tf.Add"(%x : $Tensor<Float>, %y : $Tensor<Float>) \
+    {T: $Float} : $Tensor<Float>
+
+Represents a graph program operation.
+
+``graph_op`` instructions have a name, zero or more operands, zero or more
+attributes (an identifier and SIL constant value), and one or more result
+types.
+
+This instruction is only valid in raw SIL and is rewritten and extracted by the
+graph program extraction passes (debastraction, partitioning, graph lowering).
+
 Assertion configuration
 ~~~~~~~~~~~~~~~~~~~~~~~
 

include/swift/AST/DiagnosticsParse.def

@@ -1674,6 +1674,36 @@ ERROR(expr_expected_function_to_differentiate,none,
 ERROR(gradient_expr_expected_parameter,none,
       "expected a parameter, which can be the index of a function parameter with a leading dot (e.g. '.0'), or 'self'", ())
 
+// SWIFT_ENABLE_TENSORFLOW
+//------------------------------------------------------------------------------
+// Graph operation related parsing diagnostics
+//------------------------------------------------------------------------------
+ERROR(sil_graph_op_expected_attr_name,PointsToFirstBadToken,
+      "expected 'graph_op' attribute name", ())
+ERROR(sil_graph_op_expected_attr_value,PointsToFirstBadToken,
+      "expected 'graph_op' attribute value", ())
+ERROR(sil_graph_op_unhandled_attr_value,PointsToFirstBadToken,
+      "unhandled 'graph_op' attribute value", ())
+ERROR(sil_graph_op_expected_rparen,PointsToFirstBadToken,
+      "expected ')' in 'graph_op' argument list", ())
+ERROR(sil_graph_op_expected_rbrace,PointsToFirstBadToken,
+      "expected '}' in 'graph_op' attribute list", ())
+ERROR(sil_graph_op_expected_colon_after_attr_name,PointsToFirstBadToken,
+      "expected ':' after 'graph_op' attribute name", ())
+ERROR(sil_graph_op_expected_colon_before_result_types,PointsToFirstBadToken,
+      "expected ':' before 'graph_op' result types", ())
+
+ERROR(sil_const_expected_int_datatype,PointsToFirstBadToken,
+      "expected integer datatype ('i[0-9]+', e.g. 'i32')", ())
+ERROR(sil_const_expected_int_value,PointsToFirstBadToken,
+      "expected integer value in SIL constant value", ())
+ERROR(sil_const_expected_fp_datatype,PointsToFirstBadToken,
+      "expected floating point datatype ('f32' or 'f64')", ())
+ERROR(sil_const_expected_fp_value,PointsToFirstBadToken,
+      "expected floating point value in SIL constant value", ())
+ERROR(sil_const_aggregate_expected_rsquare,PointsToFirstBadToken,
+      "expected ']' at end of aggregate 'SymbolicValue'", ())
+
 #ifndef DIAG_NO_UNDEF
 # if defined(DIAG)
 #  undef DIAG

include/swift/SIL/SILBuilder.h

@@ -1300,6 +1300,14 @@ class SILBuilder {
       SILLocation Loc, SILValue Operand,
       llvm::SmallVectorImpl<SILValue> &Result);
 
+  GraphOperationInst *createGraphOperation(
+      SILLocation loc, Identifier name, ArrayRef<SILValue> operands,
+      ArrayRef<GraphOperationAttribute> attrs, ArrayRef<SILType> resultTypes) {
+    return insert(GraphOperationInst::create(
+        getModule(), getSILDebugLocation(loc), name, operands, attrs,
+        resultTypes));
+  }
+
   ClassMethodInst *createClassMethod(SILLocation Loc, SILValue Operand,
                                      SILDeclRef Member, SILType MethodTy) {
     return insert(new (getModule()) ClassMethodInst(

include/swift/SIL/SILCloner.h

@@ -1607,6 +1607,21 @@ void SILCloner<ImplClass>::visitDestructureTupleInst(
                                           getOpValue(Inst->getOperand())));
 }
 
+// SWIFT_ENABLE_TENSORFLOW
+template <typename ImplClass>
+void SILCloner<ImplClass>::visitGraphOperationInst(GraphOperationInst *Inst) {
+  getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
+  auto arguments =
+    getOpValueArray<4>(OperandValueArrayRef(Inst->getArguments()));
+  SmallVector<SILType, 4> resultTypes;
+  for (auto result : Inst->getResults())
+    resultTypes.push_back(result->getType());
+  doPostProcess(Inst,
+      getBuilder().createGraphOperation(getOpLocation(Inst->getLoc()),
+                                        Inst->getName(), arguments,
+                                        Inst->getAttributes(), resultTypes));
+}
+
 template <typename ImplClass>
 void SILCloner<ImplClass>::visitClassMethodInst(ClassMethodInst *Inst) {
   getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));

include/swift/SIL/SILInstruction.h

@@ -29,6 +29,8 @@
 #include "swift/Basic/Range.h"
 #include "swift/SIL/Consumption.h"
 #include "swift/SIL/SILAllocated.h"
+// SWIFT_ENABLE_TENSORFLOW
+#include "swift/SIL/SILConstants.h"
 #include "swift/SIL/SILDeclRef.h"
 #include "swift/SIL/SILFunctionConventions.h"
 #include "swift/SIL/SILLocation.h"
@@ -51,6 +53,8 @@ class MultipleValueInstruction;
 class MultipleValueInstructionResult;
 class DestructureTupleInst;
 class DestructureStructInst;
+// SWIFT_ENABLE_TENSORFLOW
+class GraphOperationInst;
 class NonValueInstruction;
 class SILBasicBlock;
 class SILBuilder;
@@ -8141,6 +8145,128 @@ inline DestructureTupleInst *DestructureTupleResult::getParent() {
   return cast<DestructureTupleInst>(Parent);
 }
 
+/// SWIFT_ENABLE_TENSORFLOW
+/// A result for the graph_op instruction. See documentation for
+/// graph_op for more information.
+class GraphOperationResult final : public MultipleValueInstructionResult {
+public:
+  GraphOperationResult(unsigned Index, SILType Type,
+                       ValueOwnershipKind OwnershipKind)
+  : MultipleValueInstructionResult(ValueKind::GraphOperationResult, Index,
+                                   Type, OwnershipKind) {}
+
+  static bool classof(const SILNode *N) {
+    return N->getKind() == SILNodeKind::GraphOperationResult;
+  }
+
+  GraphOperationInst *getParent() {
+    auto *Parent = MultipleValueInstructionResult::getParent();
+    return cast<GraphOperationInst>(Parent);
+  };
+
+  const GraphOperationInst *getParent() const {
+    return const_cast<GraphOperationResult *>(this)->getParent();
+  }
+};
+
+/// SWIFT_ENABLE_TENSORFLOW
+/// A graph operation attribute. Attributes have a name and a constant value.
+struct GraphOperationAttribute {
+  Identifier name;
+  SymbolicValue value;
+};
+
+/// SWIFT_ENABLE_TENSORFLOW
+/// A graph operation. This instruction will be extracted to a graph program
+/// via graph program extraction passes.
+class GraphOperationInst final
+  : public InstructionBase<
+               SILInstructionKind::GraphOperationInst,
+               MultipleValueInstruction>,
+    public MultipleValueInstructionTrailingObjects<
+               GraphOperationInst, GraphOperationResult,
+               InitialTrailingObjects<>,
+               FinalTrailingObjects<Operand, GraphOperationAttribute>> {
+  friend TrailingObjects;
+
+  /// The name of the graph operation.
+  Identifier Name;
+  /// The number of operands.
+  unsigned NumOperands;
+  /// The number of attributes.
+  unsigned NumAttributes;
+
+  GraphOperationInst(SILModule &M, SILDebugLocation loc, Identifier name,
+                     ArrayRef<SILValue> arguments,
+                     ArrayRef<GraphOperationAttribute> attrs,
+                     ArrayRef<SILType> resultTypes,
+                     ArrayRef<ValueOwnershipKind> resultOwnerships) :
+    InstructionBase(loc),
+    MultipleValueInstructionTrailingObjects(this, resultTypes,
+                                            resultOwnerships),
+    Name(name), NumOperands(arguments.size()), NumAttributes(attrs.size()) {
+      auto allOperands = getAllOperands();
+      for (unsigned i : indices(arguments))
+        new (&allOperands[i]) Operand(this, arguments[i]);
+      std::uninitialized_copy(attrs.begin(), attrs.end(),
+                              getAttributes().data());
+    }
+
+public:
+  using MultipleValueInstructionTrailingObjects::numTrailingObjects;
+  using MultipleValueInstructionTrailingObjects::totalSizeToAlloc;
+
+  ~GraphOperationInst() {
+    for (auto &operand : getAllOperands())
+      operand.~Operand();
+  }
+
+  static GraphOperationInst *create(SILModule &M, SILDebugLocation loc,
+                                    Identifier name,
+                                    ArrayRef<SILValue> arguments,
+                                    ArrayRef<GraphOperationAttribute> attrs,
+                                    ArrayRef<SILType> resultTypes);
+
+  Identifier getName() const { return Name; }
+  unsigned getNumOperands() const { return NumOperands; }
+  unsigned getNumAttributes() const { return NumAttributes; }
+
+  unsigned numTrailingObjects(OverloadToken<Operand>) const {
+    return NumOperands;
+  }
+
+  ArrayRef<Operand> getAllOperands() const {
+    return { getTrailingObjects<Operand>(), NumOperands };
+  }
+
+  MutableArrayRef<Operand> getAllOperands() {
+    return { getTrailingObjects<Operand>(), NumOperands };
+  }
+
+  OperandValueArrayRef getArguments() const {
+    return OperandValueArrayRef(getAllOperands());
+  }
+
+  ArrayRef<GraphOperationAttribute> getAttributes() const {
+    return { getTrailingObjects<GraphOperationAttribute>(), NumAttributes };
+  }
+
+  MutableArrayRef<GraphOperationAttribute> getAttributes() {
+    return { getTrailingObjects<GraphOperationAttribute>(), NumAttributes };
+  }
+
+  Optional<GraphOperationAttribute> getAttribute(StringRef name) {
+    for (auto attr : getAttributes())
+      if (attr.name.is(name))
+        return attr;
+    return None;
+  };
+
+  static bool classof(const SILNode *N) {
+    return N->getKind() == SILNodeKind::GraphOperationInst;
+  }
+};
+
 inline SILType *AllocRefInstBase::getTypeStorage() {
   // If the size of the subclasses are equal, then all of this compiles away.
   if (auto I = dyn_cast<AllocRefInst>(this))

include/swift/SIL/SILNodes.def

@@ -218,7 +218,9 @@ ABSTRACT_VALUE(MultipleValueInstructionResult, ValueBase)
   MULTIPLE_VALUE_INST_RESULT(BeginApplyResult, MultipleValueInstructionResult)
   MULTIPLE_VALUE_INST_RESULT(DestructureStructResult, MultipleValueInstructionResult)
   MULTIPLE_VALUE_INST_RESULT(DestructureTupleResult, MultipleValueInstructionResult)
-  VALUE_RANGE(MultipleValueInstructionResult, BeginApplyResult, DestructureTupleResult)
+  // SWIFT_ENABLE_TENSORFLOW
+  MULTIPLE_VALUE_INST_RESULT(GraphOperationResult, MultipleValueInstructionResult)
+  VALUE_RANGE(MultipleValueInstructionResult, BeginApplyResult, GraphOperationResult)
 
 VALUE(SILUndef, ValueBase)
 
@@ -668,10 +670,13 @@ MULTIPLE_VALUE_INST(DestructureStructInst, destructure_struct,
                     SILInstruction, None, DoesNotRelease)
 MULTIPLE_VALUE_INST(DestructureTupleInst, destructure_tuple,
                     SILInstruction, None, DoesNotRelease)
-INST_RANGE(MultipleValueInstruction, BeginApplyInst, DestructureTupleInst)
+// SWIFT_ENABLE_TENSORFLOW
+MULTIPLE_VALUE_INST(GraphOperationInst, graph_op,
+                    SILInstruction, None, DoesNotRelease)
+INST_RANGE(MultipleValueInstruction, BeginApplyInst, GraphOperationInst)
 
-NODE_RANGE(SILInstruction, AllocStackInst, DestructureTupleInst)
-NODE_RANGE(SILNode, SILPHIArgument, DestructureTupleInst)
+NODE_RANGE(SILInstruction, AllocStackInst, GraphOperationInst)
+NODE_RANGE(SILNode, SILPHIArgument, GraphOperationInst)
 
 #undef SINGLE_VALUE_INST_RANGE
 #undef INST_RANGE

include/swift/Serialization/ModuleFormat.h

@@ -55,7 +55,7 @@ const uint16_t VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t VERSION_MINOR = 415; // SWIFT_ENABLE_TENSORFLOW: gradient.
+const uint16_t VERSION_MINOR = 416; // SWIFT_ENABLE_TENSORFLOW: graph_op. 
 
 using DeclIDField = BCFixed<31>;
 

lib/IRGen/IRGenSIL.cpp

@@ -907,6 +907,11 @@ class IRGenSILFunction :
     llvm_unreachable("gradient is not valid in canonical SIL");
   }
 
+  // SWIFT_ENABLE_TENSORFLOW
+  void visitGraphOperationInst(GraphOperationInst *i) {
+    llvm_unreachable("graph_op is not valid in canonical SIL");
+  }
+
   void visitFunctionRefInst(FunctionRefInst *i);
   void visitAllocGlobalInst(AllocGlobalInst *i);
   void visitGlobalAddrInst(GlobalAddrInst *i);