rework to clone further in pipeline

iree-org · Feb 19, 2025 · 9e5e53b · 9e5e53b
1 parent 33a770e
commit 9e5e53b
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 23 deletions.
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/Affinity.cpp b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Affinity.cpp
@@ -723,26 +723,11 @@ void OpAffinityPVS::initializeOperation(Operation *op, DFX::Solver &solver) {
 ChangeStatus OpAffinityPVS::updateOperation(Operation *op,
                                             DFX::Solver &solver) {
   StateType newState;
-
-  const bool consumesAny = llvm::any_of(
-      op->getOperandTypes(), +[](Type type) {
-        return isa<IREE::Stream::AffinityTypeInterface>(type);
-      });
-  if (consumesAny) {
-    for (auto operand : op->getOperands()) {
-      if (isa<IREE::Stream::AffinityTypeInterface>(operand.getType())) {
-        auto valuePVS = solver.getElementFor<ValueProducerAffinityPVS>(
-            *this, Position::forValue(operand), DFX::Resolution::REQUIRED);
-        newState ^= valuePVS;
-      }
-    }
-  } else {
-    for (auto result : op->getResults()) {
-      if (isa<IREE::Stream::AffinityTypeInterface>(result.getType())) {
-        auto valuePVS = solver.getElementFor<ValueConsumerAffinityPVS>(
-            *this, Position::forValue(result), DFX::Resolution::REQUIRED);
-        newState ^= valuePVS;
-      }
+  for (auto operand : op->getOperands()) {
+    if (isa<IREE::Stream::AffinityTypeInterface>(operand.getType())) {
+      auto valuePVS = solver.getElementFor<ValueProducerAffinityPVS>(
+          *this, Position::forValue(operand), DFX::Resolution::REQUIRED);
+      newState ^= valuePVS;
     }
   }
 

diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning/ReferencePartitioning.cpp b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning/ReferencePartitioning.cpp
@@ -99,11 +99,11 @@ partitionStreamableOpsReference(IREE::Stream::PartitioningConfigAttr config,
     IREE::Stream::AffinityAttr affinityAttr;
     if (auto affinityOp = dyn_cast<IREE::Stream::AffinityOpInterface>(op))
       affinityAttr = affinityOp.getAffinityAttr();
-    if (!IREE::Stream::AffinityAttr::canExecuteTogether(
+    bool preferCloneToConsumers = streamableOp.preferCloneToConsumers();
+    if (!preferCloneToConsumers && !IREE::Stream::AffinityAttr::canExecuteTogether(
             affinityAttr, builders[partitionOrdinal]->affinity))
       return false;
 
-    bool preferCloneToConsumers = streamableOp.preferCloneToConsumers();
     llvm::BitVector *opHazards = nullptr;
     llvm::BitVector opHazardsInCandidatePartition;
     if (preferCloneToConsumers) {

diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
@@ -2644,6 +2644,15 @@ static void printDispatchOperands(OpAsmPrinter &p, Operation *op,
   p << ")";
 }
 
+bool AsyncDispatchOp::preferCloneToConsumers() {
+  for (auto operand : getResourceOperands()) {
+    if (isa<Stream::ResourceType>(operand.getType())) {
+      return false;
+    }
+  }
+  return true;
+}
+
 LogicalResult AsyncDispatchOp::verify() {
   AsyncDispatchOp op = *this;
   if (failed(verifyOpValueSizes(op, op.getResourceOperands(),

diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
@@ -2493,7 +2493,9 @@ def Stream_AsyncDispatchOp : Stream_Op<"async.dispatch", [
   DeclareOpInterfaceMethods<SymbolUserOpInterface>,
   Stream_AffinityOp,
   Stream_AsyncPhaseOp,
-  Stream_StreamableOp,
+  DeclareOpInterfaceMethods<Stream_StreamableOp, [
+    "preferCloneToConsumers",
+  ]>,
   DeclareOpInterfaceMethods<Stream_AsyncAccessOp, [
     "getAsyncAccessRanges",
   ]>,