Fix an issue in ReferencePatitioning. (iree-org#21343)

When cloning operations into partitions, we intentionally only check the
the users inside those partitions for dependencies. However if we cannot
clone the operation into that particular partition for any other reason
(hazard) then we have to rely on a dependency from another partition, we
were not checking for hazards in this case.

This solution is somewhat crude, although should be correct in all
cases. If we are not able to clone
an operation into ALL consumers, then we make sure for any partition it
IS going to be cloned into, there will not be a data-hazard.

A more elegant solution would be to leave it in all possible cloned
partitions, but then for any other consunmers make sure they have at
least one clone that would not make a cycle otherwise put it into a new
partition. However that would require a re-work of the sorting logic
which treats all cloned ops as interchangeable.

---------

Signed-off-by: Andrew Woloszyn <[email protected]>
Co-authored-by: Ben Vanik <[email protected]>

Author: AWoloszyn

compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning.cpp

@@ -23,14 +23,14 @@ void dumpPartition(Partition &partition, AsmState &asmState) {
     partition.affinity.dump();
     llvm::dbgs() << "\n";
   }
-  llvm::dbgs() << " INS:\n  ";
-  llvm::interleaveComma(partition.ins, llvm::dbgs(), [&](Value in) {
-    in.printAsOperand(llvm::dbgs(), asmState);
-  });
-  llvm::dbgs() << "\n OUTS:\n  ";
-  llvm::interleaveComma(partition.outs, llvm::dbgs(), [&](Value out) {
-    out.printAsOperand(llvm::dbgs(), asmState);
-  });
+  llvm::dbgs() << " INS:\n    ";
+  llvm::interleave(
+      partition.ins, llvm::dbgs(),
+      [&](Value in) { in.print(llvm::dbgs(), asmState); }, "\n    ");
+  llvm::dbgs() << "\n OUTS:\n    ";
+  llvm::interleave(
+      partition.outs, llvm::dbgs(),
+      [&](Value out) { out.print(llvm::dbgs(), asmState); }, "\n    ");
   llvm::dbgs() << "\n OPS:\n";
   for (auto *op : llvm::reverse(partition.ops)) {
     llvm::dbgs() << "  ";

compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning/ReferencePartitioning.cpp

@@ -37,6 +37,36 @@ static std::unique_ptr<AsmState> getRootAsmState(Block *block) {
   return nullptr;
 }
 
+struct OpInfo {
+  // Which partitions the op is contained within.
+  llvm::BitVector membership;
+  // Which partitions transitively depend on this operation.
+  llvm::BitVector hazards;
+};
+
+struct PartitionBuilder {
+  unsigned ordinal;
+  // Affinity of the partition.
+  IREE::Stream::AffinityAttr affinity;
+  // Ops present in the partition; ops may be present in multiple partitions.
+  SetVector<Operation *> ops;
+  // Ops that were cloned and are known not to have their values escape.
+  DenseSet<Operation *> clonedOps;
+  // Which partitions transitively depend on this partition.
+  llvm::BitVector hazards;
+  void insert(Operation *op, OpInfo &opInfo) {
+    if (auto affinityOp = dyn_cast<IREE::Stream::AffinityOpInterface>(op)) {
+      affinity = affinity ? affinity.joinAND(affinityOp.getAffinityAttr())
+                          : affinityOp.getAffinityAttr();
+    }
+    opInfo.membership.set(ordinal);
+    if (opInfo.hazards.size() > ordinal)
+      opInfo.hazards.reset(ordinal);
+    ops.insert(op);
+    hazards |= opInfo.hazards;
+  }
+};
+
 // This is terrible. See Stream/Analysis/Partition.h for a description of what
 // a real implementation would do. We want cost modeling for tie breakers when
 // an op could be in multiple partitions, cloning for ops that are not worth
@@ -46,36 +76,8 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config,
                                 Block *block) {
   PartitionSet partitionSet;
 
-  struct OpInfo {
-    // Which partitions the op is contained within.
-    llvm::BitVector membership;
-    // Which partitions transitively depend on this operation.
-    llvm::BitVector hazards;
-  };
   DenseMap<Operation *, OpInfo> opInfos;
 
-  struct PartitionBuilder {
-    unsigned ordinal;
-    // Affinity of the partition.
-    IREE::Stream::AffinityAttr affinity;
-    // Ops present in the partition; ops may be present in multiple partitions.
-    SetVector<Operation *> ops;
-    // Ops that were cloned and are known not to have their values escape.
-    DenseSet<Operation *> clonedOps;
-    // Which partitions transitively depend on this partition.
-    llvm::BitVector hazards;
-    void insert(Operation *op, OpInfo &opInfo) {
-      if (auto affinityOp = dyn_cast<IREE::Stream::AffinityOpInterface>(op)) {
-        affinity = affinity ? affinity.joinAND(affinityOp.getAffinityAttr())
-                            : affinityOp.getAffinityAttr();
-      }
-      opInfo.membership.set(ordinal);
-      if (opInfo.hazards.size() > ordinal)
-        opInfo.hazards.reset(ordinal);
-      ops.insert(op);
-      hazards |= opInfo.hazards;
-    }
-  };
   SmallVector<std::unique_ptr<PartitionBuilder>> builders;
   llvm::BitVector usableBuilders;
 
@@ -93,7 +95,8 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config,
   };
 
   auto canAddOpToPartition = [&](Operation &op, OpInfo &opInfo,
-                                 unsigned partitionOrdinal) {
+                                 unsigned partitionOrdinal,
+                                 bool check_for_clones = true) {
     auto streamableOp = dyn_cast<IREE::Stream::StreamableOpInterface>(op);
     if (!streamableOp) {
       return false;
@@ -111,7 +114,8 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config,
       return false;
     }
 
-    bool preferCloneToConsumers = streamableOp.preferCloneToConsumers();
+    bool preferCloneToConsumers =
+        check_for_clones && streamableOp.preferCloneToConsumers();
     llvm::BitVector *opHazards = nullptr;
     llvm::BitVector opHazardsInCandidatePartition;
     if (preferCloneToConsumers) {
@@ -154,6 +158,13 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config,
   llvm::DenseMap<Operation *, llvm::SmallVector<Operation *>> syncOps;
 
   for (auto &op : llvm::reverse(*block)) {
+
+    LLVM_DEBUG({
+      llvm::dbgs() << "====\nPartitioning op:\n";
+      op.print(llvm::dbgs(), *asmState);
+      llvm::dbgs() << "\n";
+    });
+
     // Skip constants; they just add noise (and since they are heavily CSE'd
     // they have lots of users to test).
     if (op.hasTrait<OpTrait::ConstantLike>()) {
@@ -190,6 +201,11 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config,
           syncOps[producer] = llvm::SmallVector<Operation *>();
         }
         syncOps[producer].push_back(&op);
+        LLVM_DEBUG({
+          llvm::dbgs() << "Skipping sync op for now \n";
+          op.print(llvm::dbgs(), *asmState);
+          llvm::dbgs() << "\n";
+        });
         continue;
       }
     }
@@ -203,12 +219,6 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config,
     opInfo.hazards.reserve(builders.size() + 1);
     opInfo.hazards.resize(builders.size(), /*t=*/false);
 
-    LLVM_DEBUG({
-      llvm::dbgs() << "====\nPartitioning op:\n";
-      op.print(llvm::dbgs(), *asmState);
-      llvm::dbgs() << "\n";
-    });
-
     // Set bits for each partition this op may be able to be placed into.
     // We prune the set based on whether the users are part of a transitive
     // dependency chain down the use-def chain to a partition.
@@ -242,6 +252,19 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config,
           continue;
         }
         auto &userInfo = userInfoIt->second;
+
+        LLVM_DEBUG({
+          llvm::dbgs() << "Testing sync user:\n";
+          user->print(llvm::dbgs(), *asmState);
+          llvm::dbgs() << "\n";
+          for (auto membershipOrdinal : userInfo.membership.set_bits()) {
+            llvm::dbgs() << "  member of partition " << membershipOrdinal
+                         << "\n";
+          }
+          for (auto hazardOrdinal : userInfo.hazards.set_bits()) {
+            llvm::dbgs() << "  hazard w/ partition " << hazardOrdinal << "\n";
+          }
+        });
         opInfo.hazards |= userInfo.membership;
         opInfo.hazards |= userInfo.hazards;
         consumers.reset();
@@ -282,6 +305,34 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config,
       continue;
     }
 
