Support cloning operand-less dispatches

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;
     }
   }
 

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

@@ -56,9 +56,12 @@ void PartitionSet::dump(AsmState &asmState) {}
 LogicalResult Partition::verify(Location loc) {
   // Ensure all ops are compatible with the partition affinity.
   for (auto *op : ops) {
+    auto streamableOp =
+        dyn_cast_if_present<IREE::Stream::StreamableOpInterface>(op);
+    bool cloned = streamableOp && streamableOp.preferCloneToConsumers();
     if (auto affinityOp = dyn_cast<IREE::Stream::AffinityOpInterface>(op)) {
-      if (!IREE::Stream::AffinityAttr::areCompatible(
-              affinity, affinityOp.getAffinityAttr())) {
+      if (!cloned && !IREE::Stream::AffinityAttr::areCompatible(
+                         affinity, affinityOp.getAffinityAttr())) {
         return op->emitError("op affinity ")
                << affinityOp.getAffinityAttr()
                << " is not compatible with the partition affinity " << affinity;

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

@@ -99,11 +99,12 @@ 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) {

compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp

@@ -2189,6 +2189,15 @@ TensorDispatchOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
                                   getWorkload(), symbolTable);
 }
 
+bool TensorDispatchOp::preferCloneToConsumers() {
+  for (auto operand : getMixedOperands()) {
+    if (isa<Stream::ResourceType>(operand.getType())) {
+      return false;
+    }
+  }
+  return true;
+}
+
 std::pair<unsigned, unsigned>
 TensorDispatchOp::getTiedOperandsIndexAndLength() {
   return getODSOperandIndexAndLength(1); // $operands
@@ -2644,6 +2653,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(),

compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td

@@ -1744,7 +1744,9 @@ def Stream_TensorDispatchOp : Stream_Op<"tensor.dispatch", [
   DeclareOpInterfaceMethods<SymbolUserOpInterface>,
   Stream_AffinityOp,
   Stream_TensorPhaseOp,
-  Stream_StreamableOp,
+  DeclareOpInterfaceMethods<Stream_StreamableOp, [
+    "preferCloneToConsumers",
+  ]>,
   Util_SizeAwareOp,
   DeclareOpInterfaceMethods<Util_TiedOpInterface, [
     "getTiedOperandsIndexAndLength",
@@ -2493,7 +2495,9 @@ def Stream_AsyncDispatchOp : Stream_Op<"async.dispatch", [
   DeclareOpInterfaceMethods<SymbolUserOpInterface>,
   Stream_AffinityOp,
   Stream_AsyncPhaseOp,
-  Stream_StreamableOp,
+  DeclareOpInterfaceMethods<Stream_StreamableOp, [
+    "preferCloneToConsumers",
+  ]>,
   DeclareOpInterfaceMethods<Stream_AsyncAccessOp, [
     "getAsyncAccessRanges",
   ]>,

compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleExecution.cpp

@@ -162,6 +162,12 @@ struct ExecutePartitionBuilder {
       if (affinityOp.getAffinityAttr() == partition->affinity) {
         affinityOp.setAffinityAttr(nullptr);
       }
+
+      auto streamableOp =
+          dyn_cast<IREE::Stream::StreamableOpInterface>(clonedOp);
+      if (streamableOp && streamableOp.preferCloneToConsumers()) {
+        affinityOp.setAffinityAttr(nullptr);
+      }
     }
 
     return true;