Improve deabstraction SSA promotion logic, fixing SR-8395

lib/SILOptimizer/Mandatory/TFDeabstraction.cpp

@@ -354,14 +354,13 @@ static BuiltinInst *simplifyOperands(BuiltinInst *inst, TFDeabstraction &TFDA) {
       if (!decl || !isa<StructDecl>(decl)) return nullptr;
 
       // Check to see if there is a single stored field.
-      auto fieldIt = decl->getStoredProperties().begin();
-      if (fieldIt == decl->getStoredProperties().end()) return nullptr;
+      auto field = tf::getFieldIfContainsSingleField(decl);
+      if (!field) return nullptr;
 
       // If this is the top level of the struct, retain the field decl.
-      if (result == nullptr) result = *fieldIt;
+      if (result == nullptr) result = field;
 
-      type = (*fieldIt++)->getType();
-      if (fieldIt != decl->getStoredProperties().end()) return nullptr;
+      type = field->getType();
 
       // If we unwrapped a level and got to a builtin type, then this is a
       // wrapper.
@@ -1190,10 +1189,33 @@ void TFDeabstraction::prepareStackAllocForPromotion(AllocStackInst *alloc) {
   // we have tensor values mixed in with other random values that shouldn't
   // (or can't) be loaded.  For now, we can just fail to deabstract these
   // cases.
+
+  // Our first scan will look for begin_access instructions and remove them,
+  // allowing the second pass to be simpler.
+  for (auto UI = alloc->use_begin(); UI != alloc->use_end();) {
+    auto *begin = dyn_cast<BeginAccessInst>((*UI++)->getUser());
+    if (!begin)
+      continue;
+
+    // If we have a begin_access instruction, replace uses of begin_access with
+    // uses of the original value and remove the end_access.
+    for (auto UI = begin->use_begin(); UI != begin->use_end();) {
+      auto *use = *UI++;
+      auto inst = use->getUser();
+      if (isa<EndAccessInst>(inst))
+        inst->eraseFromParent();
+      else
+        use->set(alloc);
+    }
+    begin->eraseFromParent();
+  }
+
+  // Our second pass looks for aggregate operations and struct_element_addrs
+  // that poke inside the allocation.
   for (auto UI = alloc->use_begin(); UI != alloc->use_end();) {
-    auto inst = (*UI)->getUser();
+    auto *inst = (*UI)->getUser();
 
-    if (auto sea = dyn_cast<StructElementAddrInst>(inst))
+    if (auto *sea = dyn_cast<StructElementAddrInst>(inst)) {
       if (auto *use = sea->getSingleUse()) {
         // If we have a load(struct_element_addr(alloc)) turn it into
         // struct_extract(load(alloc)).
@@ -1210,7 +1232,31 @@ void TFDeabstraction::prepareStackAllocForPromotion(AllocStackInst *alloc) {
           sea->eraseFromParent();
           continue;
         }
+
+        // If we have a store(x ->struct_element_addr(alloc)), turn it into a
+        // load of the whole value, a bunch of extracts, then a struct_inst
+        // to rebuild the whole value, then a store of the whole thing.
+        //
+        // TODO: For now, we only handle a single element struct, which is
+        // considerably simpler.
+        //
+        if (auto *store = dyn_cast<StoreInst>(use->getUser())) {
+          if (use->getOperandNumber() == 1 &&  // store TO the alloca.
+              tf::getFieldIfContainsSingleField(sea->getStructDecl())) {
+            SILBuilder B(store);
+            auto *newStruct = B.createStruct(store->getLoc(),
+                                             alloc->getType().getObjectType(),
+                                             store->getOperand(0));
+            B.createStore(store->getLoc(), newStruct, sea->getOperand(),
+                          store->getOwnershipQualifier());
+            store->eraseFromParent();
+            ++UI;
+            sea->eraseFromParent();
+            continue;
+          }
+        }
       }
+    }
 
     // Explode aggregate by-address instructions like copy-addr.
     if (explodeAggregateInst(inst, /*all types*/nullptr)) {
@@ -1219,28 +1265,8 @@ void TFDeabstraction::prepareStackAllocForPromotion(AllocStackInst *alloc) {
       continue;
     }
 
-    // If we have an instruction other than begin_access, remember it.
-    auto *begin = dyn_cast<BeginAccessInst>(inst);
-    if (!begin) {
-      ++UI;
-      continue;
-    }
-
-    // If we have a begin_access instruction, look through it.  Add all of the
-    // users to the users list, and replace uses of begin_access with uses of
-    // the original value.  Finally, ignore and remove the end_access.
-    for (auto UI = begin->use_begin(); UI != begin->use_end();) {
-      auto *use = *UI++;
-      auto inst = use->getUser();
-      if (isa<EndAccessInst>(inst)) {
-        inst->eraseFromParent();
-      } else {
-        use->set(alloc);
-      }
-    }
-
+    // Otherwise we have something else, leave it alone.
     ++UI;
-    begin->eraseFromParent();
   }
 }
 

lib/SILOptimizer/Mandatory/TFLowerGraph.cpp

@@ -2119,6 +2119,7 @@ TFGraphFunctionLowering::visitGraphOperationInst(GraphOperationInst *inst) {
       // The scalar case is very simple, the shape of a scalar is 0d, and the
       // data type comes from an attr that should already be processed.
       SmallVector<int64_t, 4> shape;
+      attrValue = attrValue.lookThroughSingleElementAggregates();
       if (attrValue.getKind() == SymbolicValue::Integer ||
           attrValue.getKind() == SymbolicValue::Float) {
         if (addScalar(attrValue, elements))

lib/SILOptimizer/Mandatory/TFPartition.cpp

@@ -82,13 +82,9 @@ static bool isUserIgnoredByPartitioning(SILInstruction *inst) {
 /// type of the single member, asserting and aborting if we get something
 /// unexpected.
 static CanType getSingleElementDeclFieldType(NominalTypeDecl *decl) {
-  auto fieldIt = decl->getStoredProperties().begin();
-  assert(fieldIt != decl->getStoredProperties().end() &&
-         "Struct should have one member");
-  auto fieldType = (*fieldIt++)->getType()->getCanonicalType();
-  assert(fieldIt == decl->getStoredProperties().end() &&
-         "Struct should have one member");
-  return fieldType;
+  auto *field = tf::getFieldIfContainsSingleField(decl);
+  assert(field && "Struct should have one member");
+  return field->getType()->getCanonicalType();
 }
 
 /// Classification of instructions that are interesting to the partitioning

lib/SILOptimizer/Mandatory/TFUtilities.cpp

@@ -84,6 +84,19 @@ llvm::raw_ostream *tf::getTFDumpIntermediateStream() {
   return &fileStream;
 }
 
+/// If the specified decl has a single stored field, return it.  Otherwise
+/// return null.
+VarDecl *tf::getFieldIfContainsSingleField(NominalTypeDecl *decl) {
+  // Check to see if there is a single stored field.
+  auto fieldIt = decl->getStoredProperties().begin();
+  if (fieldIt == decl->getStoredProperties().end())
+    return nullptr;
+  auto result = *fieldIt++;
+  if (fieldIt != decl->getStoredProperties().end())
+    return nullptr;
+  return result;
+}
+
 bool tf::isTensorHandle(SILType ty) {
   return (bool)isTensorHandle(ty.getASTType());
 }

lib/SILOptimizer/Mandatory/TFUtilities.h

@@ -37,6 +37,10 @@ namespace tf {
   /// return null.  This is used for integration unit tests and debugging.
   llvm::raw_ostream *getTFDumpIntermediateStream();
 
+  /// If the specified decl has a single stored field, return it.  Otherwise
+  /// return null.
+  VarDecl *getFieldIfContainsSingleField(NominalTypeDecl *decl);
+
   /// If the specified type is the well-known TensorHandle<T> type, then return
   /// "T".  If not, return a null type.
   bool isTensorHandle(SILType ty);

test/TensorFlow/debugging.swift

@@ -18,9 +18,7 @@ public func debugValuesInLoop(_ x: Tensor<Float>) {
 // CHECK-LABEL: --- TFPartition Accelerator Result: {{.*}}basicDebugValues{{.*}}
 // CHECK: @{{.*}}basicDebugValues{{.*}}.tf
 // CHECK: [[ONE:%.*]] = graph_op "Const"
-// CHECK-NEXT: graph_op "tfc.SendToHost,i"
 // CHECK: [[ADD_RESULT:%.*]] = graph_op "Add,i,i"
-// CHECK-NEXT: graph_op "tfc.SendToHost,i"([[ADD_RESULT]] : $TensorHandle<Float>)
 // CHECK: graph_op "Square,i"([[ADD_RESULT]] : $TensorHandle<Float>) {T: $Float, __device: "/device:CPU:0"} : $TensorHandle<Float>
 
 

test/TensorFlow/optimization-disabled.swift

@@ -2,17 +2,14 @@
 import TensorFlow
 
 public func testArrayValues() -> Tensor<Float> {
-  // expected-warning @+1 14 {{value implicitly copied to the host}}
   let x: Tensor<Float> = [[1, 2], [3, 4]]
   return (matmul(x, x) + x).toHost()
-// expected-warning @-1 {{value implicitly copied to the host}}
 }
 
 /*
 CHECK-LABEL: --- TFPartition Accelerator Result: {{.*}}testArrayValues
 CHECK: %0 = graph_op "Const"() {dtype: $Float, value$tensor: f32 0x3F800000 /* 1 */, __device: "ALL_DEVICES"} : $TensorHandle<Float>
-CHECK: %1 = graph_op "tfc.SendToHost,i"(%0 : $TensorHandle<Float>) {tensorId: i32 0, __device: "/device:CPU:0"}
-CHECK-NOT: tfc.RecvFromHost
+CHECK: %1 = graph_op "Const"() {dtype: $Float, value$tensor: f32 0x40000000 /* 2 */
 CHECK-LABEL: ----
 */