Adding tryLookupResourceUsageAffinity. (#20891)

This will be used to lower into the new `#hal.device.optimal` attr added
in #20879.

Progress on #20855 (still needs usage, this is just the analysis).
Progress on #20851.

Author: benvanik

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

@@ -385,6 +385,15 @@ TraversalResult ValueConsumerAffinityPVS::updateFromUse(Value value,
           return TraversalResult::INCOMPLETE;
         }
       })
+      .Case([&](IREE::Stream::YieldOp op) {
+        auto parentOp = op->getParentOp();
+        auto &resultPVS = solver.getElementFor<ValueConsumerAffinityPVS>(
+            *this,
+            Position::forValue(parentOp->getResult(operand.getOperandNumber())),
+            DFX::Resolution::REQUIRED);
+        newState ^= resultPVS.getState();
+        return TraversalResult::COMPLETE;
+      })
       .Case([&](IREE::Util::ReturnOp op) {
         return solver.getExplorer().walkIncomingCalls(
             op->getParentOfType<mlir::CallableOpInterface>(),
@@ -985,6 +994,41 @@ bool AffinityAnalysis::tryLookupResourceAffinity(
   return true;
 }
 
+bool AffinityAnalysis::tryLookupResourceUsageAffinity(
+    Value value, SmallVectorImpl<IREE::Stream::AffinityAttr> &affinities) {
+  auto producerPVS = solver.lookupElementFor<ValueProducerAffinityPVS>(
+      Position::forValue(value));
+  if (producerPVS && producerPVS->isValidState() &&
+      !producerPVS->isUndefContained()) {
+    if (!producerPVS->getAssumedSet().empty()) {
+      llvm::append_range(affinities, producerPVS->getAssumedSet());
+    }
+  }
+  auto consumerPVS = solver.lookupElementFor<ValueConsumerAffinityPVS>(
+      Position::forValue(value));
+  if (consumerPVS && consumerPVS->isValidState() &&
+      !consumerPVS->isUndefContained()) {
+    if (!consumerPVS->getAssumedSet().empty()) {
+      // There's probably a better way to do this but the template f*ckery
+      // required is too painful. I suspect a concat + sort + unique would do
+      // it, but not be any more efficient as we'd have to do the sort twice. We
+      // should probably change the functions to take a SetVectorImpl instead.
+      for (auto affinity : consumerPVS->getAssumedSet()) {
+        if (!producerPVS->getAssumedSet().contains(affinity)) {
+          affinities.push_back(affinity);
+        }
+      }
+    }
+  }
+  if (affinities.empty()) {
+    // Analysis completed but no affinity was specified; try to find a default.
+    return tryLookupDefaultAffinity(value.getParentBlock()->getParentOp(),
+                                    affinities);
+  }
+  sortAffinities(affinities);
+  return true;
+}
+
 LogicalResult AffinityAnalysis::run() {
   // Initialize globals so that we can assign them affinity.
   explorer.forEachGlobal([&](const auto *globalInfo) {
@@ -1029,6 +1073,8 @@ LogicalResult AffinityAnalysis::run() {
         if (isa<IREE::Stream::AffinityTypeInterface>(result.getType())) {
           solver.getOrCreateElementFor<ValueProducerAffinityPVS>(
               Position::forValue(result));
+          solver.getOrCreateElementFor<ValueConsumerAffinityPVS>(
+              Position::forValue(result));
         }
       }
     });

compiler/src/iree/compiler/Dialect/Stream/Analysis/Affinity.h

@@ -86,11 +86,16 @@ class AffinityAnalysis {
   // If all affinities are compatible one will be chosen in an unspecified way.
   IREE::Stream::AffinityAttr lookupResourceAffinity(Value value);
 
-  // Populates all potential affinities of |value| in |affinities|.
+  // Populates all potential producer affinities of |value| in |affinities|.
   // Returns false if analysis failed and the set of affinities is unknown.
   bool tryLookupResourceAffinity(
       Value value, SmallVectorImpl<IREE::Stream::AffinityAttr> &affinities);
 
+  // Populates all potential consumer affinities of |value| in |affinities|.
+  // Returns false if analysis failed and the set of affinities is unknown.
+  bool tryLookupResourceUsageAffinity(
+      Value value, SmallVectorImpl<IREE::Stream::AffinityAttr> &affinities);
+
 private:
   Explorer explorer;
   llvm::BumpPtrAllocator allocator;

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

@@ -528,8 +528,9 @@ def Stream_Resource : TypeDef<Stream_Dialect, "Resource", [
     them does not extend outside of the invocation they are provided to.
 
     Stream values are not usable directly outside of a stream execution or
-    transfer operation. If the contents of the value are needed they must first
-    be transferred via `stream.transfer` - which may incur a copy.
+    transfer operation. If the contents of the value are needed for host access
+    or access from a device without memory compatibility they must first
+    be transferred via a transfer operation (e.g. `stream.async.transfer`).
   }];
 
   let parameters = (ins

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

@@ -55,6 +55,7 @@ static void annotateOperandsAndResults(Operation *op,
                                        AffinityAnalysis &affinityAnalysis) {
   auto emptyArray = ArrayAttr::get(op->getContext(), {});
   SmallVector<Attribute> operandAttrs;
+  SmallVector<Attribute> operandUsageAttrs;
   for (auto operand : op->getOperands()) {
     if (isa<IREE::Stream::AffinityTypeInterface>(operand.getType())) {
       SmallVector<IREE::Stream::AffinityAttr> affinities;
@@ -64,9 +65,18 @@ static void annotateOperandsAndResults(Operation *op,
       } else {
         operandAttrs.push_back(emptyArray);
       }
+      SmallVector<IREE::Stream::AffinityAttr> usageAffinities;
+      if (affinityAnalysis.tryLookupResourceUsageAffinity(operand,
+                                                          usageAffinities)) {
+        operandUsageAttrs.push_back(ArrayAttr::get(
+            op->getContext(), llvm::to_vector_of<Attribute>(usageAffinities)));
+      } else {
+        operandUsageAttrs.push_back(emptyArray);
+      }
     }
   }
   SmallVector<Attribute> resultAttrs;
+  SmallVector<Attribute> resultUsageAttrs;
   for (auto result : op->getResults()) {
     if (isa<IREE::Stream::AffinityTypeInterface>(result.getType())) {
       SmallVector<IREE::Stream::AffinityAttr> affinities;
@@ -76,16 +86,32 @@ static void annotateOperandsAndResults(Operation *op,
       } else {
         resultAttrs.push_back(emptyArray);
       }
+      SmallVector<IREE::Stream::AffinityAttr> usageAffinities;
+      if (affinityAnalysis.tryLookupResourceUsageAffinity(result,
+                                                          usageAffinities)) {
+        resultUsageAttrs.push_back(ArrayAttr::get(
+            op->getContext(), llvm::to_vector_of<Attribute>(usageAffinities)));
+      } else {
+        resultUsageAttrs.push_back(emptyArray);
+      }
     }
   }
   if (!operandAttrs.empty()) {
     op->setAttr("stream.affinities.operands",
                 ArrayAttr::get(op->getContext(), operandAttrs));
   }
+  if (!operandUsageAttrs.empty()) {
+    op->setAttr("stream.affinities.operands.usage",
+                ArrayAttr::get(op->getContext(), operandUsageAttrs));
+  }
   if (!resultAttrs.empty()) {
     op->setAttr("stream.affinities.results",
                 ArrayAttr::get(op->getContext(), resultAttrs));
   }
+  if (!resultUsageAttrs.empty()) {
+    op->setAttr("stream.affinities.results.usage",
+                ArrayAttr::get(op->getContext(), resultUsageAttrs));
+  }
 }
 
 static void annotateFuncOp(FunctionOpInterface funcOp,