fix double-deabstraction bug (#19593)

Author: marcrasi

include/swift/SIL/GraphOperationBuilder.h

@@ -44,6 +44,10 @@ class GraphOperationBuilder {
   void addListArgument(llvm::ArrayRef<SILValue> arguments,
                        llvm::StringRef name = llvm::StringRef());
 
+  /// Add a list argument to the GraphOperationInst, with an optional name.
+  void addListArgument(OperandValueArrayRef arguments,
+                       llvm::StringRef name = llvm::StringRef());
+
   /// Add an attribute with known constant value to the GraphOperationInst.
   /// Returns a reference to the attribute, valid for the lifetime of the
   /// GraphOperationBuilder, that you can use to mutate the attribute before

lib/AST/GraphOperationBuilder.cpp

@@ -43,6 +43,17 @@ void GraphOperationBuilder::addListArgument(ArrayRef<SILValue> arguments,
   }
 }
 
+/// Add a list argument to the GraphOperationInst, with an optional name.
+void GraphOperationBuilder::addListArgument(OperandValueArrayRef arguments,
+                                            StringRef name) {
+  MangledName += ",L";
+  MangledName += name;
+  for (auto argument : arguments) {
+    MangledName += ",e";
+    Operands.push_back(argument);
+  }
+}
+
 /// Add an attribute with known constant value to the GraphOperationInst.
 /// Returns a reference to the attribute, valid for the lifetime of the
 /// GraphOperationBuilder, that you can use to mutate the attribute before

lib/SILOptimizer/Mandatory/TFDeabstraction.cpp

@@ -384,22 +384,31 @@ static GraphOperationInst *simplifyOperands(GraphOperationInst *origInst,
   };
 
   // Predicate that returns true if an argument of the specified type should be
-  // rewritten - either to load an address argument or expand a struct
-  // parameter.
-  auto canSimplifyOperand =
-      [&](SILType type, GraphOperationInfo::ArgumentLowering lowering) -> bool {
+  // rewritten to load an address argument or expand a struct parameter.
+  auto canSimplifyArgumentType = [&](SILType type) -> bool {
     return isLoadableAddressType(type) ||
-           getPrimitiveStructField(type.getASTType()) != nullptr ||
-           lowering == GraphOperationInfo::ArgumentLowering::Out;
+           getPrimitiveStructField(type.getASTType()) != nullptr;
+  };
+
+  // Predicate that returns true if an argument should be rewritten.
+  auto canSimplifyArgument =
+      [&](const GraphOperationInfo::StructuredArgument &argument) -> bool {
+    switch (argument.getKind()) {
+    case GraphOperationInfo::SAK_Single:
+      return canSimplifyArgumentType(argument.getSingleArgument()->getType()) ||
+             std::get<1>(argument.getArgumentNameAndLowering()) ==
+                 GraphOperationInfo::ArgumentLowering::Out;
+    case GraphOperationInfo::SAK_List:
+      // We can get SAK_List arguments from inlining functions that have already
+      // been deabstracted. These arguments do not need further simplification.
+      return false;
+    }
   };
 
   // If we don't have to change any arguments, don't rewrite the graph_op.
   bool mustChangeGraphOp = false;
   for (auto &argument : opInfo.getStructuredArguments()) {
-    assert(argument.getKind() == GraphOperationInfo::SAK_Single &&
-           "SILGen should not have generated a list argument");
-    if (canSimplifyOperand(argument.getSingleArgument()->getType(),
-                           std::get<1>(argument.getArgumentNameAndLowering()))) {
+    if (canSimplifyArgument(argument)) {
       mustChangeGraphOp = true;
       break;
     }
@@ -413,10 +422,21 @@ static GraphOperationInst *simplifyOperands(GraphOperationInst *origInst,
   // Okay, we do have to simplify something.  Scan through and rewrite arguments.
   SILBuilder B(origInst);
   GraphOperationBuilder opBuilder(opInfo.getOperationName());
+  // Pass attributes through.
+  for (auto &attr : origInst->getAttributes())
+    opBuilder.addAttribute(attr);
   SILValue outParameterAddress;
   for (auto &argument : opInfo.getStructuredArguments()) {
+    if (argument.getKind() == GraphOperationInfo::SAK_List) {
+      // We can get SAK_List arguments from inlining functions that have already
+      // been deabstracted. Pass these arguments through.
+      opBuilder.addListArgument(argument.getArgumentList(),
+                                argument.getArgumentNameWithSuffix());
+      continue;
+    }
+
     assert(argument.getKind() == GraphOperationInfo::SAK_Single &&
-           "SILGen should not have generated a list argument");
+           "should have already handled all other argument kinds");
     auto argumentValue = argument.getSingleArgument();
     auto argumentLowering = std::get<1>(argument.getArgumentNameAndLowering());
 
@@ -2126,6 +2146,14 @@ void TFDeabstraction::evaluateAttributesAndDoPacking(
   // Find the device attribute specified for the instruction if present.
   StringRef opDevice;
 
+  // Pass attributes through.
+  for (auto &attr : origInst->getAttributes()) {
+    if (attr.name.str() == DEVICE_ATTR) {
+      opDevice = attr.value.getStringValue();
+    }
+    opBuilder.addAttribute(attr);
+  }
+
   // It is common to have input lists with repeated elements.  These will
   // generally be uniqued on entry to this routine.  We cache the projections in
   // these maps so that we can reuse them and avoid code bloat.
@@ -2137,8 +2165,17 @@ void TFDeabstraction::evaluateAttributesAndDoPacking(
 
   for (auto i : range(opInfo.getStructuredArguments().size())) {
     auto argument = opInfo.getStructuredArguments()[i];
+
+    if (argument.getKind() == GraphOperationInfo::SAK_List) {
+      // We can get SAK_List arguments from inlining functions that have already
+      // been deabstracted. Pass these arguments through.
+      opBuilder.addListArgument(argument.getArgumentList(),
+                                argument.getArgumentNameWithSuffix());
+      continue;
+    }
+
     assert(argument.getKind() == GraphOperationInfo::SAK_Single &&
-           "SILGen should not have generated a list argument");
+           "should have already handled all other argument kinds");
     auto argumentValue = argument.getSingleArgument();
     auto argumentTy = argumentValue->getType();
     auto argumentNameAndLowering = argument.getArgumentNameAndLowering();

test/TensorFlow/crashers.swift

@@ -372,3 +372,21 @@ public func SR8419(iterationCount: Int) {
     }
   }
 }
+
+// If `deabstractedCallee` gets deabstracted before `inlineDeabstracted_*`,
+// then the insts in `deabstractedCallee` get deabstracted twice. There was
+// a bug where the compiler crashed when deabstracting certain graph_ops twice.
+// There is no guaranteed deabstraction order, so this test isn't guaranteed to
+// catch the problem. Sandwiching `deabstractedCallee` between two callers
+// makes this test catch the problem as long as the order happens to be linear
+// up or down.
+public func inlineDeabstracted_a() -> Tensor<Float> {
+  return deabstractedCallee([1, 2, 3])
+}
+// expected-warning @+1 {{implicitly copied}}
+public func deabstractedCallee(_ t: Tensor<Float>) -> Tensor<Float> {
+  return t ++ Tensor<Float>([1, 2, 3]) // expected-note {{value used here}}
+}
+public func inlineDeabstracted_b() -> Tensor<Float> {
+  return deabstractedCallee([1, 2, 3])
+}