Improve deabstraction SSA promotion logic, fixing SR-8395 (#18400)

Improve deabstraction SSA promotion logic, fixing SR-8395
 - Handle the simple case of storing into a struct, to handle the common
   pattern of a store through a struct_element_addr into a Tensor.
 - Fix handling of begin_access to be a separate loop.  Before, we were
   adding new entries to the use list that we're iterating over.  This
   was almost fine (no invalid iterators or anything) but the entries
   are added to the start of the list so we wouldn't see them.
 - Fix a bug this exposed in graph lowering where we wouldn't handle
   aggregate wrappers around elements in a Tensor initializer.

Author: lattner

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: ----
 */