+    // If we prefer to clone to our consumers, but we are
+    // only cloning to a subset, we have to re-check our
+    // partitions as they may generate cycles.
+    if (streamableOp.preferCloneToConsumers()) {
+      auto tempCandidates = candidates;
+      tempCandidates &= consumers;
+      if (tempCandidates.count() != consumers.count()) {
+        // Prune candidates that do not have a compatible affinity.
+        for (auto ordinal : candidates.set_bits()) {
+          if (!canAddOpToPartition(op, opInfo, ordinal, false)) {
+            LLVM_DEBUG(llvm::dbgs() << "Candidate partition " << ordinal
+                                    << " incompatible for clone\n");
+            candidates.reset(ordinal);
+          }
+        }
+
+        for (auto syncOp : syncOps[&op]) {
+          for (auto ordinal : candidates.set_bits()) {
+            if (!canAddOpToPartition(*syncOp, opInfo, ordinal, false)) {
+              LLVM_DEBUG(llvm::dbgs() << "Candidate partition " << ordinal
+                                      << " incompatible for clone\n");
+              candidates.reset(ordinal);
+            }
+          }
+        }
+      }
+    }
+
     // First see which partitions are consuming this that we can also safely
     // move in to.
     consumers &= candidates;
@@ -321,6 +372,10 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config,
 
     // If we have synchronization operations we can place in the last block:
     for (auto syncOp : syncOps[&op]) {
+      LLVM_DEBUG(llvm::dbgs() << "Moving sync to candidate partition "
+                              << firstCandidateOrdinal << ":\n    ");
+      LLVM_DEBUG(syncOp->print(llvm::dbgs(), *asmState));
+      LLVM_DEBUG(llvm::dbgs() << "\n");
       builder->insert(syncOp, opInfo);
     }