diff --git a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowInterfaces.td b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowInterfaces.td
index 5a1227e..dcb0b0f 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/IR/FlowInterfaces.td
+++ b/compiler/src/iree/compiler/Dialect/Flow/IR/FlowInterfaces.td
@@ -9,8 +9,4 @@
 
 include "iree/compiler/Dialect/Util/IR/UtilBase.td"
 
-//===----------------------------------------------------------------------===//
-// IREE::Flow::StreamableOpInterface
-//===----------------------------------------------------------------------===//
-
 #endif  // IREE_DIALECT_FLOW_INTERFACES
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/Affinity.cpp b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Affinity.cpp
new file mode 100644
index 0000000..ac3c166
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Affinity.cpp
@@ -0,0 +1,1050 @@
+// Copyright 2024 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Dialect/Stream/Analysis/Affinity.h"
+
+#include <utility>
+
+#include "iree/compiler/Dialect/Stream/IR/StreamDialect.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
+#include "iree/compiler/Dialect/Util/Analysis/DFX/Element.h"
+#include "iree/compiler/Dialect/Util/Analysis/DFX/State.h"
+#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
+#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Debug.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Diagnostics.h"
+
+#define DEBUG_TYPE "iree-util-dfx"
+
+namespace mlir::iree_compiler::IREE::Stream {
+
+//===----------------------------------------------------------------------===//
+// Utilities
+//===----------------------------------------------------------------------===//
+
+static const std::string getAffinitySetAsStr(
+    const DFX::PotentialValuesState<IREE::Stream::AffinityAttr> &state,
+    AsmState &asmState) {
+  std::string str;
+  llvm::raw_string_ostream sstream(str);
+  sstream << "pvs: ";
+  if (state.isValidState()) {
+    sstream << "[";
+    if (state.isUndefContained()) {
+      sstream << "undef, ";
+    }
+    llvm::interleaveComma(state.getAssumedSet(), sstream,
+                          [&](IREE::Stream::AffinityAttr value) {
+                            cast<Attribute>(value).print(sstream);
+                          });
+    sstream << "]";
+  } else {
+    sstream << "(invalid)";
+  }
+  sstream.flush();
+  return str;
+}
+
+//===----------------------------------------------------------------------===//
+// Analysis elements
+//===----------------------------------------------------------------------===//
+
+class ValueProducerAffinityPVS
+    : public DFX::StateWrapper<
+          DFX::PotentialValuesState<IREE::Stream::AffinityAttr>,
+          DFX::ValueElement> {
+public:
+  using BaseType =
+      DFX::StateWrapper<DFX::PotentialValuesState<IREE::Stream::AffinityAttr>,
+                        DFX::ValueElement>;
+  using BaseType::BaseType;
+
+  static ValueProducerAffinityPVS &createForPosition(const Position &pos,
+                                                     DFX::Solver &solver) {
+    return *(new (solver.getAllocator()) ValueProducerAffinityPVS(pos));
+  }
+
+  // Identity definitions.
+  const std::string getName() const override {
+    return "ValueProducerAffinityPVS";
+  }
+  const void *getID() const override { return &ID; }
+  static bool classof(const DFX::AbstractElement *element) {
+    return (element->getID() == &ID);
+  }
+  static const char ID;
+
+  const std::string getAsStr(AsmState &asmState) const override {
+    return getAffinitySetAsStr(getState(), asmState);
+  }
+
+private:
+  void initializeValue(Value value, DFX::Solver &solver) override;
+  ChangeStatus updateValue(Value value, DFX::Solver &solver) override;
+  void updateFromUse(Value value, OpOperand &operand, StateType &newState,
+                     DFX::Solver &solver);
+
+  // Operations that the value is pinned to.
+  SetVector<Operation *> pinnedOps;
+};
+const char ValueProducerAffinityPVS::ID = 0;
+
+class GlobalAffinityPVS
+    : public DFX::StateWrapper<
+          DFX::PotentialValuesState<IREE::Stream::AffinityAttr>,
+          DFX::TypedOperationElement<IREE::Util::GlobalOpInterface>> {
+public:
+  using BaseType = DFX::StateWrapper<
+      DFX::PotentialValuesState<IREE::Stream::AffinityAttr>,
+      DFX::TypedOperationElement<IREE::Util::GlobalOpInterface>>;
+  using BaseType::BaseType;
+
+  static GlobalAffinityPVS &createForPosition(const Position &pos,
+                                              DFX::Solver &solver) {
+    return *(new (solver.getAllocator()) GlobalAffinityPVS(pos));
+  }
+
+  // Identity definitions.
+  const std::string getName() const override { return "GlobalAffinityPVS"; }
+  const void *getID() const override { return &ID; }
+  static bool classof(const DFX::AbstractElement *element) {
+    return (element->getID() == &ID);
+  }
+  static const char ID;
+
+  const std::string getAsStr(AsmState &asmState) const override {
+    return getAffinitySetAsStr(getState(), asmState);
+  }
+
+private:
+  void initializeOperation(IREE::Util::GlobalOpInterface globalOp,
+                           DFX::Solver &solver) override;
+  ChangeStatus updateOperation(IREE::Util::GlobalOpInterface globalOp,
+                               DFX::Solver &solver) override;
+};
+const char GlobalAffinityPVS::ID = 0;
+
+class OpAffinityPVS : public DFX::StateWrapper<
+                          DFX::PotentialValuesState<IREE::Stream::AffinityAttr>,
+                          DFX::OperationElement> {
+public:
+  using BaseType =
+      DFX::StateWrapper<DFX::PotentialValuesState<IREE::Stream::AffinityAttr>,
+                        DFX::OperationElement>;
+  using BaseType::BaseType;
+
+  static OpAffinityPVS &createForPosition(const Position &pos,
+                                          DFX::Solver &solver) {
+    return *(new (solver.getAllocator()) OpAffinityPVS(pos));
+  }
+
+  // Identity definitions.
+  const std::string getName() const override { return "OpAffinityPVS"; }
+  const void *getID() const override { return &ID; }
+  static bool classof(const DFX::AbstractElement *element) {
+    return (element->getID() == &ID);
+  }
+  static const char ID;
+
+  const std::string getAsStr(AsmState &asmState) const override {
+    return getAffinitySetAsStr(getState(), asmState);
+  }
+
+private:
+  void initializeOperation(Operation *op, DFX::Solver &solver) override;
+  ChangeStatus updateOperation(Operation *op, DFX::Solver &solver) override;
+};
+const char OpAffinityPVS::ID = 0;
+
+//===----------------------------------------------------------------------===//
+// ValueConsumerAffinityPVS
+//===----------------------------------------------------------------------===//
+
+class ValueConsumerAffinityPVS
+    : public DFX::StateWrapper<
+          DFX::PotentialValuesState<IREE::Stream::AffinityAttr>,
+          DFX::ValueElement> {
+public:
+  using BaseType =
+      DFX::StateWrapper<DFX::PotentialValuesState<IREE::Stream::AffinityAttr>,
+                        DFX::ValueElement>;
+  using BaseType::BaseType;
+
+  static ValueConsumerAffinityPVS &createForPosition(const Position &pos,
+                                                     DFX::Solver &solver) {
+    return *(new (solver.getAllocator()) ValueConsumerAffinityPVS(pos));
+  }
+
+  // Identity definitions.
+  const std::string getName() const override {
+    return "ValueConsumerAffinityPVS";
+  }
+  const void *getID() const override { return &ID; }
+  static bool classof(const DFX::AbstractElement *element) {
+    return (element->getID() == &ID);
+  }
+  static const char ID;
+
+  const std::string getAsStr(AsmState &asmState) const override {
+    return getAffinitySetAsStr(getState(), asmState);
+  }
+
+private:
+  void initializeValue(Value value, DFX::Solver &solver) override;
+  ChangeStatus updateValue(Value value, DFX::Solver &solver) override;
+  TraversalResult updateFromUse(Value value, OpOperand &operand,
+                                StateType &newState, DFX::Solver &solver);
+};
+const char ValueConsumerAffinityPVS::ID = 0;
+
+void ValueConsumerAffinityPVS::initializeValue(Value value,
+                                               DFX::Solver &solver) {}
+
+ChangeStatus ValueConsumerAffinityPVS::updateValue(Value value,
+                                                   DFX::Solver &solver) {
+  StateType newState;
+  auto traversalResult = TraversalResult::COMPLETE;
+
+  // Walk into all consumers of the SSA value.
+  // Note that we may end up at multiple global stores of different globals
+  // by walking down through calls/branches/etc.
+  traversalResult |= solver.getExplorer().walkTransitiveUses(
+      value,
+      [&](OpOperand &operand) {
+        traversalResult |= updateFromUse(value, operand, newState, solver);
+        return WalkResult::advance();
+      },
+      (TraversalBehavior::DEFAULT | TraversalBehavior::DONT_WALK_TIED_VALUES));
+
+  if (traversalResult == TraversalResult::INCOMPLETE) {
+    // Incomplete traversal because of external call graph edges or pointers.
+    newState.unionAssumedWithUndef();
+    newState.indicatePessimisticFixpoint();
+  }
+  return DFX::clampStateAndIndicateChange(getState(), newState);
+}
+
+TraversalResult ValueConsumerAffinityPVS::updateFromUse(Value value,
+                                                        OpOperand &operand,
+                                                        StateType &newState,
+                                                        DFX::Solver &solver) {
+  // If the value is consumed by an affinity-aware op then we can directly use
+  // the affinity specified on the op. A majority of the values we care about at
+  // the stream level are consumed by affinity-aware ops and earlier in the
+  // pipeline dialects may have transfer ops that define affinities we can
+  // anchor on.
+  if (auto affinityOp =
+          dyn_cast<IREE::Stream::AffinityOpInterface>(operand.getOwner())) {
+    auto opPVS = solver.getElementFor<OpAffinityPVS>(
+        *this, Position::forOperation(operand.getOwner()),
+        DFX::Resolution::REQUIRED);
+    LLVM_DEBUG({
+      llvm::dbgs() << "[ValueConsumerAffinityPVS] value ";
+      value.printAsOperand(llvm::dbgs(), solver.getAsmState());
+      llvm::dbgs() << " affinity using consumer affinity from ";
+      operand.get().printAsOperand(llvm::dbgs(), solver.getAsmState());
+      llvm::dbgs() << " as ";
+      opPVS.print(llvm::dbgs(), solver.getAsmState());
+      llvm::dbgs() << "\n";
+    });
+    newState ^= opPVS;
+  }
+
+  // If the consumer op has the operand tied to one or more results then we walk
+  // through to track the transitive consumers. When this analysis runs we are
+  // usually still prior to baking out copy-on-write behavior so it's possible
+  // that the results of the tied operation end up in different places.
+  if (auto tiedOp = dyn_cast<IREE::Util::TiedOpInterface>(operand.getOwner())) {
+    auto tiedResults = tiedOp.getOperandTiedResults(operand.getOperandNumber());
+    for (auto tiedResult : tiedResults) {
+      auto resultPVS = solver.getElementFor<ValueConsumerAffinityPVS>(
+          *this, Position::forValue(tiedResult), DFX::Resolution::REQUIRED);
+      LLVM_DEBUG({
+        llvm::dbgs() << "[ValueConsumerAffinityPVS] value ";
+        value.printAsOperand(llvm::dbgs(), solver.getAsmState());
+        llvm::dbgs() << " affinity referencing tied operand ";
+        operand.get().printAsOperand(llvm::dbgs(), solver.getAsmState());
+        llvm::dbgs() << " result ";
+        tiedResult.printAsOperand(llvm::dbgs(), solver.getAsmState());
+        llvm::dbgs() << " as ";
+        resultPVS.print(llvm::dbgs(), solver.getAsmState());
+        llvm::dbgs() << "\n";
+      });
+      newState ^= resultPVS;
+    }
+  }
+
+  // Handle consumers that are not affinity aware - this should have any control
+  // flow ops so that we can track values that flow through the program.
+  return TypeSwitch<Operation *, TraversalResult>(operand.getOwner())
+      .Case([&](mlir::arith::SelectOp op) {
+        auto &resultPVS = solver.getElementFor<ValueConsumerAffinityPVS>(
+            *this, Position::forValue(op.getResult()),
+            DFX::Resolution::REQUIRED);
+        newState ^= resultPVS.getState();
+        return TraversalResult::COMPLETE;
+      })
+      .Case([&](mlir::BranchOpInterface op) {
+        return solver.getExplorer().walkOutgoingBranchOperandArguments(
+            op, operand.getOperandNumber(),
+            [&](Block *targetBlock, BlockArgument arg) {
+              auto &argUsage = solver.getElementFor<ValueConsumerAffinityPVS>(
+                  *this, Position::forValue(arg), DFX::Resolution::OPTIONAL);
+              newState ^= argUsage;
+              return WalkResult::advance();
+            });
+      })
+      .Case([&](mlir::scf::ForOp op) {
+        if (operand.getOperandNumber() >= op.getNumControlOperands()) {
+          int64_t blockIdx =
+              operand.getOperandNumber() - op.getNumControlOperands();
+          auto &beforeUsage = solver.getElementFor<ValueConsumerAffinityPVS>(
+              *this, Position::forValue(op.getRegionIterArg(blockIdx)),
+              DFX::Resolution::REQUIRED);
+          newState ^= beforeUsage.getState();
+        }
+        return TraversalResult::COMPLETE;
+      })
+      .Case([&](mlir::scf::WhileOp op) {
+        auto &beforeUsage = solver.getElementFor<ValueConsumerAffinityPVS>(
+            *this,
+            Position::forValue(
+                op.getBeforeBody()->getArgument(operand.getOperandNumber())),
+            DFX::Resolution::REQUIRED);
+        newState ^= beforeUsage.getState();
+        return TraversalResult::COMPLETE;
+      })
+      .Case([&](mlir::scf::ConditionOp op) {
+        auto &parentUsage = solver.getElementFor<ValueConsumerAffinityPVS>(
+            *this,
+            Position::forValue(
+                op->getParentOp()->getResult(operand.getOperandNumber() - 1)),
+            DFX::Resolution::REQUIRED);
+        newState ^= parentUsage.getState();
+        if (auto whileOp =
+                dyn_cast_or_null<mlir::scf::WhileOp>(op->getParentOp())) {
+          auto value = Position::forValue(
+              whileOp.getAfter().getArgument(operand.getOperandNumber() - 1));
+          auto &valueUsage = solver.getElementFor<ValueConsumerAffinityPVS>(
+              *this, value, DFX::Resolution::REQUIRED);
+          newState ^= valueUsage.getState();
+        }
+        return TraversalResult::COMPLETE;
+      })
+      .Case([&](mlir::scf::YieldOp op) {
+        if (isa<mlir::scf::IfOp>(op->getParentOp())) {
+          auto &operandUsage = solver.getElementFor<ValueConsumerAffinityPVS>(
+              *this,
+              Position::forValue(op->getOperand(operand.getOperandNumber())),
+              DFX::Resolution::REQUIRED);
+          newState ^= operandUsage.getState();
+          auto &parentUsage = solver.getElementFor<ValueConsumerAffinityPVS>(
+              *this,
+              Position::forValue(
+                  op->getParentOp()->getResult(operand.getOperandNumber())),
+              DFX::Resolution::REQUIRED);
+          newState ^= parentUsage.getState();
+          return TraversalResult::COMPLETE;
+        } else if (auto whileOp =
+                       dyn_cast<mlir::scf::WhileOp>(op->getParentOp())) {
+          auto value = Position::forValue(
+              whileOp.getBefore().getArgument(operand.getOperandNumber()));
+          auto &valueUsage = solver.getElementFor<ValueConsumerAffinityPVS>(
+              *this, value, DFX::Resolution::REQUIRED);
+          newState ^= valueUsage.getState();
+          auto &parentUsage = solver.getElementFor<ValueConsumerAffinityPVS>(
+              *this,
+              Position::forValue(
+                  whileOp->getResult(operand.getOperandNumber())),
+              DFX::Resolution::REQUIRED);
+          newState ^= parentUsage.getState();
+          return TraversalResult::COMPLETE;
+        } else if (auto forOp = dyn_cast<mlir::scf::ForOp>(op->getParentOp())) {
+          auto value = Position::forValue(
+              forOp.getRegionIterArg(operand.getOperandNumber()));
+          auto &valueUsage = solver.getElementFor<ValueConsumerAffinityPVS>(
+              *this, value, DFX::Resolution::REQUIRED);
+          newState ^= valueUsage.getState();
+          auto &parentUsage = solver.getElementFor<ValueConsumerAffinityPVS>(
+              *this,
+              Position::forValue(forOp->getResult(operand.getOperandNumber())),
+              DFX::Resolution::REQUIRED);
+          newState ^= parentUsage.getState();
+          return TraversalResult::COMPLETE;
+        } else {
+          assert(false && "unhandled scf yield parent");
+          return TraversalResult::INCOMPLETE;
+        }
+      })
+      .Case([&](IREE::Util::ReturnOp op) {
+        return solver.getExplorer().walkIncomingCalls(
+            op->getParentOfType<mlir::CallableOpInterface>(),
+            [&](mlir::CallOpInterface callOp) {
+              auto &argUsage = solver.getElementFor<ValueConsumerAffinityPVS>(
+                  *this,
+                  Position::forValue(
+                      callOp->getResult(operand.getOperandNumber())),
+                  DFX::Resolution::OPTIONAL);
+              getState() ^= argUsage;
+              return WalkResult::advance();
+            });
+      })
+      .Case([&](IREE::Util::OptimizationBarrierOp op) {
+        auto &resultPVS = solver.getElementFor<ValueConsumerAffinityPVS>(
+            *this, Position::forValue(op.getResult(operand.getOperandNumber())),
+            DFX::Resolution::REQUIRED);
+        newState ^= resultPVS.getState();
+        return TraversalResult::COMPLETE;
+      })
+      .Case([&](IREE::Util::GlobalStoreOpInterface op) {
+        auto *globalInfo =
+            solver.getExplorer().queryGlobalInfoFrom(op.getGlobalName(), op);
+        auto &globalPVS = solver.getElementFor<GlobalAffinityPVS>(
+            *this, Position::forOperation(globalInfo->op),
+            DFX::Resolution::REQUIRED);
+        newState ^= globalPVS.getState();
+        return TraversalResult::COMPLETE;
+      })
+      .Default([&](Operation *op) { return TraversalResult::COMPLETE; });
+}
+
+//===----------------------------------------------------------------------===//
+// ValueProducerAffinityPVS
+//===----------------------------------------------------------------------===//
+
+void ValueProducerAffinityPVS::initializeValue(Value value,
+                                               DFX::Solver &solver) {
+  solver.getExplorer().walkDefiningOps(value, [&](OpResult result) {
+    if (!isa<IREE::Stream::AffinityTypeInterface>(result.getType())) {
+      return WalkResult::skip();
+    }
+    if (auto affinityOp =
+            dyn_cast_if_present<IREE::Stream::AffinityOpInterface>(
+                result.getOwner())) {
+      if (affinityOp.pinsValueAffinity()) {
+        pinnedOps.insert(result.getOwner());
+      }
+    }
+    return WalkResult::advance();
+  });
+  solver.getExplorer().walkTransitiveUses(value, [&](OpOperand &operand) {
+    if (!isa<IREE::Stream::AffinityTypeInterface>(operand.get().getType())) {
+      return WalkResult::skip();
+    }
+    if (auto affinityOp =
+            dyn_cast_if_present<IREE::Stream::AffinityOpInterface>(
+                operand.getOwner())) {
+      if (affinityOp.pinsValueAffinity()) {
+        pinnedOps.insert(operand.getOwner());
+      }
+    }
+    return WalkResult::advance();
+  });
+}
+
+ChangeStatus ValueProducerAffinityPVS::updateValue(Value value,
+                                                   DFX::Solver &solver) {
+  StateType newState;
+
+  // If there are any ops that produce the value and pin to a specific affinity
+  // then we take those directly and ignore all others.
+  if (!pinnedOps.empty()) {
+    for (auto pinnedOp : pinnedOps) {
+      auto &opPVS = solver.getElementFor<OpAffinityPVS>(
+          *this, Position::forOperation(pinnedOp), DFX::Resolution::REQUIRED);
+      newState ^= opPVS;
+    }
+    return DFX::clampStateAndIndicateChange(getState(), newState);
+  }
+
+  // We special case some ops that act as barriers in the program. This prevents
+  // us from walking past boundaries that are not profitable to do so with; for
+  // example, globals are usually stored in independent contexts from where they
+  // are consumed.
+  if (auto barrierOp = dyn_cast_if_present<IREE::Util::OptimizationBarrierOp>(
+          value.getDefiningOp())) {
+    auto operand =
+        barrierOp.getOperand(cast<OpResult>(value).getResultNumber());
+    auto operandPVS = solver.getElementFor<ValueProducerAffinityPVS>(
+        *this, Position::forValue(operand), DFX::Resolution::REQUIRED);
+    LLVM_DEBUG({
+      llvm::dbgs() << "[ValueProducerAffinityPVS] value ";
+      value.printAsOperand(llvm::dbgs(), solver.getAsmState());
+      llvm::dbgs() << " affinity using barrier op operand as ";
+      operandPVS.print(llvm::dbgs(), solver.getAsmState());
+      llvm::dbgs() << "\n";
+    });
+    newState ^= operandPVS;
+    return DFX::clampStateAndIndicateChange(getState(), newState);
+  } else if (auto loadOp =
+                 dyn_cast_if_present<IREE::Util::GlobalLoadOpInterface>(
+                     value.getDefiningOp())) {
+    auto *globalInfo = solver.getExplorer().queryGlobalInfoFrom(
+        loadOp.getGlobalName(), loadOp);
+    auto &globalPVS = solver.getElementFor<GlobalAffinityPVS>(
+        *this, Position::forOperation(globalInfo->op),
+        DFX::Resolution::REQUIRED);
+    LLVM_DEBUG({
+      llvm::dbgs() << "[ValueProducerAffinityPVS] value ";
+      value.printAsOperand(llvm::dbgs(), solver.getAsmState());
+      llvm::dbgs() << " affinity using global op affinity from "
+                   << loadOp.getGlobalName() << " as ";
+      globalPVS.print(llvm::dbgs(), solver.getAsmState());
+      llvm::dbgs() << "\n";
+    });
+    newState ^= globalPVS.getState();
+    return DFX::clampStateAndIndicateChange(getState(), newState);
+  }
+
+  // Walk the program up into any possible producers of the value.
+  auto traversalResult = TraversalResult::COMPLETE;
+  traversalResult |= solver.getExplorer().walkDefiningOps(
+      value,
+      [&](OpResult result) {
+        if (isa<CallOpInterface>(result.getOwner())) {
+          return WalkResult::advance();
+        }
+
+        // If coming from an affinity-aware op that pins the value storage to a
+        // particular affinity that overrides all other logic.
+        if (auto affinityOp =
+                dyn_cast_if_present<IREE::Stream::AffinityOpInterface>(
+                    result.getDefiningOp())) {
+          if (affinityOp.pinsValueAffinity()) {
+            auto &opPVS = solver.getElementFor<OpAffinityPVS>(
+                *this, Position::forOperation(affinityOp),
+                DFX::Resolution::REQUIRED);
+            LLVM_DEBUG({
+              llvm::dbgs() << "[ValueProducerAffinityPVS] value ";
+              value.printAsOperand(llvm::dbgs(), solver.getAsmState());
+              llvm::dbgs() << " affinity using assuming pinned affinity from ";
+              result.printAsOperand(llvm::dbgs(), solver.getAsmState());
+              llvm::dbgs() << " as ";
+              opPVS.print(llvm::dbgs(), solver.getAsmState());
+              llvm::dbgs() << "\n";
+            });
+            newState ^= opPVS;
+            newState.indicateOptimisticFixpoint();
+            return WalkResult::advance();
+          }
+        }
+
+        // If the result value is tied to an operand of the defining op then
+        // inherit the operand affinity.
+        if (auto tiedOp = dyn_cast_if_present<IREE::Util::TiedOpInterface>(
+                result.getDefiningOp())) {
+          auto operand = tiedOp.getTiedResultOperand(result);
+          if (operand) {
+            auto &valuePVS = solver.getElementFor<ValueProducerAffinityPVS>(
+                *this, Position::forValue(operand), DFX::Resolution::OPTIONAL);
+            LLVM_DEBUG({
+              llvm::dbgs() << "[ValueProducerAffinityPVS] value ";
+              value.printAsOperand(llvm::dbgs(), solver.getAsmState());
+              llvm::dbgs() << " affinity referencing tied operand ";
+              operand.printAsOperand(llvm::dbgs(), solver.getAsmState());
+              llvm::dbgs() << " as ";
+              valuePVS.print(llvm::dbgs(), solver.getAsmState());
+              llvm::dbgs() << "\n";
+            });
+            newState ^= valuePVS;
+            return WalkResult::advance();
+          }
+        }
+
+        // If the value is produced by the defining op then assume that the
+        // execution affinity dictates the result affinity.
+        if (auto affinityOp =
+                dyn_cast_if_present<IREE::Stream::AffinityOpInterface>(
+                    result.getDefiningOp())) {
+          auto &opPVS = solver.getElementFor<OpAffinityPVS>(
+              *this, Position::forOperation(result.getOwner()),
+              DFX::Resolution::OPTIONAL);
+          LLVM_DEBUG({
+            llvm::dbgs() << "[ValueProducerAffinityPVS] value ";
+            value.printAsOperand(llvm::dbgs(), solver.getAsmState());
+            llvm::dbgs() << " affinity using op affinity from result ";
+            result.printAsOperand(llvm::dbgs(), solver.getAsmState());
+            llvm::dbgs() << " as ";
+            opPVS.print(llvm::dbgs(), solver.getAsmState());
+            llvm::dbgs() << "\n";
+          });
+          newState ^= opPVS;
+          return WalkResult::advance();
+        }
+
+        // Special handling for specific ops.
+        TypeSwitch<Operation *>(result.getOwner())
+            .Case<IREE::Util::GlobalLoadOpInterface>([&](auto loadOp) {
+              auto *globalInfo = solver.getExplorer().queryGlobalInfoFrom(
+                  loadOp.getGlobalName(), loadOp);
+              auto &globalPVS = solver.getElementFor<GlobalAffinityPVS>(
+                  *this, Position::forOperation(globalInfo->op),
+                  DFX::Resolution::REQUIRED);
+              LLVM_DEBUG({
+                llvm::dbgs() << "[ValueProducerAffinityPVS] value ";
+                value.printAsOperand(llvm::dbgs(), solver.getAsmState());
+                llvm::dbgs()
+                    << " affinity using global op affinity from result ";
+                result.printAsOperand(llvm::dbgs(), solver.getAsmState());
+                llvm::dbgs() << " as ";
+                globalPVS.print(llvm::dbgs(), solver.getAsmState());
+                llvm::dbgs() << "\n";
+              });
+              newState ^= globalPVS.getState();
+            })
+            .Case<mlir::arith::SelectOp>([&](auto op) {
+              auto &truePVS = solver.getElementFor<ValueProducerAffinityPVS>(
+                  *this, Position::forValue(op.getTrueValue()),
+                  DFX::Resolution::REQUIRED);
+              newState ^= truePVS.getState();
+              auto &falsePVS = solver.getElementFor<ValueProducerAffinityPVS>(
+                  *this, Position::forValue(op.getFalseValue()),
+                  DFX::Resolution::REQUIRED);
+              newState ^= falsePVS.getState();
+            })
+            .Default([&](auto op) {
+              auto valuePVS = solver.getElementFor<ValueProducerAffinityPVS>(
+                  *this, Position::forValue(result), DFX::Resolution::OPTIONAL);
+              newState ^= valuePVS;
+            });
+        return WalkResult::advance();
+      },
+      (TraversalBehavior::DEFAULT | TraversalBehavior::DONT_WALK_TIED_VALUES));
+
+  if (traversalResult == TraversalResult::INCOMPLETE) {
+    // Incomplete traversal because of external call graph edges or pointers.
+    newState.unionAssumedWithUndef();
+    newState.indicatePessimisticFixpoint();
+  }
+  return DFX::clampStateAndIndicateChange(getState(), newState);
+}
+
+//===----------------------------------------------------------------------===//
+// GlobalAffinityPVS
+//===----------------------------------------------------------------------===//
+
+void GlobalAffinityPVS::initializeOperation(
+    IREE::Util::GlobalOpInterface globalOp, DFX::Solver &solver) {
+  // If an affinity is explicitly specified we take that over all analysis.
+  if (auto affinityAttr = IREE::Stream::AffinityAttr::lookup(globalOp)) {
+    LLVM_DEBUG({
+      llvm::dbgs() << "[GlobalAffinityPVS] global @"
+                   << globalOp.getGlobalName().getValue()
+                   << " affinity explicitly specified as ";
+      affinityAttr.print(llvm::dbgs());
+      llvm::dbgs() << "\n";
+    });
+    unionAssumed(affinityAttr);
+    indicateOptimisticFixpoint();
+    return;
+  }
+}
+
+ChangeStatus
+GlobalAffinityPVS::updateOperation(IREE::Util::GlobalOpInterface globalOp,
+                                   DFX::Solver &solver) {
+  StateType newState;
+  auto traversalResult = TraversalResult::COMPLETE;
+
+  const auto *globalInfo = solver.getExplorer().getGlobalInfo(globalOp);
+  if (globalInfo->isIndirect) {
+    traversalResult = TraversalResult::INCOMPLETE;
+  }
+
+  // Traverse all transitive uses of the global.
+  // We try to place globals where they are used as the common case is weights
+  // or parameters that are read more frequently than they are written.
+  // The reasoning is that if there are more writes than reads there's unneeded
+  // work being done and otherwise there's always at least one read per write
+  // or more reads than writes.
+  bool anyLoads = false;
+  for (auto loadOp : globalInfo->getLoads()) {
+    anyLoads = true;
+    auto &valuePVS = solver.getElementFor<ValueConsumerAffinityPVS>(
+        *this, Position::forValue(loadOp.getLoadedGlobalValue()),
+        DFX::Resolution::OPTIONAL);
+    if (valuePVS.isValidState()) {
+      newState ^= valuePVS;
+    }
+  }
+
+  // If there were no loads then take the affinity from stores.
+  // This is not common but can arise in tests or where the globals may be used
+  // to model side-effecting behavior.
+  if (!anyLoads) {
+    for (auto storeOp : globalInfo->getStores()) {
+      auto &valuePVS = solver.getElementFor<ValueProducerAffinityPVS>(
+          *this, Position::forValue(storeOp.getStoredGlobalValue()),
+          DFX::Resolution::OPTIONAL);
+      if (valuePVS.isValidState()) {
+        newState ^= valuePVS;
+      }
+    }
+  }
+
+  if (traversalResult == TraversalResult::INCOMPLETE) {
+    // Incomplete traversal because of external call graph edges or pointers.
+    newState.unionAssumedWithUndef();
+    newState.indicatePessimisticFixpoint();
+  }
+  return DFX::clampStateAndIndicateChange(getState(), newState);
+}
+
+//===----------------------------------------------------------------------===//
+// OpAffinityPVS
+//===----------------------------------------------------------------------===//
+
+void OpAffinityPVS::initializeOperation(Operation *op, DFX::Solver &solver) {
+  // If an affinity is explicitly specified we take that over all analysis.
+  if (auto affinityAttr = IREE::Stream::AffinityAttr::lookup(op)) {
+    LLVM_DEBUG({
+      llvm::dbgs() << "[OpAffinityPVS] op ";
+      op->getName().print(llvm::dbgs());
+      llvm::dbgs() << " affinity explicitly specified as ";
+      affinityAttr.print(llvm::dbgs());
+      llvm::dbgs() << "\n";
+    });
+    unionAssumed(affinityAttr);
+    indicateOptimisticFixpoint();
+    return;
+  }
+}
+
+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;
+      }
+    }
+  }
+
+  return DFX::clampStateAndIndicateChange(getState(), newState);
+}
+
+//===----------------------------------------------------------------------===//
+// AffinityAnalysis
+//===----------------------------------------------------------------------===//
+
+// Tries to find a default affinity specified on an ancestor of |fromOp| and
+// adds it to |affinities|. Returns true if an affinity was found.
+static bool tryLookupDefaultAffinity(
+    Operation *fromOp,
+    SmallVectorImpl<IREE::Stream::AffinityAttr> &affinities) {
+  while (fromOp) {
+    auto affinityAttr = fromOp->getAttrOfType<IREE::Stream::AffinityAttr>(
+        "stream.affinity.default");
+    if (affinityAttr) {
+      affinities.push_back(affinityAttr);
+      return true;
+    }
+    fromOp = fromOp->getParentOp();
+  }
+  return false;
+}
+
+// Returns the first affinity if all affinities are compatible and otherwise
+// returns nullptr.
+static IREE::Stream::AffinityAttr
+trySelectLeadAffinity(ArrayRef<IREE::Stream::AffinityAttr> affinities) {
+  if (affinities.empty()) {
+    return {};
+  }
+  auto leadAffinityAttr = affinities.front();
+  for (size_t i = 1; i < affinities.size(); ++i) {
+    if (!IREE::Stream::AffinityAttr::areCompatible(affinities[i],
+                                                   leadAffinityAttr)) {
+      return {};
+    }
+  }
+  return leadAffinityAttr;
+}
+
+// Sorts |affinities| in the natural affinity sort order.
+// We unfortunately have to do this as the PVS elements we source from are
+// unsorted.
+static void
+sortAffinities(SmallVectorImpl<IREE::Stream::AffinityAttr> &affinities) {
+  // HACK: this should probably do a type id ordering followed by a
+  // type-specific ordering (interface compare method?). We just need this to be
+  // stable as the affinities come from multiple DenseSets that have run-to-run
+  // ordering variance. This is very inefficient but is only used when there are
+  // multiple possible affinities and we try to avoid that anyway.
+  if (affinities.size() <= 1) {
+    return;
+  }
+  llvm::stable_sort(affinities, [](IREE::Stream::AffinityAttr lhs,
+                                   IREE::Stream::AffinityAttr rhs) {
+    std::string lhsStr;
+    llvm::raw_string_ostream lhsStream(lhsStr);
+    lhs.print(lhsStream);
+    std::string rhsStr;
+    llvm::raw_string_ostream rhsStream(rhsStr);
+    rhs.print(rhsStream);
+    return lhsStr < rhsStr;
+  });
+}
+
+AffinityAnalysis::AffinityAnalysis(Operation *rootOp)
+    : explorer(rootOp, TraversalAction::RECURSE), solver(explorer, allocator) {
+  explorer.setOpInterfaceAction<mlir::FunctionOpInterface>(
+      TraversalAction::RECURSE);
+
+  explorer.setDialectAction<mlir::scf::SCFDialect>(TraversalAction::RECURSE);
+
+  explorer.setDialectAction<IREE::Stream::StreamDialect>(
+      TraversalAction::RECURSE);
+  explorer.setOpAction<IREE::Stream::ExecutableOp>(TraversalAction::IGNORE);
+
+  explorer.initialize();
+}
+
+AffinityAnalysis::~AffinityAnalysis() = default;
+
+IREE::Stream::AffinityAttr
+AffinityAnalysis::lookupGlobalAffinity(Operation *op) {
+  SmallVector<IREE::Stream::AffinityAttr> affinities;
+  if (!tryLookupGlobalAffinity(op, affinities) || affinities.empty()) {
+    return {};
+  }
+  if (affinities.size() == 1) {
+    return affinities.front();
+  }
+  return trySelectLeadAffinity(affinities);
+}
+
+bool AffinityAnalysis::tryLookupGlobalAffinity(
+    Operation *op, SmallVectorImpl<IREE::Stream::AffinityAttr> &affinities) {
+  auto globalPVS =
+      solver.lookupElementFor<GlobalAffinityPVS>(Position::forOperation(op));
+  if (!globalPVS || !globalPVS->isValidState() ||
+      globalPVS->isUndefContained()) {
+    // Analysis failed.
+    return false;
+  }
+  if (globalPVS->getAssumedSet().empty()) {
+    // Analysis completed but no affinity was specified; try to find a default.
+    return tryLookupDefaultAffinity(op, affinities);
+  }
+  for (auto affinityAttr : globalPVS->getAssumedSet()) {
+    affinities.push_back(affinityAttr);
+  }
+  sortAffinities(affinities);
+  return true;
+}
+
+IREE::Stream::AffinityAttr
+AffinityAnalysis::lookupExecutionAffinity(Operation *op) {
+  SmallVector<IREE::Stream::AffinityAttr> affinities;
+  if (!tryLookupExecutionAffinity(op, affinities) || affinities.empty()) {
+    return {};
+  }
+  if (affinities.size() == 1) {
+    return affinities.front();
+  }
+  return trySelectLeadAffinity(affinities);
+}
+
+bool AffinityAnalysis::tryLookupExecutionAffinity(
+    Operation *op, SmallVectorImpl<IREE::Stream::AffinityAttr> &affinities) {
+  auto opPVS =
+      solver.lookupElementFor<OpAffinityPVS>(Position::forOperation(op));
+  if (!opPVS || !opPVS->isValidState() || opPVS->isUndefContained()) {
+    // Analysis failed.
+    return false;
+  }
+  if (opPVS->getAssumedSet().empty()) {
+    // Analysis completed but no affinity was specified; try to find a default.
+    return tryLookupDefaultAffinity(op, affinities);
+  }
+  for (auto affinityAttr : opPVS->getAssumedSet()) {
+    affinities.push_back(affinityAttr);
+  }
+  sortAffinities(affinities);
+  return true;
+}
+
+IREE::Stream::AffinityAttr
+AffinityAnalysis::inferExecutionAffinity(Operation *op) {
+  SmallVector<IREE::Stream::AffinityAttr> affinities;
+  if (!tryInferExecutionAffinity(op, affinities) || affinities.empty()) {
+    return {};
+  }
+  if (affinities.size() == 1) {
+    return affinities.front();
+  }
+  return trySelectLeadAffinity(affinities);
+}
+
+bool AffinityAnalysis::tryInferExecutionAffinity(
+    Operation *op, SmallVectorImpl<IREE::Stream::AffinityAttr> &affinities) {
+  if (auto affinityOp = dyn_cast<IREE::Stream::AffinityOpInterface>(op)) {
+    return tryLookupExecutionAffinity(op, affinities);
+  }
+  DFX::PotentialValuesState<IREE::Stream::AffinityAttr> opPVS;
+  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.lookupElementFor<ValueProducerAffinityPVS>(
+            Position::forValue(operand), nullptr, DFX::Resolution::REQUIRED);
+        if (valuePVS && valuePVS->isValidState()) {
+          opPVS.unionAssumed(valuePVS->getState());
+        } else {
+          return false;
+        }
+      }
+    }
+  } else {
+    for (auto result : op->getResults()) {
+      if (isa<IREE::Stream::AffinityTypeInterface>(result.getType())) {
+        auto valuePVS = solver.lookupElementFor<ValueConsumerAffinityPVS>(
+            Position::forValue(result), nullptr, DFX::Resolution::REQUIRED);
+        if (valuePVS && valuePVS->isValidState()) {
+          opPVS.unionAssumed(valuePVS->getState());
+        } else {
+          return false;
+        }
+      }
+    }
+  }
+  if (!opPVS.isValidState() || opPVS.isUndefContained()) {
+    // Analysis failed.
+    return false;
+  }
+  if (opPVS.getAssumedSet().empty()) {
+    // Analysis completed but no affinity was specified; try to find a default.
+    return tryLookupDefaultAffinity(op, affinities);
+  }
+  for (auto affinityAttr : opPVS.getAssumedSet()) {
+    affinities.push_back(affinityAttr);
+  }
+  sortAffinities(affinities);
+  return true;
+}
+
+IREE::Stream::AffinityAttr
+AffinityAnalysis::lookupResourceAffinity(Value value) {
+  SmallVector<IREE::Stream::AffinityAttr> affinities;
+  if (!tryLookupResourceAffinity(value, affinities) || affinities.empty()) {
+    return {};
+  }
+  if (affinities.size() == 1) {
+    return affinities.front();
+  }
+  return trySelectLeadAffinity(affinities);
+}
+
+bool AffinityAnalysis::tryLookupResourceAffinity(
+    Value value, SmallVectorImpl<IREE::Stream::AffinityAttr> &affinities) {
+  auto valuePVS = solver.lookupElementFor<ValueProducerAffinityPVS>(
+      Position::forValue(value));
+  if (!valuePVS || !valuePVS->isValidState() || valuePVS->isUndefContained()) {
+    // Analysis failed.
+    return false;
+  }
+  if (valuePVS->getAssumedSet().empty()) {
+    // Analysis completed but no affinity was specified; try to find a default.
+    return tryLookupDefaultAffinity(value.getParentBlock()->getParentOp(),
+                                    affinities);
+  }
+  for (auto affinityAttr : valuePVS->getAssumedSet()) {
+    affinities.push_back(affinityAttr);
+  }
+  sortAffinities(affinities);
+  return true;
+}
+
+LogicalResult AffinityAnalysis::run() {
+  // Initialize globals so that we can assign them affinity.
+  explorer.forEachGlobal([&](const auto *globalInfo) {
+    if (isa<IREE::Stream::AffinityTypeInterface>(
+            globalInfo->op.getGlobalType())) {
+      solver.getOrCreateElementFor<GlobalAffinityPVS>(
+          Position::forOperation(globalInfo->op));
+    }
+  });
+
+  // Initialize op execution affinities for any ops that use tracked types.
+  //
+  // TODO(benvanik): avoid doing this initialization for the entire module and
+  // instead rely on DFX to automatically populate the required abstract values.
+  // There's some missing logic in the element initialization, though, and by
+  // initializing all values we side-step that and work with test programs that
+  // may not have I/O edges that we could easily latch on to here.
+  explorer.forEachFunctionLikeOp([&](FunctionOpInterface funcOp) {
+    for (auto &block : funcOp.getBlocks()) {
+      for (auto arg : block.getArguments()) {
+        if (isa<IREE::Stream::AffinityTypeInterface>(arg.getType())) {
+          solver.getOrCreateElementFor<ValueProducerAffinityPVS>(
+              Position::forValue(arg));
+        }
+      }
+    }
+    funcOp.walk([&](Operation *op) {
+      if (auto regionOp = dyn_cast<RegionBranchOpInterface>(op)) {
+        for (auto &region : regionOp->getRegions()) {
+          for (auto arg : region.getArguments()) {
+            if (isa<IREE::Stream::AffinityTypeInterface>(arg.getType())) {
+              solver.getOrCreateElementFor<ValueProducerAffinityPVS>(
+                  Position::forValue(arg));
+            }
+          }
+        }
+      }
+      if (auto affinityOp = dyn_cast<IREE::Stream::AffinityOpInterface>(op)) {
+        solver.getOrCreateElementFor<OpAffinityPVS>(Position::forOperation(op));
+      }
+      for (auto result : op->getResults()) {
+        if (isa<IREE::Stream::AffinityTypeInterface>(result.getType())) {
+          solver.getOrCreateElementFor<ValueProducerAffinityPVS>(
+              Position::forValue(result));
+        }
+      }
+    });
+  });
+
+  if (failed(solver.run())) {
+    return failure(); // did not converge
+  }
+
+  LLVM_DEBUG({
+    llvm::dbgs()
+        << "\n\n[Analysis] affinity analysis results for the whole module:\n";
+    solver.print(llvm::dbgs());
+    llvm::dbgs() << "\n";
+  });
+
+  return success();
+}
+
+} // namespace mlir::iree_compiler::IREE::Stream
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/Affinity.h b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Affinity.h
new file mode 100644
index 0000000..3642a53
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Affinity.h
@@ -0,0 +1,102 @@
+// Copyright 2024 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_COMPILER_DIALECT_STREAM_ANALYSIS_AFFINITY_H_
+#define IREE_COMPILER_DIALECT_STREAM_ANALYSIS_AFFINITY_H_
+
+#include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
+#include "iree/compiler/Dialect/Util/Analysis/DFX/Solver.h"
+#include "iree/compiler/Dialect/Util/Analysis/Explorer.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Diagnostics.h"
+
+namespace mlir::iree_compiler::IREE::Stream {
+
+//===----------------------------------------------------------------------===//
+// Affinity analysis
+//===----------------------------------------------------------------------===//
+
+// Performs whole-program analysis of resource and tensor value affinity.
+// All `!stream.resource` and `tensor` SSA values will be analyzed and their
+// affinities where used will be available for querying via the lookup
+// functions.
+class AffinityAnalysis {
+public:
+  explicit AffinityAnalysis(Operation *rootOp);
+  ~AffinityAnalysis();
+
+  // Runs analysis and populates the resource usage map.
+  // May fail if analysis cannot be completed due to unsupported or unknown IR.
+  LogicalResult run();
+
+  // Returns the affinity of the global |op| based on its loads.
+  // The global storage should be allocated with this affinity and available for
+  // fast access from any compatible affinity.
+  //
+  // If an explicit affinity is provided via a stream.affinity attribute then
+  // that will be used in place of analysis. If there are more than one consumer
+  // (such as multiple loads) with differing affinities or analysis fails then
+  // no affinity is returned. If all affinities are compatible one will be
+  // chosen in an unspecified way.
+  IREE::Stream::AffinityAttr lookupGlobalAffinity(Operation *op);
+
+  // Populates all potential affinities of the global |op| in |affinities|.
+  // Returns false if analysis failed and the set of affinities is unknown.
+  bool tryLookupGlobalAffinity(
+      Operation *op, SmallVectorImpl<IREE::Stream::AffinityAttr> &affinities);
+
+  // Returns the affinity of the executable |op| based on the op-specific rules
+  // as to whether its operands or results control placement. The operation
+  // should be scheduled to execute with this affinity and efficiently consume
+  // or produce resources that share a compatible affinity.
+  //
+  // If an explicit affinity is provided via stream.affinity attrs or the
+  // affinity op interface then that will be used in place of analysis. If there
+  // are multiple possible affinities or analysis fails no affinity is returned.
+  // If all affinities are compatible one will be chosen in an unspecified way.
+  IREE::Stream::AffinityAttr lookupExecutionAffinity(Operation *op);
+
+  // Populates all potential execution affinities of |op| in |affinities|.
+  // Returns false if analysis failed and the set of affinities is unknown.
+  bool tryLookupExecutionAffinity(
+      Operation *op, SmallVectorImpl<IREE::Stream::AffinityAttr> &affinities);
+
+  // Returns the affinity of |op| as if it were executable even if it is not.
+  // This relies on analysis of operands and results having resolved and
+  // otherwise returns nullptr indicating the op has no assumed affinity.
+  IREE::Stream::AffinityAttr inferExecutionAffinity(Operation *op);
+
+  // Populates all inferred potential execution affinities of |op| in
+  // |affinities|. This relies on analysis of operands and results having
+  // resolved and otherwise returns nullptr indicating the op has no assumed
+  // affinity.
+  // Returns false if analysis failed and the set of affinities is unknown.
+  bool tryInferExecutionAffinity(
+      Operation *op, SmallVectorImpl<IREE::Stream::AffinityAttr> &affinities);
+
+  // Returns the affinity of |value| based on its producers.
+  // The resource should be allocated with this affinity and be usable by any
+  // compatible affinity.
+  //
+  // If there are more than one producer of the value (such as multiple callers)
+  // with differing affinities or analysis fails then no affinity is returned.
+  // 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|.
+  // Returns false if analysis failed and the set of affinities is unknown.
+  bool tryLookupResourceAffinity(
+      Value value, SmallVectorImpl<IREE::Stream::AffinityAttr> &affinities);
+
+private:
+  Explorer explorer;
+  llvm::BumpPtrAllocator allocator;
+  DFX::Solver solver;
+};
+
+} // namespace mlir::iree_compiler::IREE::Stream
+
+#endif // IREE_COMPILER_DIALECT_STREAM_ANALYSIS_AFFINITY_H_
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Stream/Analysis/BUILD.bazel
index 4e1421b..3cbb5b5 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Analysis/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Stream/Analysis/BUILD.bazel
@@ -15,12 +15,14 @@
 iree_compiler_cc_library(
     name = "Analysis",
     srcs = [
+        "Affinity.cpp",
         "Partitioning.cpp",
         "Partitioning/ReferencePartitioning.cpp",
         "ResourceHazards.cpp",
         "ResourceUsage.cpp",
     ],
     hdrs = [
+        "Affinity.h",
         "Partitioning.h",
         "ResourceHazards.h",
         "ResourceUsage.h",
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Stream/Analysis/CMakeLists.txt
index f1b0fc8..c2dd74c 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Analysis/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Stream/Analysis/CMakeLists.txt
@@ -14,10 +14,12 @@
   NAME
     Analysis
   HDRS
+    "Affinity.h"
     "Partitioning.h"
     "ResourceHazards.h"
     "ResourceUsage.h"
   SRCS
+    "Affinity.cpp"
     "Partitioning.cpp"
     "Partitioning/ReferencePartitioning.cpp"
     "ResourceHazards.cpp"
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/ResourceUsage.cpp b/compiler/src/iree/compiler/Dialect/Stream/Analysis/ResourceUsage.cpp
index 1708782..4ff656c 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Analysis/ResourceUsage.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Analysis/ResourceUsage.cpp
@@ -17,7 +17,6 @@
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
@@ -416,10 +415,14 @@
           // TODO(benvanik): remove kFavorTransients.
           bool isSourceExternal = !sourceUsage.isAssumed(NOT_EXTERNAL);
           bool isTargetInternal = isAssumed(NOT_EXTERNAL);
-          if (kFavorTransients && isSourceExternal && isTargetInternal) {
+          bool deviceChange =
+              op.getSourceAffinityAttr() != op.getResultAffinityAttr();
+          if ((kFavorTransients || deviceChange) && isSourceExternal &&
+              isTargetInternal) {
             LLVM_DEBUG({
-              llvm::dbgs() << "[ValueResourceUsage] skipping forward prop of "
-                              "external into internal: ";
+              llvm::dbgs()
+                  << "[ValueResourceUsage] skipping forward prop of external "
+                     "into internal due to kFavorTransients/device-change: ";
               op.print(llvm::dbgs(), solver.getAsmState());
               llvm::dbgs() << "\n";
             });
@@ -529,7 +532,6 @@
               *this,
               Position::forValue(op.getBeforeBody()->getArgument(operandIdx)),
               DFX::Resolution::REQUIRED);
-
           getState() ^= beforeUsage.getState();
         })
         .Case([&](mlir::scf::ConditionOp op) {
@@ -562,29 +564,30 @@
                 Position::forValue(op->getParentOp()->getResult(operandIdx)),
                 DFX::Resolution::REQUIRED);
             getState() ^= parentUsage.getState();
-          } else if (auto whileOp =
-                         dyn_cast_or_null<scf::WhileOp>(op->getParentOp())) {
+          } else if (auto whileOp = dyn_cast<scf::WhileOp>(op->getParentOp())) {
             auto value =
                 Position::forValue(whileOp.getBefore().getArgument(operandIdx));
             auto &valueUsage = solver.getElementFor<ValueResourceUsage>(
                 *this, value, DFX::Resolution::REQUIRED);
             getState() ^= valueUsage.getState();
-          } else if (auto forOp =
-                         dyn_cast_or_null<scf::ForOp>(op->getParentOp())) {
+            auto &parentUsage = solver.getElementFor<ValueResourceUsage>(
+                *this, Position::forValue(whileOp->getResult(operandIdx)),
+                DFX::Resolution::REQUIRED);
+            getState() ^= parentUsage.getState();
+          } else if (auto forOp = dyn_cast<scf::ForOp>(op->getParentOp())) {
             auto value = Position::forValue(forOp.getRegionIterArg(operandIdx));
             auto &valueUsage = solver.getElementFor<ValueResourceUsage>(
                 *this, value, DFX::Resolution::REQUIRED);
             getState() ^= valueUsage.getState();
-
             auto &parentUsage = solver.getElementFor<ValueResourceUsage>(
                 *this, Position::forValue(forOp->getResult(operandIdx)),
                 DFX::Resolution::REQUIRED);
             getState() ^= parentUsage.getState();
           } else {
-            assert(false && "Unsupported test case");
+            assert(false && "unhandled scf yield parent");
           }
         })
-        .Case([&](mlir::func::ReturnOp op) {
+        .Case([&](IREE::Util::ReturnOp op) {
           auto &operandUsage = solver.getElementFor<ValueResourceUsage>(
               *this, Position::forValue(op.getOperand(operandIdx)),
               DFX::Resolution::REQUIRED);
@@ -734,11 +737,14 @@
           // TODO(benvanik): remove kFavorTransients.
           bool isSourceInternal = isAssumed(NOT_EXTERNAL);
           bool isTargetExternal = !resultUsage.isAssumed(NOT_EXTERNAL);
-          if (kFavorTransients && isSourceInternal && isTargetExternal) {
+          bool deviceChange =
+              op.getSourceAffinityAttr() != op.getResultAffinityAttr();
+          if ((kFavorTransients || deviceChange) && isSourceInternal &&
+              isTargetExternal) {
             LLVM_DEBUG({
               llvm::dbgs()
                   << "[ValueResourceUsage] skipping back prop of external into "
-                     "internal due to kFavorTransients: ";
+                     "internal due to kFavorTransients/device-change: ";
               op.print(llvm::dbgs(), solver.getAsmState());
               llvm::dbgs() << "\n";
             });
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamBase.td b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamBase.td
index bfcca44..4c8fb8d 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamBase.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamBase.td
@@ -504,6 +504,7 @@
 }
 
 def Stream_Resource : TypeDef<Stream_Dialect, "Resource", [
+  Stream_AffinityType,
   Util_ReferenceType,
   Util_SizeAwareType,
   DeclareTypeInterfaceMethods<Util_GlobalType, [
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamInterfaces.td b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamInterfaces.td
index f17b753..2b686b7 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamInterfaces.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamInterfaces.td
@@ -97,7 +97,13 @@
     // Returns an affinity active for the given operation.
     // This will recursively walk parent operations until one with the
     // `stream.affinity` attribute is found.
-    static AffinityAttr lookup(Operation *op);
+    static AffinityAttr lookup(Operation *fromOp);
+
+    // Returns an affinity active for the given operation or the fallback
+    // default if none is specified.
+    // This will recursively walk parent operations until one with the
+    // `stream.affinity` attribute is found.
+    static AffinityAttr lookupOrDefault(Operation *fromOp);
 
     // TODO(benvanik): replace with more fine-grained compatibility checks.
     // "Compatible" can mean a lot of things: are they cache-coherent, are they
@@ -116,10 +122,24 @@
 }
 
 //===----------------------------------------------------------------------===//
+// IREE::Stream::AffinityTypeInterface
+//===----------------------------------------------------------------------===//
+
+def Stream_AffinityType : TypeInterface<"AffinityTypeInterface"> {
+  let cppNamespace = "::mlir::iree_compiler::IREE::Stream";
+
+  let description = [{
+    Indicates a type represents a resource that has its affinity tracked.
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // IREE::Stream::AffinityOpInterface
 //===----------------------------------------------------------------------===//
 
 def Stream_AffinityOp : OpInterface<"AffinityOpInterface"> {
+  let cppNamespace = "::mlir::iree_compiler::IREE::Stream";
+
   let description = [{
     TBD. Used to denote a stream affinity for ops and specify the kind of
     environment the ops are expected run in.
@@ -142,6 +162,19 @@
     >,
     InterfaceMethod<
       /*desc=*/[{
+        Returns true if the operands and results should be pinned to the
+        affinity of the op. This overrides all automatic placement logic.
+      }],
+      /*retTy=*/"bool",
+      /*methodName=*/"pinsValueAffinity",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return false;
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
         Returns the stream affinity for the op, indicating where it should run.
       }],
       /*retTy=*/"IREE::Stream::AffinityAttr",
@@ -149,7 +182,7 @@
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return dyn_cast_or_null<IREE::Stream::AffinityAttr>($_self->getAttr("affinity"));
+        return dyn_cast_or_null<IREE::Stream::AffinityAttr>($_op->getAttr("affinity"));
       }]
     >,
     InterfaceMethod<
@@ -161,8 +194,8 @@
       /*args=*/(ins "IREE::Stream::AffinityAttr":$value),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        if (value) $_self->setAttr("affinity", value);
-        else $_self->removeAttr("affinity");
+        if (value) $_op->setAttr("affinity", value);
+        else $_op->removeAttr("affinity");
       }]
     >,
   ];
@@ -173,6 +206,8 @@
 //===----------------------------------------------------------------------===//
 
 def Stream_StreamableOp : OpInterface<"StreamableOpInterface"> {
+  let cppNamespace = "::mlir::iree_compiler::IREE::Stream";
+
   let description = [{
     Interface for ops that can be asynchronous executed in a streaming context.
   }];
@@ -212,6 +247,8 @@
 //===----------------------------------------------------------------------===//
 
 def Stream_AsyncAccessOp : OpInterface<"AsyncAccessOpInterface"> {
+  let cppNamespace = "::mlir::iree_compiler::IREE::Stream";
+
   let description = [{
     Interface for stream.async.* ops that access subviews of resources.
     This allows for some basic analysis and is only valid prior to allocation.
@@ -240,6 +277,8 @@
 //===----------------------------------------------------------------------===//
 
 def Stream_SubviewEffectOp : OpInterface<"SubviewEffectOpInterface"> {
+  let cppNamespace = "::mlir::iree_compiler::IREE::Stream";
+
   let description = [{
     Interface for ops that operate on subviews of resources used to query the
     memory effects for subviews on operands.
@@ -258,6 +297,8 @@
 //===----------------------------------------------------------------------===//
 
 def Stream_TimelineOp : OpInterface<"TimelineOpInterface"> {
+  let cppNamespace = "::mlir::iree_compiler::IREE::Stream";
+
   let description = [{
     Interface for ops that operate in an ordered sequence defined by timepoints.
   }];
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
index 84cf1eb..9360eea 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOpFolders.cpp
@@ -1231,7 +1231,7 @@
         constantOp, constantOp.getResult().getType(), splatOp.getResult(),
         resultSize, resultSize,
         /*source_affinity=*/constantOp.getAffinityAttr(),
-        /*result_affinity=*/nullptr);
+        /*result_affinity=*/constantOp.getAffinityAttr());
     return success();
   }
 };
@@ -1452,9 +1452,9 @@
   LogicalResult matchAndRewrite(AsyncConstantOp constantOp,
                                 PatternRewriter &rewriter) const override {
     auto value = dyn_cast<ElementsAttr>(constantOp.getValue());
-    if (!value || !value.isSplat())
+    if (!value || !value.isSplat()) {
       return failure();
-
+    }
     auto splatElementAttr =
         llvm::dyn_cast<SplatElementsAttr>(value).getSplatValue<TypedAttr>();
     auto splatValue = rewriter.create<arith::ConstantOp>(
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
index 698c7b9..358d1fd 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
@@ -2020,14 +2020,57 @@
 }
 
 IREE::Stream::AffinityAttr AsyncTransferOp::getAffinityAttr() {
-  return getResultAffinityAttr();
+  auto sourceType = cast<IREE::Stream::ResourceType>(getSource().getType());
+  auto resultType = cast<IREE::Stream::ResourceType>(getResult().getType());
+  if (sourceType.getLifetime() == IREE::Stream::Lifetime::Staging &&
+      resultType.getLifetime() == IREE::Stream::Lifetime::Staging) {
+    // TODO(multi-device): figure out how to model staging->staging transfers.
+    return getSourceAffinityAttr();
+  } else if (sourceType.getLifetime() == IREE::Stream::Lifetime::Staging) {
+    // If source is staging then the op should execute on the consumer.
+    return getResultAffinityAttr();
+  } else if (resultType.getLifetime() == IREE::Stream::Lifetime::Staging) {
+    // If result is staging then the op should execute on the producer.
+    return getSourceAffinityAttr();
+  } else {
+    // Default to result affinity.
+    return getResultAffinityAttr();
+  }
 }
 
 void AsyncTransferOp::setAffinityAttr(IREE::Stream::AffinityAttr value) {
-  if (value)
-    setResultAffinityAttr(value);
-  else
-    removeResultAffinityAttr();
+  auto sourceType = cast<IREE::Stream::ResourceType>(getSource().getType());
+  auto resultType = cast<IREE::Stream::ResourceType>(getResult().getType());
+  if (sourceType.getLifetime() == IREE::Stream::Lifetime::Staging &&
+      resultType.getLifetime() == IREE::Stream::Lifetime::Staging) {
+    // TODO(multi-device): figure out how to model staging->staging transfers.
+    if (value) {
+      setSourceAffinityAttr(value);
+    } else {
+      removeSourceAffinityAttr();
+    }
+  } else if (sourceType.getLifetime() == IREE::Stream::Lifetime::Staging) {
+    // If source is staging then the op should execute on the consumer.
+    if (value) {
+      setResultAffinityAttr(value);
+    } else {
+      removeResultAffinityAttr();
+    }
+  } else if (resultType.getLifetime() == IREE::Stream::Lifetime::Staging) {
+    // If result is staging then the op should execute on the producer.
+    if (value) {
+      setSourceAffinityAttr(value);
+    } else {
+      removeSourceAffinityAttr();
+    }
+  } else {
+    // Default to result affinity.
+    if (value) {
+      setResultAffinityAttr(value);
+    } else {
+      removeResultAffinityAttr();
+    }
+  }
 }
 
 void AsyncTransferOp::getAsyncAccessRanges(
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
index dbe5207..871e3bb 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
@@ -1520,7 +1520,7 @@
 
   let assemblyFormat = [{
     (`on` `(` $affinity^ `)`)?
-    $value `,` $target `[` $start_indices `for` $lengths `]` `:`
+    $value `,` $target (`[` $start_indices `for` $lengths^ `]`)? `:`
     type($value)
     `->`
     $target_encoding (`` `{` $target_encoding_dims^ `}`)?
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.cpp
index 19c2410..82b8609 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.cpp
@@ -335,20 +335,40 @@
 // #stream.affinity
 //===----------------------------------------------------------------------===//
 
-AffinityAttr AffinityAttr::lookup(Operation *op) {
-  auto attrId = StringAttr::get(op->getContext(), "stream.affinity");
-  while (op) {
-    if (auto affinityOp = llvm::dyn_cast<AffinityOpInterface>(op)) {
-      auto affinity = affinityOp.getAffinityAttr();
-      if (affinity)
+// static
+AffinityAttr AffinityAttr::lookup(Operation *fromOp) {
+  auto attrId = StringAttr::get(fromOp->getContext(), "stream.affinity");
+  while (fromOp) {
+    if (auto affinityOp = llvm::dyn_cast<AffinityOpInterface>(fromOp)) {
+      if (auto affinity = affinityOp.getAffinityAttr()) {
         return affinity;
+      }
     }
-    auto attr = op->getAttrOfType<AffinityAttr>(attrId);
-    if (attr)
+    if (auto attr = fromOp->getAttrOfType<AffinityAttr>(attrId)) {
       return attr;
-    op = op->getParentOp();
+    }
+    fromOp = fromOp->getParentOp();
   }
-  return {}; // No affinity found; let caller decide what to do.
+  // No affinity found; let caller decide what to do.
+  return {};
+}
+
+// static
+AffinityAttr AffinityAttr::lookupOrDefault(Operation *fromOp) {
+  if (auto affinityAttr = AffinityAttr::lookup(fromOp)) {
+    return affinityAttr; // found a specified affinity
+  }
+  auto attrId =
+      StringAttr::get(fromOp->getContext(), "stream.affinity.default");
+  while (fromOp) {
+    if (auto affinityAttr =
+            fromOp->getAttrOfType<IREE::Stream::AffinityAttr>(attrId)) {
+      return affinityAttr;
+    }
+    fromOp = fromOp->getParentOp();
+  }
+  // No affinity or default found; let caller decide what to do.
+  return {};
 }
 
 // static
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.h b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.h
index 42b8424..d69e226 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.h
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.h
@@ -69,9 +69,7 @@
 
 #include "iree/compiler/Dialect/Stream/IR/StreamAttrInterfaces.h.inc" // IWYU pragma: export
 
-namespace mlir::iree_compiler::IREE::Stream {
 #include "iree/compiler/Dialect/Stream/IR/StreamTypeInterfaces.h.inc" // IWYU pragma: export
-} // namespace mlir::iree_compiler::IREE::Stream
 
 // clang-format off: must be included after all LLVM/MLIR headers.
 #define GET_TYPEDEF_CLASSES
@@ -99,8 +97,12 @@
                          const AsyncAccessRange &rhs);
 };
 
+} // namespace mlir::iree_compiler::IREE::Stream
+
 #include "iree/compiler/Dialect/Stream/IR/StreamOpInterfaces.h.inc" // IWYU pragma: export
 
+namespace mlir::iree_compiler::IREE::Stream {
+
 //===----------------------------------------------------------------------===//
 // custom<ParameterReference>($scope, $key)
 //===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/AnnotateAffinities.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/AnnotateAffinities.cpp
new file mode 100644
index 0000000..62b9db2
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/AnnotateAffinities.cpp
@@ -0,0 +1,127 @@
+// Copyright 2024 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Dialect/Stream/Analysis/Affinity.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamDialect.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
+#include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
+#include "iree/compiler/Dialect/Stream/Transforms/Passes.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir::iree_compiler::IREE::Stream {
+
+#define GEN_PASS_DEF_ANNOTATEAFFINITIESPASS
+#include "iree/compiler/Dialect/Stream/Transforms/Passes.h.inc"
+
+namespace {
+
+//===----------------------------------------------------------------------===//
+// --iree-stream-annotate-affinities
+//===----------------------------------------------------------------------===//
+
+static void annotateOp(Operation *op,
+                       ArrayRef<IREE::Stream::AffinityAttr> affinities) {
+  auto affinityOp = dyn_cast<IREE::Stream::AffinityOpInterface>(op);
+  if (!affinityOp || !affinityOp.requiresAffinity()) {
+    return;
+  }
+  if (!affinities.empty()) {
+    op->setAttr("stream.affinities",
+                ArrayAttr::get(op->getContext(),
+                               llvm::to_vector_of<Attribute>(affinities)));
+  }
+}
+
+static void annotateGlobalOp(IREE::Util::GlobalOpInterface globalOp,
+                             AffinityAnalysis &affinityAnalysis) {
+  if (!isa<IREE::Stream::AffinityTypeInterface>(globalOp.getGlobalType())) {
+    return;
+  }
+  SmallVector<IREE::Stream::AffinityAttr> affinities;
+  if (affinityAnalysis.tryLookupGlobalAffinity(globalOp, affinities)) {
+    annotateOp(globalOp, affinities);
+  }
+}
+
+static void annotateOperandsAndResults(Operation *op,
+                                       AffinityAnalysis &affinityAnalysis) {
+  auto emptyArray = ArrayAttr::get(op->getContext(), {});
+  SmallVector<Attribute> operandAttrs;
+  for (auto operand : op->getOperands()) {
+    if (isa<IREE::Stream::AffinityTypeInterface>(operand.getType())) {
+      SmallVector<IREE::Stream::AffinityAttr> affinities;
+      if (affinityAnalysis.tryLookupResourceAffinity(operand, affinities)) {
+        operandAttrs.push_back(ArrayAttr::get(
+            op->getContext(), llvm::to_vector_of<Attribute>(affinities)));
+      } else {
+        operandAttrs.push_back(emptyArray);
+      }
+    }
+  }
+  SmallVector<Attribute> resultAttrs;
+  for (auto result : op->getResults()) {
+    if (isa<IREE::Stream::AffinityTypeInterface>(result.getType())) {
+      SmallVector<IREE::Stream::AffinityAttr> affinities;
+      if (affinityAnalysis.tryLookupResourceAffinity(result, affinities)) {
+        resultAttrs.push_back(ArrayAttr::get(
+            op->getContext(), llvm::to_vector_of<Attribute>(affinities)));
+      } else {
+        resultAttrs.push_back(emptyArray);
+      }
+    }
+  }
+  if (!operandAttrs.empty()) {
+    op->setAttr("stream.affinities.operands",
+                ArrayAttr::get(op->getContext(), operandAttrs));
+  }
+  if (!resultAttrs.empty()) {
+    op->setAttr("stream.affinities.results",
+                ArrayAttr::get(op->getContext(), resultAttrs));
+  }
+}
+
+static void annotateFuncOp(FunctionOpInterface funcOp,
+                           AffinityAnalysis &affinityAnalysis) {
+  funcOp.walk([&](Operation *op) {
+    SmallVector<IREE::Stream::AffinityAttr> affinities;
+    if (affinityAnalysis.tryLookupExecutionAffinity(op, affinities)) {
+      annotateOp(op, affinities);
+    }
+    annotateOperandsAndResults(op, affinityAnalysis);
+  });
+}
+
+struct AnnotateAffinitiesPass
+    : public IREE::Stream::impl::AnnotateAffinitiesPassBase<
+          AnnotateAffinitiesPass> {
+  void runOnOperation() override {
+    // Run affinity analysis on the whole module.
+    AffinityAnalysis affinityAnalysis(getOperation());
+    if (failed(affinityAnalysis.run())) {
+      return signalPassFailure();
+    }
+
+    // Annotate all ops with derived affinities.
+    for (auto &op : getOperation().getOps()) {
+      if (op.hasTrait<OpTrait::IREE::Util::ObjectLike>())
+        continue;
+      if (auto globalOp = dyn_cast<IREE::Util::GlobalOpInterface>(op)) {
+        annotateGlobalOp(globalOp, affinityAnalysis);
+      } else if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
+        annotateFuncOp(funcOp, affinityAnalysis);
+      }
+    }
+  }
+};
+
+} // namespace
+
+} // namespace mlir::iree_compiler::IREE::Stream
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel
index 1a1d2e2..d2f326c 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel
@@ -15,6 +15,7 @@
 iree_compiler_cc_library(
     name = "Transforms",
     srcs = [
+        "AnnotateAffinities.cpp",
         "AnnotateDispatchArguments.cpp",
         "ConvertToStream.cpp",
         "DumpStatistics.cpp",
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt
index 9d78c8e..5eb3d27 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt
@@ -16,6 +16,7 @@
   HDRS
     "Passes.h"
   SRCS
+    "AnnotateAffinities.cpp"
     "AnnotateDispatchArguments.cpp"
     "ConvertToStream.cpp"
     "DumpStatistics.cpp"
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp
index b99b792..31a5bb6 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp
@@ -16,6 +16,12 @@
 #include "mlir/Pass/PassRegistry.h"
 #include "mlir/Transforms/Passes.h"
 
+static llvm::cl::opt<bool> clAnnotateInputAffinities(
+    "iree-stream-annotate-input-affinities",
+    llvm::cl::desc("Annotates all tensor/resource affinities on the input to "
+                   "the pipeline for debugging."),
+    llvm::cl::init(false));
+
 namespace mlir::iree_compiler::IREE::Stream {
 
 using FunctionLikeNest =
@@ -68,6 +74,13 @@
   // Conversion
   //----------------------------------------------------------------------------
 
+  // Annotate all ops/resources with the analyzed affinities.
+  // This should have no behavioral changes during conversion but allows for
+  // debugging of analysis errors in end-user tooling.
+  if (clAnnotateInputAffinities) {
+    passManager.addPass(IREE::Stream::createAnnotateAffinitiesPass());
+  }
+
   // Converts from all input dialects into various levels of the stream dialect.
   // Tensor-like things go to stream.tensor.* ops while lower level buffer-like
   // things will go to stream.async.* ops.
@@ -81,6 +94,9 @@
   // Constant/variable optimization
   //----------------------------------------------------------------------------
 
+  // Run inlining after having baked out affinities.
+  passManager.addPass(mlir::createInlinerPass());
+
   // Cleanup globals that were created during conversion.
   addCleanupPatterns(passManager);
 
@@ -96,10 +112,15 @@
   // TODO(benvanik): annotate all dispatches with preferred executable affinity.
   // TODO(benvanik): DFA to specify all value affinities and pin dispatches.
 
+  // TODO(multi-device): it's really nice to be able to verify here but it
+  // prevents compiling to stream without devices specified or continuation at
+  // various phases. It'd be nice to find a way to enable this when the user
+  // expects it to work and otherwise not.
+  //
   // Verify that all ops that may require affinities have them assigned or
   // available (on a parent scope, etc). This allows subsequent passes to trust
   // that an affinity lookup will always return a valid affinity.
-  passManager.addPass(IREE::Stream::createVerifyAffinitiesPass());
+  // passManager.addPass(IREE::Stream::createVerifyAffinitiesPass());
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
index ca2ec3a..f5ee39f 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
@@ -457,6 +457,11 @@
 // Diagnostics
 //===----------------------------------------------------------------------===//
 
+def AnnotateAffinitiesPass :
+    Pass<"iree-stream-annotate-affinities", "mlir::ModuleOp"> {
+  let summary = "Annotates affinities on all ops for debugging.";
+}
+
 def DumpStatisticsPass :
     Pass<"iree-stream-dump-statistics", "mlir::ModuleOp"> {
   let summary = "Dumps stream dialect usage information to a file.";
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleAllocation.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleAllocation.cpp
index 1bec564..c8510a6 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleAllocation.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleAllocation.cpp
@@ -668,7 +668,8 @@
     return llvm::cast<IREE::Stream::ResourceType>(value.getType())
                .getLifetime() == IREE::Stream::Lifetime::Staging;
   };
-  auto currentAffinityAttr = IREE::Stream::AffinityAttr::lookup(asyncOp);
+  auto currentAffinityAttr =
+      IREE::Stream::AffinityAttr::lookupOrDefault(asyncOp);
   bool transferIn = asyncOp.getSourceAffinityAttr() != currentAffinityAttr ||
                     isStaging(asyncOp.getSource());
   bool transferOut = asyncOp.getResultAffinityAttr() != currentAffinityAttr ||
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/VerifyAffinities.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/VerifyAffinities.cpp
index 042bbb8..7579244 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/VerifyAffinities.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/VerifyAffinities.cpp
@@ -27,7 +27,7 @@
 verifyAffinityAssigned(IREE::Stream::AffinityOpInterface op) {
   if (!op.requiresAffinity()) {
     return success(); // does not require an affinity
-  } else if (IREE::Stream::AffinityAttr::lookup(op)) {
+  } else if (IREE::Stream::AffinityAttr::lookupOrDefault(op)) {
     return success(); // has an affinity
   }
   return op->emitOpError()
@@ -55,7 +55,10 @@
                   return WalkResult::interrupt();
                 }
               }
-              return WalkResult::advance();
+              return (op->hasTrait<OpTrait::IREE::Util::ObjectLike>() ||
+                      op->hasTrait<OpTrait::SymbolTable>())
+                         ? WalkResult::skip()
+                         : WalkResult::advance();
             })
             .wasInterrupted())
       return signalPassFailure();
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel
index 524a1ce..362d672 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel
@@ -16,6 +16,7 @@
     name = "lit",
     srcs = enforce_glob(
         [
+            "annotate_affinities.mlir",
             "annotate_dispatch_arguments.mlir",
             "convert_to_stream.mlir",
             "dump_statistics.mlir",
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
index 5ea9811..fe83ee6 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
@@ -14,6 +14,7 @@
   NAME
     lit
   SRCS
+    "annotate_affinities.mlir"
     "annotate_dispatch_arguments.mlir"
     "convert_to_stream.mlir"
     "dump_statistics.mlir"
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/annotate_affinities.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/annotate_affinities.mlir
new file mode 100644
index 0000000..c3e1f1e
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/annotate_affinities.mlir
@@ -0,0 +1,1549 @@
+// RUN: iree-opt --split-input-file --iree-stream-annotate-affinities %s | FileCheck %s
+
+// Tests that we can track affinity through optimization barriers. They're meant
+// to block optimization but we really can't do much if we don't track affinity.
+// We could change this in the future but tests would be harder to write and
+// there's not a lot that can be done with an unassigned resource.
+
+// CHECK-LABEL: @optimization_barrier_consumer
+util.func private @optimization_barrier_consumer() -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: util.optimization_barrier
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_dno = util.optimization_barrier %cst : tensor<1xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.transfer %cst_dno : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %cst_a : tensor<1xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @optimization_barrier_producer
+util.func private @optimization_barrier_producer() -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: util.optimization_barrier
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a_dno = util.optimization_barrier %cst_a : tensor<1xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %cst_a_dno : tensor<1xi32>
+}
+
+// -----
+
+// Tests that constant-like ops get placed with their consumer(s).
+// We want to replicate constants where they are consumed instead of performing
+// transfers at runtime to move them around and by placing with consumers we
+// can know when we need to do that early on.
+
+// CHECK-LABEL: @constant_op
+util.func private @constant_op() -> (tensor<1xi32>, tensor<1xi32>) {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  %cst = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.transfer %cst : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %cst_b = flow.tensor.transfer %cst : tensor<1xi32> to #hal.device.promise<@dev_b>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_b>]]
+  util.return %cst_a, %cst_b : tensor<1xi32>, tensor<1xi32>
+}
+
+// -----
+
+// Tests that splats (not constant-like but no consumed values) are placed with
+// their consumer(s). These are always best to rematerialize where they are
+// consumed to avoid allocating/transfering a bunch of repeated values.
+
+// CHECK-LABEL: @splat_op
+util.func private @splat_op() -> tensor<1xi32> {
+  %splat_value = arith.constant 123 : i32
+  // CHECK: flow.tensor.splat
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %splat = flow.tensor.splat %splat_value : tensor<1xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %splat_a = flow.tensor.transfer %splat : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %splat_a : tensor<1xi32>
+}
+
+// -----
+
+// Tests that imported tensor placement is inherited.
+// Frontends can use this to declare where they expect their arguments to
+// be living at the time the functions are invoked. Imports do not perform
+// transfers so we must use whatever is declared.
+
+// CHECK-LABEL: @imported_tensor
+util.func public @imported_tensor(%buffer_view: !hal.buffer_view, %fence: !hal.fence) -> tensor<1xi32> {
+  // CHECK: hal.tensor.import
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %tensor = hal.tensor.import on(#hal.device.promise<@dev_a>) wait(%fence) => %buffer_view "input" : !hal.buffer_view -> tensor<1xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %tensor : tensor<1xi32>
+}
+
+// -----
+
+// Tests that consumer-placed ops exported to buffers are properly placed.
+// Frontends can use this to explicitly define where exported tensors must live.
+// With consumer-placed ops like constants or splats we place them directly on
+// the export target.
+
+// CHECK-LABEL: @exported_constant
+util.func public @exported_constant(%fence: !hal.fence) -> !hal.buffer_view {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: hal.tensor.barrier
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_ready = hal.tensor.barrier join(%cst : tensor<1xi32>) => %fence : !hal.fence
+  // CHECK: hal.tensor.export
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  %buffer_view = hal.tensor.export on(#hal.device.promise<@dev_a>) %cst_ready "output" : tensor<1xi32> -> !hal.buffer_view
+  util.return %buffer_view : !hal.buffer_view
+}
+
+// -----
+
+// Tests that producer-placed ops exported to buffers get the appropriate
+// affinity on both devices. Frontends can use this to explicitly define where
+// exported tensors must live. Transfers may need to be inserted in order to
+// respect the required affinities. Note here that the operand to the export
+// is on @dev_a instead of the requested @dev_b.
+
+// CHECK-LABEL: @exported_producer
+util.func public @exported_producer(%fence: !hal.fence) -> !hal.buffer_view {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.transfer %cst : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: flow.tensor.clone
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %clone_a = flow.tensor.clone %cst_a : tensor<1xi32>
+  // CHECK: hal.tensor.barrier
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %clone_ready_a = hal.tensor.barrier join(%clone_a : tensor<1xi32>) => %fence : !hal.fence
+  // CHECK: hal.tensor.export
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  %buffer_view = hal.tensor.export on(#hal.device.promise<@dev_b>) %clone_ready_a "output" : tensor<1xi32> -> !hal.buffer_view
+  // CHECK: util.return
+  util.return %buffer_view : !hal.buffer_view
+}
+
+// -----
+
+// Test in-place aliased storage for results.
+// Frontends require that the storage be placed as indicated even if that means
+// introducing transfers such that the operation is not in-place.
+
+// CHECK-LABEL: @aliased_storage
+util.func public @aliased_storage(%view: !hal.buffer_view, %storage: !hal.buffer, %fence: !hal.fence) {
+  // CHECK: hal.tensor.import
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %arg_a = hal.tensor.import on(#hal.device.promise<@dev_a>) %view : !hal.buffer_view -> tensor<4xi32>
+  // CHECK: flow.dispatch @dispatch
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %ret_b = flow.dispatch @dispatch(%arg_a) : (tensor<4xi32>) -> tensor<4xi32>
+  // CHECK: hal.tensor.alias
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %alias_b = hal.tensor.alias on(#hal.device.promise<@dev_b>) %ret_b : tensor<4xi32> to %storage : !hal.buffer
+  // CHECK: hal.tensor.barrier
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  hal.tensor.barrier join(%alias_b : tensor<4xi32>) => %fence : !hal.fence
+  util.return
+}
+
+// -----
+
+// Tests aliased storage through tied dispatches.
+
+// CHECK-LABEL: @tied_aliased_storage
+util.func public @tied_aliased_storage(%view: !hal.buffer_view, %storage: !hal.buffer, %fence: !hal.fence) {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst = flow.tensor.constant dense<123> : tensor<4xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.transfer %cst : tensor<4xi32> to #hal.device.promise<@dev_a>
+  // CHECK: flow.dispatch @dispatch0
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %t0 = flow.dispatch @dispatch0(%cst) : (tensor<4xi32>) -> tensor<4xi32>
+  // CHECK: flow.dispatch @dispatch1
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %t1 = flow.dispatch @dispatch1(%t0) : (tensor<4xi32>) -> %t0
+  // CHECK: hal.tensor.alias
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %alias = hal.tensor.alias on(#hal.device.promise<@dev_b>) %t1 : tensor<4xi32> to %storage : !hal.buffer
+  // CHECK: hal.tensor.barrier
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  hal.tensor.barrier join(%alias : tensor<4xi32>) => %fence : !hal.fence
+  util.return
+}
+
+// -----
+
+// Tests that consumer-placed ops that pass through tied ops get attributed to
+// a single consumer.
+
+// CHECK-LABEL: @tied_constant
+util.func private @tied_constant() -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: flow.dispatch @a
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %tied = flow.dispatch @a(%cst) : (tensor<1xi32>) -> %cst
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %tied_a = flow.tensor.transfer %tied : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %tied_a : tensor<1xi32>
+}
+
+// -----
+
+// Tests that consumer-placed ops that pass through tied ops get attributed to
+// transitive consumers. This is not ideal but allows the application of
+// replication policies.
+
+// CHECK-LABEL: @tied_constant_multi_consumer
+util.func private @tied_constant_multi_consumer() -> (tensor<1xi32>, tensor<1xi32>) {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  %cst = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: flow.dispatch @a
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  %tied_0 = flow.dispatch @a(%cst) : (tensor<1xi32>) -> %cst
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %tied_0_a = flow.tensor.transfer %tied_0 : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: flow.dispatch @b
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  %tied_1 = flow.dispatch @b(%cst) : (tensor<1xi32>) -> %cst
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %tied_1_b = flow.tensor.transfer %tied_1 : tensor<1xi32> to #hal.device.promise<@dev_b>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_b>]]
+  util.return %tied_0_a, %tied_1_b : tensor<1xi32>, tensor<1xi32>
+}
+
+// -----
+
+// Tests the proper transfer of consumer-placed values prior to multiple tied
+// uses don't pollute the execution affinity of ops after transfers. Note that
+// the constant will still have multiple affinities to allow for policies that
+// replicate the constant.
+
+// CHECK-LABEL: @tied_transfer_constant_multi_consumer
+util.func private @tied_transfer_constant_multi_consumer() -> (tensor<1xi32>, tensor<1xi32>) {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  %cst = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.transfer %cst : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: flow.dispatch @a
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %tied_0 = flow.dispatch @a(%cst_a) : (tensor<1xi32>) -> %cst_a
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %tied_0_a = flow.tensor.transfer %tied_0 : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %cst_b = flow.tensor.transfer %cst : tensor<1xi32> to #hal.device.promise<@dev_b>
+  // CHECK: flow.dispatch @b
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %tied_1 = flow.dispatch @b(%cst_b) : (tensor<1xi32>) -> %cst_b
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %tied_1_b = flow.tensor.transfer %tied_1 : tensor<1xi32> to #hal.device.promise<@dev_b>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_b>]]
+  util.return %tied_0_a, %tied_1_b : tensor<1xi32>, tensor<1xi32>
+}
+
+// -----
+
+// Tests that implicitly placed consumers use their transfer execution affinity.
+
+// CHECK-LABEL: @transfer_execution_affinity
+util.func private @transfer_execution_affinity() -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %cst_b = flow.tensor.transfer %cst_a : tensor<1xi32> to #hal.device.promise<@dev_b>
+  // CHECK: flow.dispatch @dispatch
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %dispatch_b = flow.dispatch @dispatch(%cst_b) : (tensor<1xi32>) -> tensor<1xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  util.return %dispatch_b : tensor<1xi32>
+}
+
+// -----
+
+// Tests that explicitly placed consumers use their explicit execution affinity.
+
+// CHECK-LABEL: @explicit_execution_affinity
+util.func private @explicit_execution_affinity() -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: flow.dispatch @dispatch
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %dispatch_b = flow.dispatch @dispatch(%cst_a) {stream.affinity = #hal.device.promise<@dev_b>} : (tensor<1xi32>) -> tensor<1xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  util.return %dispatch_b : tensor<1xi32>
+}
+
+// -----
+
+// Tests that consumers of operands with multiple affinities inherit those
+// affinities for execution. This allows policies to determine where they want
+// to execute out of the resources they may be consuming.
+
+// CHECK-LABEL: @consume_multi_affinities
+util.func private @consume_multi_affinities() -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %cst_b = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_b>} dense<456> : tensor<1xi32>
+  // CHECK: flow.dispatch @dispatch
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  %dispatch_ab = flow.dispatch @dispatch(%cst_a, %cst_b) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  util.return %dispatch_ab : tensor<1xi32>
+}
+
+// -----
+
+// Tests that globals are placed where they are loaded.
+
+// CHECK: util.global private @consumed_global_a
+// CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+util.global private @consumed_global_a : tensor<1xi32>
+util.func private @consumer_fn() -> tensor<1xi32> {
+  // CHECK: util.global.load @consumed_global_a
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %load = util.global.load @consumed_global_a : tensor<1xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %load_a = flow.tensor.transfer %load : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %load_a : tensor<1xi32>
+}
+
+// -----
+
+// Tests that a global loaded from two locations is attributed to both
+// affinities. This allows policies to decide whether to replicate the global.
+
+// CHECK: util.global private @consumed_global_ab
+// CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]
+util.global private @consumed_global_ab : tensor<1xi32>
+util.func private @consumer_fn_a() -> tensor<1xi32> {
+  // CHECK: util.global.load @consumed_global_ab
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  %load = util.global.load @consumed_global_ab : tensor<1xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %load_a = flow.tensor.transfer %load : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %load_a : tensor<1xi32>
+}
+util.func private @consumer_fn_b() -> tensor<1xi32> {
+  // CHECK: util.global.load @consumed_global_ab
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  %load = util.global.load @consumed_global_ab : tensor<1xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %load_b = flow.tensor.transfer %load : tensor<1xi32> to #hal.device.promise<@dev_b>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  util.return %load_b : tensor<1xi32>
+}
+
+// -----
+
+// Tests that consumer-placed ops track through global loads.
+
+// CHECK: util.global private mutable @global_b
+// CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+util.global private mutable @global_b : tensor<1xi32>
+util.func private @producer_fn() {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: util.global.store
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.global.store %cst_a, @global_b : tensor<1xi32>
+  util.return
+}
+util.func private @consumer_fn() -> tensor<1xi32> {
+  // CHECK: util.global.load
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %load = util.global.load @global_b : tensor<1xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %load_b = flow.tensor.transfer %load : tensor<1xi32> to #hal.device.promise<@dev_b>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  util.return %load_b : tensor<1xi32>
+}
+
+// -----
+
+// Tests that globals that are only stored take the fallback placement of
+// their producer. This is silly but can arise prior to global optimization
+// passes that may elide them.
+
+// CHECK: util.global private mutable @global_a
+// CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+util.global private mutable @global_a : tensor<1xi32>
+util.func private @producer_fn() {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: util.global.store
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.global.store %cst_a, @global_a : tensor<1xi32>
+  util.return
+}
+
+// -----
+
+// Tests that global consumers that take on consumed affinity track the global.
+
+// CHECK: util.global private @global_a
+// CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+util.global private @global_a {stream.affinity = #hal.device.promise<@dev_a>} : tensor<1xi32>
+// CHECK: util.global private @global_b
+// CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+util.global private @global_b {stream.affinity = #hal.device.promise<@dev_b>} : tensor<1xi32>
+util.func private @consumer_fn() -> tensor<1xi32> {
+  // CHECK: util.global.load @global_a
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %load_a = util.global.load @global_a : tensor<1xi32>
+  // CHECK: util.global.load @global_b
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %load_b = util.global.load @global_b : tensor<1xi32>
+  // CHECK: flow.dispatch @dispatch
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  %result_ab = flow.dispatch @dispatch(%load_a, %load_b) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  util.return %result_ab : tensor<1xi32>
+}
+
+// -----
+
+// Tests a global update tick that operates on the global from multiple
+// affinities.
+
+// CHECK: util.global private mutable @global_a
+// CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+util.global private mutable @global_a {stream.affinity = #hal.device.promise<@dev_a>} = dense<123> : tensor<1xi32>
+util.func private @step(%arg0: tensor<2xi32>) -> tensor<2xi32> {
+  // CHECK: util.global.load @global_a
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %load_a = util.global.load @global_a : tensor<1xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %arg0_b = flow.tensor.transfer %arg0 : tensor<2xi32> to #hal.device.promise<@dev_b>
+  // CHECK: flow.dispatch @dispatch
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>], [#hal.device.promise<@dev_b>]]
+  %result_b:2 = flow.dispatch @dispatch(%load_a, %arg0_b) {stream.affinity = #hal.device.promise<@dev_b>} : (tensor<1xi32>, tensor<2xi32>) -> (tensor<1xi32>, tensor<2xi32>)
+  // CHECK: util.global.store
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  util.global.store %result_b#0, @global_a : tensor<1xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  util.return %result_b#1 : tensor<2xi32>
+}
+
+// -----
+
+// Tests that constants passed through selects are placed on the consumer.
+
+// CHECK-LABEL: @select_constants_consumed
+util.func private @select_constants_consumed(%cond: i1) -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_123 = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_456 = flow.tensor.constant dense<456> : tensor<1xi32>
+  // CHECK: arith.select
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst = arith.select %cond, %cst_123, %cst_456 : tensor<1xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.transfer %cst : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %cst_a : tensor<1xi32>
+}
+
+// -----
+
+// Tests that placed operands passed through selects are tracked on consumers.
+
+// CHECK-LABEL: @select_constants_placed
+util.func private @select_constants_placed(%cond: i1) -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %cst_b = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_b>} dense<456> : tensor<1xi32>
+  // CHECK: arith.select
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  %cst_ab = arith.select %cond, %cst_a, %cst_b : tensor<1xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  util.return %cst_ab : tensor<1xi32>
+}
+
+// -----
+
+// Tests that a callee that does not touch an argument still tracks the
+// affinity through it.
+
+// CHECK-LABEL: @passthrough_caller
+util.func private @passthrough_caller() -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: util.call @passthrough_callee
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %result_a = util.call @passthrough_callee(%cst_a) : (tensor<1xi32>) -> tensor<1xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %result_a : tensor<1xi32>
+}
+// CHECK: util.func private @passthrough_callee
+util.func private @passthrough_callee(%arg0: tensor<1xi32>) -> tensor<1xi32> {
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %arg0 : tensor<1xi32>
+}
+
+// -----
+
+// Tests that callees that consumer-placed arguments that are passed to callees
+// get placed based on callee usage.
+
+// CHECK-LABEL: @consumer_placement_caller
+util.func private @consumer_placement_caller() -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: util.call @consumer_placement_callee
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %result_a = util.call @consumer_placement_callee(%cst) : (tensor<1xi32>) -> tensor<1xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %result_a : tensor<1xi32>
+}
+// CHECK: util.func private @consumer_placement_callee
+util.func private @consumer_placement_callee(%arg0: tensor<1xi32>) -> tensor<1xi32> {
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %arg0_a = flow.tensor.transfer %arg0 : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %arg0_a : tensor<1xi32>
+}
+
+// -----
+
+// Tests that multiple potential affinities are propagated across call edges.
+
+// CHECK-LABEL: @select_caller
+util.func private @select_caller(%arg0: tensor<1xi32>, %cond: i1) -> tensor<1xi32> {
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %arg0_a = flow.tensor.transfer %arg0 : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: util.call @select_callee
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  %result_ab = util.call @select_callee(%arg0_a, %cond) : (tensor<1xi32>, i1) -> tensor<1xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  util.return %result_ab : tensor<1xi32>
+}
+// CHECK: util.func private @select_callee
+util.func private @select_callee(%arg0_a: tensor<1xi32>, %cond: i1) -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %cst_b = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_b>} dense<123> : tensor<1xi32>
+  // CHECK: arith.select
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  %select_ab = arith.select %cond, %arg0_a, %cst_b : tensor<1xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  util.return %select_ab : tensor<1xi32>
+}
+
+// -----
+
+// Tests that consumer-placed ops are propagated across call edges.
+
+// CHECK-LABEL: @consumer_multi_placement_caller
+util.func private @consumer_multi_placement_caller() -> (tensor<1xi32>, tensor<1xi32>) {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_c>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_c>]]
+  %cst = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: util.call @consumer_multi_placement_callee
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_c>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_c>]]
+  %result_0_c = util.call @consumer_multi_placement_callee(%cst) : (tensor<1xi32>) -> tensor<1xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_c>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %result_0_a = flow.tensor.transfer %result_0_c : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: util.call @consumer_multi_placement_callee
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_c>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_c>]]
+  %result_1_c = util.call @consumer_multi_placement_callee(%cst) : (tensor<1xi32>) -> tensor<1xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_c>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %result_1_b = flow.tensor.transfer %result_1_c : tensor<1xi32> to #hal.device.promise<@dev_b>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_b>]]
+  util.return %result_0_a, %result_1_b : tensor<1xi32>, tensor<1xi32>
+}
+// CHECK: util.func private @consumer_multi_placement_callee
+util.func private @consumer_multi_placement_callee(%arg0: tensor<1xi32>) -> tensor<1xi32> {
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_c>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_c>]]
+  %arg0_c = flow.tensor.transfer %arg0 : tensor<1xi32> to #hal.device.promise<@dev_c>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_c>]]
+  util.return %arg0_c : tensor<1xi32>
+}
+
+// -----
+
+// Tests that operand/result affinities are tracked across call edges.
+
+// CHECK-LABEL: @dispatch_fn_a
+util.func private @dispatch_fn_a() -> tensor<4xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %0 = flow.tensor.constant dense<123> : tensor<4xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %1 = flow.tensor.transfer %0 : tensor<4xi32> to #hal.device.promise<@dev_a>
+  // CHECK: flow.dispatch @dispatch_a_0
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %2 = flow.dispatch @dispatch_a_0(%1) : (tensor<4xi32>) -> tensor<4xi32>
+  // CHECK: util.call @dispatch_fn_b
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %3 = util.call @dispatch_fn_b(%2) : (tensor<4xi32>) -> tensor<4xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %4 = flow.tensor.transfer %3 : tensor<4xi32> to #hal.device.promise<@dev_a>
+  // CHECK: flow.dispatch @dispatch_a_1
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %5 = flow.dispatch @dispatch_a_1(%4) : (tensor<4xi32>) -> tensor<4xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %5 : tensor<4xi32>
+}
+// CHECK: util.func private @dispatch_fn_b
+util.func private @dispatch_fn_b(%arg0: tensor<4xi32>) -> tensor<4xi32> {
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %0 = flow.tensor.transfer %arg0 : tensor<4xi32> to #hal.device.promise<@dev_b>
+  // CHECK: flow.dispatch @dispatch_b
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %1 = flow.dispatch @dispatch_b(%0) : (tensor<4xi32>) -> tensor<4xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  util.return %1 : tensor<4xi32>
+}
+
+// -----
+
+// Tests a realistic call graph with explicit transfers.
+
+// CHECK-LABEL: @dispatch_fn_a
+util.func private @dispatch_fn_a() -> tensor<4xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %0 = flow.tensor.constant dense<123> : tensor<4xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %1 = flow.tensor.transfer %0 : tensor<4xi32> to #hal.device.promise<@dev_a>
+  // CHECK: util.call @dispatch_fn_b
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %2 = util.call @dispatch_fn_b(%1) : (tensor<4xi32>) -> tensor<4xi32>
+  // CHECK: util.call @dispatch_fn_c
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_c>]]
+  %3 = util.call @dispatch_fn_c(%1) : (tensor<4xi32>) -> tensor<4xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %4 = flow.tensor.transfer %2 : tensor<4xi32> to #hal.device.promise<@dev_a>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_c>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %5 = flow.tensor.transfer %3 : tensor<4xi32> to #hal.device.promise<@dev_a>
+  // CHECK: flow.dispatch @dispatch_a
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %6 = flow.dispatch @dispatch_a(%4, %5) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %5 : tensor<4xi32>
+}
+// CHECK: util.func private @dispatch_fn_b
+util.func private @dispatch_fn_b(%arg0: tensor<4xi32>) -> tensor<4xi32> {
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %0 = flow.tensor.transfer %arg0 : tensor<4xi32> to #hal.device.promise<@dev_b>
+  // CHECK: flow.dispatch @dispatch_b
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %1 = flow.dispatch @dispatch_b(%0) : (tensor<4xi32>) -> tensor<4xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  util.return %1 : tensor<4xi32>
+}
+// CHECK: util.func private @dispatch_fn_c
+util.func private @dispatch_fn_c(%arg0: tensor<4xi32>) -> tensor<4xi32> {
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_c>]]
+  %0 = flow.tensor.transfer %arg0 : tensor<4xi32> to #hal.device.promise<@dev_c>
+  // CHECK: flow.dispatch @dispatch_c
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_c>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_c>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_c>]]
+  %1 = flow.dispatch @dispatch_c(%0) : (tensor<4xi32>) -> tensor<4xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_c>]]
+  util.return %1 : tensor<4xi32>
+}
+
+// -----
+
+// Tests that consumer-placed ops are tracked across branch edges.
+
+// CHECK-LABEL: @cfg_branch_constant_consumed
+util.func private @cfg_branch_constant_consumed() -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: cf.br ^bb1
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  cf.br ^bb1(%cst : tensor<1xi32>)
+^bb1(%bb1_arg0: tensor<1xi32>):
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.transfer %bb1_arg0 : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %cst_a : tensor<1xi32>
+}
+
+// -----
+
+// Tests that producer-placed ops are tracked across branch edges.
+
+// CHECK-LABEL: @cfg_branch_dispatch_produced
+util.func private @cfg_branch_dispatch_produced() -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: cf.br ^bb1
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  cf.br ^bb1(%cst_a : tensor<1xi32>)
+^bb1(%bb1_arg0: tensor<1xi32>):
+  // CHECK: flow.dispatch @dispatch
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %dispatch_a = flow.dispatch @dispatch(%bb1_arg0) : (tensor<1xi32>) -> tensor<1xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %dispatch_a : tensor<1xi32>
+}
+
+// -----
+
+// Tests that back edges on loops track affinity changes.
+
+// CHECK-LABEL: @cfg_loop_back_edge
+util.func private @cfg_loop_back_edge() -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: cf.br ^bb1
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  cf.br ^bb1(%cst_a : tensor<1xi32>)
+^bb1(%bb1_arg0: tensor<1xi32>):
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %bb1_arg0_b = flow.tensor.transfer %bb1_arg0 : tensor<1xi32> to #hal.device.promise<@dev_b>
+  // CHECK: util.call @step
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  %cond = util.call @step(%bb1_arg0_b) : (tensor<1xi32>) -> i1
+  // CHECK: cf.cond_br
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_b>]]
+  cf.cond_br %cond, ^bb1(%bb1_arg0 : tensor<1xi32>), ^bb2(%bb1_arg0_b : tensor<1xi32>)
+^bb2(%bb2_arg0: tensor<1xi32>):
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_c>]]
+  %bb2_arg0_c = flow.tensor.transfer %bb2_arg0 : tensor<1xi32> to #hal.device.promise<@dev_c>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_c>]]
+  util.return %bb2_arg0_c : tensor<1xi32>
+}
+util.func private @step(tensor<1xi32>) -> i1
+
+// -----
+
+// Tests that conditional branches acting as selects propagate both affinities.
+
+// CHECK-LABEL: @cfg_cond_branch_select
+util.func private @cfg_cond_branch_select(%cond: i1) -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %cst_b = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_b>} dense<456> : tensor<1xi32>
+  // CHECK: cf.cond_br
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_b>]]
+  cf.cond_br %cond, ^bb1(%cst_a : tensor<1xi32>), ^bb1(%cst_b : tensor<1xi32>)
+^bb1(%bb1_arg0: tensor<1xi32>):
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  util.return %bb1_arg0 : tensor<1xi32>
+}
+
+// -----
+
+// Tests that consumer-placed ops through conditional branches acting as selects
+// get placed on all targets.
+
+// CHECK-LABEL: @cfg_cond_branch_select_consumer
+util.func private @cfg_cond_branch_select_consumer(%cond: i1) -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  %cst = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: cf.cond_br
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>], [#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  cf.cond_br %cond, ^bb1(%cst : tensor<1xi32>), ^bb2(%cst : tensor<1xi32>)
+^bb1(%bb1_arg0: tensor<1xi32>):
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.transfer %bb1_arg0 : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %cst_a : tensor<1xi32>
+^bb2(%bb2_arg0: tensor<1xi32>):
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %cst_b = flow.tensor.transfer %bb2_arg0 : tensor<1xi32> to #hal.device.promise<@dev_b>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  util.return %cst_b : tensor<1xi32>
+}
+
+// -----
+
+// Tests scf.if capturing consumer-placed ops tracks the affinity into nested
+// regions.
+
+// CHECK-LABEL: @scf_if_capture_consumer
+util.func private @scf_if_capture_consumer(%cond: i1) -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  %cst = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: scf.if
+  %cst_ab = scf.if %cond -> tensor<1xi32> {
+    // CHECK: flow.tensor.transfer
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+    %cst_a = flow.tensor.transfer %cst : tensor<1xi32> to #hal.device.promise<@dev_a>
+    // CHECK: scf.yield
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    scf.yield %cst_a : tensor<1xi32>
+  // CHECK: else
+  } else {
+    // CHECK: flow.tensor.transfer
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+    %cst_b = flow.tensor.transfer %cst : tensor<1xi32> to #hal.device.promise<@dev_b>
+    // CHECK: scf.yield
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+    scf.yield %cst_b : tensor<1xi32>
+  // CHECK{LITERAL}: } {stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  }
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  util.return %cst_ab : tensor<1xi32>
+}
+
+// -----
+
+// Tests scf.if capturing explicitly placed ops tracks the affinity of their
+// produced results into consumers.
+
+// CHECK-LABEL: @scf_if_capture_producer
+util.func private @scf_if_capture_producer(%cond: i1) -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: scf.if
+  %cst_bc = scf.if %cond -> tensor<1xi32> {
+    // CHECK: flow.tensor.transfer
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+    %cst_b = flow.tensor.transfer %cst_a : tensor<1xi32> to #hal.device.promise<@dev_b>
+    // CHECK: scf.yield
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+    scf.yield %cst_b : tensor<1xi32>
+  // CHECK: else
+  } else {
+    // CHECK: flow.tensor.transfer
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_c>]]
+    %cst_c = flow.tensor.transfer %cst_a : tensor<1xi32> to #hal.device.promise<@dev_c>
+    // CHECK: scf.yield
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_c>]]
+    scf.yield %cst_c : tensor<1xi32>
+  // CHECK{LITERAL}: } {stream.affinities.results = [[#hal.device.promise<@dev_b>, #hal.device.promise<@dev_c>]]
+  }
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>, #hal.device.promise<@dev_c>]]
+  util.return %cst_bc : tensor<1xi32>
+}
+
+// -----
+
+// Tests scf.if returning unassigned consumer-placed operations has the affinity
+// tracked across scf.yields and assigned based on the consumer.
+
+// CHECK-LABEL: @scf_if_consumer_yield
+util.func private @scf_if_consumer_yield(%cond: i1) -> tensor<1xi32> {
+  // CHECK: scf.if
+  %cst = scf.if %cond -> tensor<1xi32> {
+    // CHECK: flow.tensor.constant
+    // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+    %cst_0 = flow.tensor.constant dense<123> : tensor<1xi32>
+    // CHECK: scf.yield
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    scf.yield %cst_0 : tensor<1xi32>
+  // CHECK: else
+  } else {
+    // CHECK: flow.tensor.constant
+    // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+    %cst_1 = flow.tensor.constant dense<456> : tensor<1xi32>
+    // CHECK: scf.yield
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    scf.yield %cst_1 : tensor<1xi32>
+  // CHECK{LITERAL}: } {stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  }
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.transfer %cst : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %cst_a : tensor<1xi32>
+}
+
+// -----
+
+// Tests that consumer-placed ops get placed based on their use in the body.
+
+// CHECK-LABEL: @scf_for_consumer_body_transfer
+util.func private @scf_for_consumer_body_transfer() -> tensor<1xi32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: scf.for
+  %for = scf.for %i = %c0 to %c2 step %c1 iter_args(%arg0 = %cst) -> tensor<1xi32> {
+    // CHECK: flow.tensor.transfer
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+    %arg0_a = flow.tensor.transfer %arg0 : tensor<1xi32> to #hal.device.promise<@dev_a>
+    // CHECK: flow.dispatch @dispatch
+    // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+    %t = flow.dispatch @dispatch(%arg0_a) : (tensor<1xi32>) -> tensor<1xi32>
+    // CHECK: scf.yield
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    scf.yield %t : tensor<1xi32>
+  // CHECK{LITERAL}: } {stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  }
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %for : tensor<1xi32>
+}
+
+// -----
+
+// Tests that scf.for ops with transfers/explicit affinities on the edges get
+// the
+
+// CHECK-LABEL: @scf_for_boundary_transfer
+util.func private @scf_for_boundary_transfer() -> (tensor<1xi32>, tensor<1xi32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  %cst = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: scf.for
+  %for:2 = scf.for %i = %c0 to %c2 step %c1 iter_args(%arg0 = %cst, %arg1 = %cst) -> (tensor<1xi32>, tensor<1xi32>) {
+    // CHECK: flow.tensor.transfer
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+    %arg0_a = flow.tensor.transfer %arg0 : tensor<1xi32> to #hal.device.promise<@dev_a>
+    // CHECK: flow.dispatch @dispatch
+    // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+    %t = flow.dispatch @dispatch(%arg0_a) : (tensor<1xi32>) -> tensor<1xi32>
+    // CHECK: scf.yield
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+    scf.yield %t, %arg1 : tensor<1xi32>, tensor<1xi32>
+  // CHECK{LITERAL}: } {stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>], [#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  }
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %for_0_b = flow.tensor.transfer %for#0 : tensor<1xi32> to #hal.device.promise<@dev_b>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %for_1_b = flow.tensor.transfer %for#1 : tensor<1xi32> to #hal.device.promise<@dev_b>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>], [#hal.device.promise<@dev_b>]]
+  util.return %for_0_b, %for_1_b : tensor<1xi32>, tensor<1xi32>
+}
+
+// -----
+
+// Tests that transfers track through iter_args.
+
+// CHECK-LABEL: @scf_for_body_transfer
+util.func private @scf_for_body_transfer() -> tensor<1xi32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: scf.for
+  %for = scf.for %i = %c0 to %c2 step %c1 iter_args(%arg0 = %cst_a) -> tensor<1xi32> {
+    // CHECK: flow.tensor.transfer
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+    %arg0_b = flow.tensor.transfer %arg0 : tensor<1xi32> to #hal.device.promise<@dev_b>
+    // CHECK: flow.dispatch @dispatch
+    // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+    %t = flow.dispatch @dispatch(%arg0_b) : (tensor<1xi32>) -> tensor<1xi32>
+    // CHECK: scf.yield
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+    scf.yield %t : tensor<1xi32>
+  // CHECK{LITERAL}: } {stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  }
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_c>]]
+  %for_c = flow.tensor.transfer %for : tensor<1xi32> to #hal.device.promise<@dev_c>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_c>]]
+  util.return %for_c : tensor<1xi32>
+}
+
+// -----
+
+// Tests that placed values track through iter_args to consumers in scf.for
+// bodies.
+
+// CHECK-LABEL: @scf_for_capture_producer
+util.func private @scf_for_capture_producer() -> tensor<1xi32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: scf.for
+  %for = scf.for %i = %c0 to %c2 step %c1 iter_args(%arg0 = %cst_a) -> tensor<1xi32> {
+    // CHECK: flow.dispatch @dispatch
+    // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+    %t = flow.dispatch @dispatch(%arg0) : (tensor<1xi32>) -> tensor<1xi32>
+    // CHECK: scf.yield
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    scf.yield %t : tensor<1xi32>
+  // CHECK{LITERAL}: } {stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  }
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %for : tensor<1xi32>
+}
+
+// -----
+
+// Tests that consumer-placed ops get placed based on their use in the body.
+
+// CHECK-LABEL: @scf_while_consumer_body_transfer
+util.func private @scf_while_consumer_body_transfer() -> tensor<1xi32> {
+  %c0 = arith.constant 0 : index
+  %c2_i32 = arith.constant 2 : i32
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: scf.while
+  %while = scf.while(%arg0 = %cst) : (tensor<1xi32>) -> tensor<1xi32> {
+    // CHECK: flow.tensor.load
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    %cond_i32 = flow.tensor.load %arg0[%c0] : tensor<1xi32>
+    %cond = arith.cmpi slt, %cond_i32, %c2_i32 : i32
+    // CHECK: scf.condition
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    scf.condition(%cond) %arg0 : tensor<1xi32>
+  } do {
+  ^bb0(%arg0: tensor<1xi32>):
+    // CHECK: flow.tensor.transfer
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+    %arg0_a = flow.tensor.transfer %arg0 : tensor<1xi32> to #hal.device.promise<@dev_a>
+    // CHECK: flow.dispatch @dispatch
+    // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+    %t = flow.dispatch @dispatch(%arg0_a) : (tensor<1xi32>) -> tensor<1xi32>
+    // CHECK: scf.yield
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    scf.yield %t : tensor<1xi32>
+  // CHECK: } attributes {
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  }
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %while : tensor<1xi32>
+}
+
+// -----
+
+// Tests that consumer-placed ops get placed based on their use as the result
+// of an scf.while body.
+
+// CHECK-LABEL: @scf_while_consumer_result_transfer
+util.func private @scf_while_consumer_result_transfer() -> tensor<1xi32> {
+  %c0 = arith.constant 0 : index
+  %c2_i32 = arith.constant 2 : i32
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst = flow.tensor.constant dense<123> : tensor<1xi32>
+  // CHECK: scf.while
+  %while = scf.while(%arg0 = %cst) : (tensor<1xi32>) -> tensor<1xi32> {
+    // CHECK: flow.tensor.load
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    %cond_i32 = flow.tensor.load %arg0[%c0] : tensor<1xi32>
+    %cond = arith.cmpi slt, %cond_i32, %c2_i32 : i32
+    // CHECK: scf.condition
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    scf.condition(%cond) %arg0 : tensor<1xi32>
+  } do {
+  ^bb0(%arg0: tensor<1xi32>):
+    // CHECK: flow.dispatch @dispatch
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+    %t = flow.dispatch @dispatch(%arg0) : (tensor<1xi32>) -> tensor<1xi32>
+    // CHECK: scf.yield
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    scf.yield %t : tensor<1xi32>
+  // CHECK: } attributes {
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  }
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %while_a = flow.tensor.transfer %while : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %while_a : tensor<1xi32>
+}
+
+// -----
+
+// Tests that transfers track through scf.while bodies.
+
+// CHECK-LABEL: @scf_while_body_transfer
+util.func private @scf_while_body_transfer() -> tensor<1xi32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2_i32 = arith.constant 2 : i32
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: scf.while
+  %while = scf.while(%arg0 = %cst_a) : (tensor<1xi32>) -> tensor<1xi32> {
+    // CHECK: flow.tensor.load
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+    %cond_i32 = flow.tensor.load %arg0[%c0] : tensor<1xi32>
+    %cond = arith.cmpi slt, %cond_i32, %c2_i32 : i32
+    // CHECK: scf.condition
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+    scf.condition(%cond) %arg0 : tensor<1xi32>
+  } do {
+  ^bb0(%arg0: tensor<1xi32>):
+    // CHECK: flow.tensor.transfer
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+    %arg0_b = flow.tensor.transfer %arg0 : tensor<1xi32> to #hal.device.promise<@dev_b>
+    // CHECK: flow.dispatch @dispatch
+    // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+    %t = flow.dispatch @dispatch(%arg0_b) : (tensor<1xi32>) -> tensor<1xi32>
+    // CHECK: scf.yield
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+    scf.yield %t : tensor<1xi32>
+  // CHECK: } attributes {
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  }
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>, #hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_c>]]
+  %while_c = flow.tensor.transfer %while : tensor<1xi32> to #hal.device.promise<@dev_c>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_c>]]
+  util.return %while_c : tensor<1xi32>
+}
+
+// -----
+
+// Tests that placed values track through to consumers in scf.while conditions.
+
+// CHECK-LABEL: @scf_while_capture_producer_condition
+util.func private @scf_while_capture_producer_condition() -> tensor<1xi32> {
+  %c0 = arith.constant 0 : index
+  %c2_i32 = arith.constant 2 : i32
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: scf.while
+  %while = scf.while(%arg0 = %cst_a) : (tensor<1xi32>) -> tensor<1xi32> {
+    // CHECK: flow.tensor.transfer
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+    %arg0_a = flow.tensor.transfer %arg0 : tensor<1xi32> to #hal.device.promise<@dev_a>
+    // CHECK: flow.tensor.load
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    %cond_i32 = flow.tensor.load %arg0_a[%c0] : tensor<1xi32>
+    %cond = arith.cmpi slt, %cond_i32, %c2_i32 : i32
+    // CHECK: scf.condition
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    scf.condition(%cond) %arg0 : tensor<1xi32>
+  } do {
+  ^bb0(%arg0: tensor<1xi32>):
+    // CHECK: flow.dispatch @dispatch
+    // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+    %t = flow.dispatch @dispatch(%arg0) : (tensor<1xi32>) -> tensor<1xi32>
+    // CHECK: scf.yield
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    scf.yield %t : tensor<1xi32>
+  // CHECK: } attributes {
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  }
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %while : tensor<1xi32>
+}
+
+// -----
+
+// Tests that placed values track through to consumers in scf.while bodies.
+
+// CHECK-LABEL: @scf_while_capture_producer_body
+util.func private @scf_while_capture_producer_body() -> tensor<1xi32> {
+  %c0 = arith.constant 0 : index
+  %c2_i32 = arith.constant 2 : i32
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst_a = flow.tensor.constant {stream.affinity = #hal.device.promise<@dev_a>} dense<123> : tensor<1xi32>
+  // CHECK: scf.while
+  %while = scf.while(%arg0 = %cst_a) : (tensor<1xi32>) -> tensor<1xi32> {
+    // CHECK: flow.tensor.load
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    %cond_i32 = flow.tensor.load %arg0[%c0] : tensor<1xi32>
+    %cond = arith.cmpi slt, %cond_i32, %c2_i32 : i32
+    // CHECK: scf.condition
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    scf.condition(%cond) %arg0 : tensor<1xi32>
+  } do {
+  ^bb0(%arg0: tensor<1xi32>):
+    // CHECK: flow.dispatch @dispatch
+    // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+    %t = flow.dispatch @dispatch(%arg0) : (tensor<1xi32>) -> tensor<1xi32>
+    // CHECK: scf.yield
+    // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+    scf.yield %t : tensor<1xi32>
+  // CHECK: } attributes {
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  }
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %while : tensor<1xi32>
+}
+
+// -----
+
+// Tests a realistic program with ABI ops.
+
+// CHECK-LABEL: @simple_program
+util.func public @simple_program(%arg0: !hal.buffer_view, %arg1: !hal.fence, %arg2: !hal.fence) -> !hal.buffer_view {
+  // CHECK: hal.tensor.import
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %0 = hal.tensor.import on(#hal.device.promise<@dev_a>) wait(%arg1) => %arg0 "input0" : !hal.buffer_view -> tensor<1xi32>
+  // CHECK: util.call @_simple_program
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %1 = util.call @_simple_program(%0) : (tensor<1xi32>) -> tensor<1xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %2 = flow.tensor.transfer %1 : tensor<1xi32> to #hal.device.promise<@dev_a>
+  // CHECK: hal.tensor.barrier
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %3 = hal.tensor.barrier join(%2 : tensor<1xi32>) => %arg2 : !hal.fence
+  // CHECK: hal.tensor.export
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  %4 = hal.tensor.export on(#hal.device.promise<@dev_a>) %3 "output0" : tensor<1xi32> -> !hal.buffer_view
+  util.return %4 : !hal.buffer_view
+}
+// CHECK: util.func private @_simple_program
+util.func private @_simple_program(%arg0: tensor<1xi32>) -> tensor<1xi32> {
+  // CHECK: util.call @dispatch_a
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %0 = util.call @dispatch_a(%arg0) : (tensor<1xi32>) -> tensor<1xi32>
+  // CHECK: flow.tensor.transfer
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %1 = flow.tensor.transfer %0 : tensor<1xi32> to #hal.device.promise<@dev_b>
+  // CHECK: util.call @dispatch_b
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %2 = util.call @dispatch_b(%1) : (tensor<1xi32>) -> tensor<1xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  util.return %2 : tensor<1xi32>
+}
+// CHECK: util.func private @dispatch_a
+util.func private @dispatch_a(%arg0: tensor<1xi32>) -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %cst = flow.tensor.constant dense<[1]> : tensor<1xi32>
+  // CHECK: flow.dispatch @dispatch_a
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_a>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>], [#hal.device.promise<@dev_a>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_a>]]
+  %0 = flow.dispatch @dispatch_a(%arg0, %cst) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_a>]]
+  util.return %0 : tensor<1xi32>
+}
+// CHECK: util.func private @dispatch_b
+util.func private @dispatch_b(%arg0: tensor<1xi32>) -> tensor<1xi32> {
+  // CHECK: flow.tensor.constant
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %cst = flow.tensor.constant dense<[2]> : tensor<1xi32>
+  // CHECK: flow.dispatch @dispatch_b
+  // CHECK-SAME{LITERAL}: stream.affinities = [#hal.device.promise<@dev_b>]
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>], [#hal.device.promise<@dev_b>]]
+  // CHECK-SAME{LITERAL}: stream.affinities.results = [[#hal.device.promise<@dev_b>]]
+  %0 = flow.dispatch @dispatch_b(%arg0, %cst) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
+  // CHECK: util.return
+  // CHECK-SAME{LITERAL}: stream.affinities.operands = [[#hal.device.promise<@dev_b>]]
+  util.return %0 : tensor<1xi32>
+}
diff --git a/compiler/src/iree/compiler/Dialect/Util/Analysis/Explorer.cpp b/compiler/src/iree/compiler/Dialect/Util/Analysis/Explorer.cpp
index e014588..bb46c83 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Analysis/Explorer.cpp
+++ b/compiler/src/iree/compiler/Dialect/Util/Analysis/Explorer.cpp
@@ -679,7 +679,8 @@
 // traversal algorithm separated from the policy here. This would let us
 // reuse the traversal for other kinds of walks that are more specific (like
 // only getting the ops or values instead of both, etc).
-TraversalResult Explorer::walkDefiningOps(Value value, ResultWalkFn fn) {
+TraversalResult Explorer::walkDefiningOps(Value value, ResultWalkFn fn,
+                                          TraversalBehavior options) {
   // Fast-path short-circuit for constants, which are like 25% of all IR.
   if (value.getDefiningOp() &&
       value.getDefiningOp()->hasTrait<OpTrait::ConstantLike>()) {
@@ -856,15 +857,17 @@
 
     // If the op is tied we may need to walk up to the operand the result is
     // tied to.
-    if (auto tiedOp = dyn_cast<IREE::Util::TiedOpInterface>(definingOp)) {
-      auto tiedOperand = tiedOp.getTiedResultOperand(resultValue);
-      if (tiedOperand) {
-        LLVM_DEBUG({
-          llvm::dbgs() << "   + queuing tied operand ";
-          tiedOperand.printAsOperand(llvm::dbgs(), asmState);
-          llvm::dbgs() << "\n";
-        });
-        worklist.insert(tiedOperand);
+    if (!bitEnumContains(options, TraversalBehavior::DONT_WALK_TIED_VALUES)) {
+      if (auto tiedOp = dyn_cast<IREE::Util::TiedOpInterface>(definingOp)) {
+        auto tiedOperand = tiedOp.getTiedResultOperand(resultValue);
+        if (tiedOperand) {
+          LLVM_DEBUG({
+            llvm::dbgs() << "   + queuing tied operand ";
+            tiedOperand.printAsOperand(llvm::dbgs(), asmState);
+            llvm::dbgs() << "\n";
+          });
+          worklist.insert(tiedOperand);
+        }
       }
     }
 
@@ -891,7 +894,8 @@
   return result;
 }
 
-TraversalResult Explorer::walkTransitiveUses(Value value, UseWalkFn fn) {
+TraversalResult Explorer::walkTransitiveUses(Value value, UseWalkFn fn,
+                                             TraversalBehavior options) {
   LLVM_DEBUG(llvm::dbgs() << "[[ Explorer::walkTransitiveUses ]]\n");
   TraversalResult result = TraversalResult::COMPLETE;
 
@@ -1090,15 +1094,17 @@
 
       // If the op is tied we may need to walk down to the results the operand
       // is tied to (multiple results can tie the same operand).
-      if (auto tiedOp = dyn_cast<IREE::Util::TiedOpInterface>(ownerOp)) {
-        for (auto tiedResult :
-             tiedOp.getOperandTiedResults(use.getOperandNumber())) {
-          LLVM_DEBUG({
-            llvm::dbgs() << "   + queuing tied result ";
-            tiedResult.printAsOperand(llvm::dbgs(), asmState);
-            llvm::dbgs() << "\n";
-          });
-          worklist.insert(tiedResult);
+      if (!bitEnumContains(options, TraversalBehavior::DONT_WALK_TIED_VALUES)) {
+        if (auto tiedOp = dyn_cast<IREE::Util::TiedOpInterface>(ownerOp)) {
+          for (auto tiedResult :
+               tiedOp.getOperandTiedResults(use.getOperandNumber())) {
+            LLVM_DEBUG({
+              llvm::dbgs() << "   + queuing tied result ";
+              tiedResult.printAsOperand(llvm::dbgs(), asmState);
+              llvm::dbgs() << "\n";
+            });
+            worklist.insert(tiedResult);
+          }
         }
       }
 
@@ -1149,14 +1155,18 @@
   return result;
 }
 
-TraversalResult Explorer::walkTransitiveUsers(Value value, OperationWalkFn fn) {
+TraversalResult Explorer::walkTransitiveUsers(Value value, OperationWalkFn fn,
+                                              TraversalBehavior options) {
   DenseSet<Operation *> visitedOwners;
-  return walkTransitiveUses(value, [&](OpOperand &use) {
-    if (visitedOwners.insert(use.getOwner()).second) {
-      return fn(use.getOwner());
-    }
-    return WalkResult::advance();
-  });
+  return walkTransitiveUses(
+      value,
+      [&](OpOperand &use) {
+        if (visitedOwners.insert(use.getOwner()).second) {
+          return fn(use.getOwner());
+        }
+        return WalkResult::advance();
+      },
+      options);
 }
 
 } // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Dialect/Util/Analysis/Explorer.h b/compiler/src/iree/compiler/Dialect/Util/Analysis/Explorer.h
index 1e975be..35ee12a 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Analysis/Explorer.h
+++ b/compiler/src/iree/compiler/Dialect/Util/Analysis/Explorer.h
@@ -37,6 +37,31 @@
   IGNORE,
 };
 
+enum class TraversalBehavior : uint32_t {
+  // When traversing defining ops any tied result will move through its tied
+  // operand. When traversing uses any tied operand will move through its tied
+  // results (as many as are tied to the operand).
+  DEFAULT = 0u,
+  // Don't traverse through tied operands or results.
+  DONT_WALK_TIED_VALUES = 1 << 0u,
+};
+inline TraversalBehavior operator~(TraversalBehavior value) {
+  return static_cast<TraversalBehavior>(~static_cast<uint32_t>(value));
+}
+inline TraversalBehavior operator|(TraversalBehavior lhs,
+                                   TraversalBehavior rhs) {
+  return static_cast<TraversalBehavior>(static_cast<uint32_t>(lhs) |
+                                        static_cast<uint32_t>(rhs));
+}
+inline TraversalBehavior operator&(TraversalBehavior lhs,
+                                   TraversalBehavior rhs) {
+  return static_cast<TraversalBehavior>(static_cast<uint32_t>(lhs) &
+                                        static_cast<uint32_t>(rhs));
+}
+inline bool bitEnumContains(TraversalBehavior bits, TraversalBehavior bit) {
+  return (static_cast<uint32_t>(bits) & static_cast<uint32_t>(bit)) != 0;
+}
+
 // Boolean operations on TraversalResult behave as though `INCOMPLETE` is
 // truthy to allow for |='ing results.
 enum class TraversalResult {
@@ -313,7 +338,9 @@
   //  Walk %2: [%2 of producer.b]
   //  Walk @some_user::%arg0: [%0 of producer.a]
   //  Walk @some_user::ret0: [%2 of producer.b]
-  TraversalResult walkDefiningOps(Value value, ResultWalkFn fn);
+  TraversalResult
+  walkDefiningOps(Value value, ResultWalkFn fn,
+                  TraversalBehavior options = TraversalBehavior::DEFAULT);
 
   // Randomly walks uses of |value| and any transitive alias of |value|.
   // The uses may come from any part of the program.
@@ -334,13 +361,17 @@
   //  Walk %arg0: [%arg0 of producer.a]
   //  Walk %0: [%0 of call @some_user, %arg0 of producer.b]
   //  Walk %2: [%2 of return, %1 of return]
-  TraversalResult walkTransitiveUses(Value value, UseWalkFn fn);
+  TraversalResult
+  walkTransitiveUses(Value value, UseWalkFn fn,
+                     TraversalBehavior options = TraversalBehavior::DEFAULT);
 
   // Randomly walks uses of |value| and any transitive alias of |value| and
   // returns each owner operation once. As a value may be used multiple times
   // by a single operation this is equivalent to a walkTransitiveUses with
   // deduplication on the owner of the use.
-  TraversalResult walkTransitiveUsers(Value value, OperationWalkFn fn);
+  TraversalResult
+  walkTransitiveUsers(Value value, OperationWalkFn fn,
+                      TraversalBehavior options = TraversalBehavior::DEFAULT);
 
 private:
   // Maps callee callable region -> call sites.
diff --git a/compiler/src/iree/compiler/ExternalInterfaces/StreamExternalModels.cpp b/compiler/src/iree/compiler/ExternalInterfaces/StreamExternalModels.cpp
index b82d599..ab1adf0 100644
--- a/compiler/src/iree/compiler/ExternalInterfaces/StreamExternalModels.cpp
+++ b/compiler/src/iree/compiler/ExternalInterfaces/StreamExternalModels.cpp
@@ -18,6 +18,35 @@
 
 namespace {
 
+template <typename OpT>
+struct OptionalOpAffinityAttrExternalModel
+    : public IREE::Stream::AffinityOpInterface::ExternalModel<
+          OptionalOpAffinityAttrExternalModel<OpT>, OpT> {
+  static void add(MLIRContext *context) {
+    OpT::template attachInterface<OptionalOpAffinityAttrExternalModel<OpT>>(
+        *context);
+  }
+
+  // Affinity only required for results that hold resources that
+  // require placement.
+  bool requiresAffinity(Operation *op) const {
+    auto resultType = cast<OpT>(op).getResult().getType();
+    return isa<TensorType>(resultType);
+  }
+
+  IREE::Stream::AffinityAttr getAffinityAttr(Operation *op) const {
+    return op->getAttrOfType<IREE::Stream::AffinityAttr>("stream.affinity");
+  }
+
+  void setAffinityAttr(Operation *op, IREE::Stream::AffinityAttr value) const {
+    if (value) {
+      op->setAttr("stream.affinity", value);
+    } else {
+      op->removeAttr("stream.affinity");
+    }
+  }
+};
+
 struct FlowTransferTargetAffinityAttrExternalModel
     : public IREE::Stream::AffinityOpInterface::ExternalModel<
           FlowTransferTargetAffinityAttrExternalModel,
@@ -29,11 +58,11 @@
 
   bool requiresAffinity(Operation *op) const { return true; }
 
-  IREE::Stream::AffinityAttr getAffinity(Operation *op) const {
+  IREE::Stream::AffinityAttr getAffinityAttr(Operation *op) const {
     return op->getAttrOfType<IREE::Stream::AffinityAttr>("target");
   }
 
-  void setAffinity(Operation *op, IREE::Stream::AffinityAttr value) const {
+  void setAffinityAttr(Operation *op, IREE::Stream::AffinityAttr value) const {
     op->setAttr("target", value);
   }
 };
@@ -49,12 +78,14 @@
 
   bool requiresAffinity(Operation *op) const { return false; }
 
-  IREE::Stream::AffinityAttr getAffinity(Operation *op) const {
+  bool pinsValueAffinity(Operation *op) const { return true; }
+
+  IREE::Stream::AffinityAttr getAffinityAttr(Operation *op) const {
     return op->getAttrOfType<IREE::Stream::AffinityAttr>("affinity");
   }
 
-  void setAffinity(Operation *op, IREE::Stream::AffinityAttr value) const {
-    if (value)
+  void setAffinityAttr(Operation *op, IREE::Stream::AffinityAttr value) const {
+    if (value) {
       op->setAttr("affinity", value);
     } else {
       op->removeAttr("affinity");
@@ -78,12 +109,12 @@
     return isa<TensorType>(globalType);
   }
 
-  IREE::Stream::AffinityAttr getAffinity(Operation *op) const {
+  IREE::Stream::AffinityAttr getAffinityAttr(Operation *op) const {
     return op->getAttrOfType<IREE::Stream::AffinityAttr>("stream.affinity");
   }
 
-  void setAffinity(Operation *op, IREE::Stream::AffinityAttr value) const {
-    if (value)
+  void setAffinityAttr(Operation *op, IREE::Stream::AffinityAttr value) const {
+    if (value) {
       op->setAttr("stream.affinity", value);
     } else {
       op->removeAttr("stream.affinity");
@@ -91,7 +122,7 @@
   }
 };
 
-template <typename OpT>
+template <typename OpT, bool kRequiresAffinity = true>
 struct AffinityOpAttrExternalModel
     : public IREE::Stream::AffinityOpInterface::ExternalModel<
           AffinityOpAttrExternalModel<OpT, kRequiresAffinity>, OpT> {
@@ -102,14 +133,14 @@
 
   // Most structural ops don't require affinities and after placement we don't
   // use the affinities even if the ops still exist.
-  bool requiresAffinity(Operation *op) const { return false; }
+  bool requiresAffinity(Operation *op) const { return kRequiresAffinity; }
 
-  IREE::Stream::AffinityAttr getAffinity(Operation *op) const {
+  IREE::Stream::AffinityAttr getAffinityAttr(Operation *op) const {
     return op->getAttrOfType<IREE::Stream::AffinityAttr>("stream.affinity");
   }
 
-  void setAffinity(Operation *op, IREE::Stream::AffinityAttr value) const {
-    if (value)
+  void setAffinityAttr(Operation *op, IREE::Stream::AffinityAttr value) const {
+    if (value) {
       op->setAttr("stream.affinity", value);
     } else {
       op->removeAttr("stream.affinity");
@@ -117,32 +148,71 @@
   }
 };
 
+struct TensorAffinityTypeExternalModel
+    : public IREE::Stream::AffinityTypeInterface::ExternalModel<
+          TensorAffinityTypeExternalModel, RankedTensorType> {
+  static void add(MLIRContext *context) {
+    RankedTensorType::attachInterface<TensorAffinityTypeExternalModel>(
+        *context);
+  }
+};
+
 } // namespace
 
 void registerStreamExternalModels(DialectRegistry &registry) {
-  registry.insert<IREE::Flow::FlowDialect>();
+  registry.addExtension(+[](MLIRContext *context) {
+    TensorAffinityTypeExternalModel::add(context);
+  });
+
+  registry.insert<arith::ArithDialect>();
   registry.addExtension(
-      +[](MLIRContext *context, IREE::Flow::FlowDialect *dialect) {
-        FlowTransferTargetAffinityAttrExternalModel::add(context);
+      +[](MLIRContext *context, arith::ArithDialect *dialect) {
+        OptionalOpAffinityAttrExternalModel<arith::ConstantOp>::add(context);
       });
 
+  registry.insert<IREE::Flow::FlowDialect>();
+  registry.addExtension(+[](MLIRContext *context,
+                            IREE::Flow::FlowDialect *dialect) {
+    FlowTransferTargetAffinityAttrExternalModel::add(context);
+    AffinityOpAttrExternalModel<IREE::Flow::DispatchRegionOp>::add(context);
+    AffinityOpAttrExternalModel<IREE::Flow::DispatchWorkgroupsOp>::add(context);
+    AffinityOpAttrExternalModel<IREE::Flow::DispatchOp>::add(context);
+    AffinityOpAttrExternalModel<IREE::Flow::CallOp>::add(context);
+    AffinityOpAttrExternalModel<IREE::Flow::TensorConstantOp>::add(context);
+    AffinityOpAttrExternalModel<IREE::Flow::TensorDynamicConstantOp>::add(
+        context);
+    AffinityOpAttrExternalModel<IREE::Flow::TensorAllocaOp>::add(context);
+    AffinityOpAttrExternalModel<IREE::Flow::TensorEmptyOp>::add(context);
+    AffinityOpAttrExternalModel<IREE::Flow::TensorSplatOp>::add(context);
+    AffinityOpAttrExternalModel<IREE::Flow::TensorCloneOp>::add(context);
+    AffinityOpAttrExternalModel<IREE::Flow::TensorSliceOp>::add(context);
+    AffinityOpAttrExternalModel<IREE::Flow::TensorUpdateOp>::add(context);
+    AffinityOpAttrExternalModel<IREE::Flow::ChannelDefaultOp>::add(context);
+    AffinityOpAttrExternalModel<IREE::Flow::CollectiveAllGatherOp>::add(
+        context);
+    AffinityOpAttrExternalModel<IREE::Flow::CollectiveAllReduceOp>::add(
+        context);
+    AffinityOpAttrExternalModel<IREE::Flow::CollectiveAllToAllOp>::add(context);
+    AffinityOpAttrExternalModel<IREE::Flow::CollectiveReduceScatterOp>::add(
+        context);
+    AffinityOpAttrExternalModel<IREE::Flow::CollectiveSendRecvOp>::add(context);
+  });
+
   registry.insert<IREE::HAL::HALDialect>();
   registry.addExtension(+[](MLIRContext *context,
                             IREE::HAL::HALDialect *dialect) {
     HALTensorAffinityAttrExternalModel<IREE::HAL::TensorImportOp>::add(context);
     HALTensorAffinityAttrExternalModel<IREE::HAL::TensorExportOp>::add(context);
     HALTensorAffinityAttrExternalModel<IREE::HAL::TensorAliasOp>::add(context);
-    HALTensorAffinityAttrExternalModel<IREE::HAL::TensorBarrierOp>::add(
-        context);
   });
 
   registry.insert<IREE::Util::UtilDialect>();
-  registry.addExtension(
-      +[](MLIRContext *context, IREE::Util::UtilDialect *dialect) {
-        GlobalOpAffinityAttrExternalModel<IREE::Util::GlobalOp>::add(context);
-        AffinityOpAttrExternalModel<IREE::Util::InitializerOp>::add(context);
-        AffinityOpAttrExternalModel<IREE::Util::FuncOp>::add(context);
-      });
+  registry.addExtension(+[](MLIRContext *context,
+                            IREE::Util::UtilDialect *dialect) {
+    GlobalOpAffinityAttrExternalModel<IREE::Util::GlobalOp>::add(context);
+    AffinityOpAttrExternalModel<IREE::Util::InitializerOp, false>::add(context);
+    AffinityOpAttrExternalModel<IREE::Util::FuncOp, false>::add(context);
+  });
 }
 
 } // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Preprocessing/Common/BUILD.bazel b/compiler/src/iree/compiler/Preprocessing/Common/BUILD.bazel
index 81e204e..13bc6bf 100644
--- a/compiler/src/iree/compiler/Preprocessing/Common/BUILD.bazel
+++ b/compiler/src/iree/compiler/Preprocessing/Common/BUILD.bazel
@@ -57,6 +57,7 @@
         "//compiler/src/iree/compiler/Dialect/HAL/Analysis",
         "//compiler/src/iree/compiler/Dialect/HAL/IR",
         "//compiler/src/iree/compiler/Dialect/LinalgExt/IR",
+        "//compiler/src/iree/compiler/Dialect/Stream/Analysis",
         "//compiler/src/iree/compiler/Dialect/Stream/IR",
         "//compiler/src/iree/compiler/Dialect/Util/IR",
         "@llvm-project//llvm:Support",
diff --git a/compiler/src/iree/compiler/Preprocessing/Common/CMakeLists.txt b/compiler/src/iree/compiler/Preprocessing/Common/CMakeLists.txt
index 723dacc..3764d49 100644
--- a/compiler/src/iree/compiler/Preprocessing/Common/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Preprocessing/Common/CMakeLists.txt
@@ -68,6 +68,7 @@
     iree::compiler::Dialect::HAL::Analysis
     iree::compiler::Dialect::HAL::IR
     iree::compiler::Dialect::LinalgExt::IR
+    iree::compiler::Dialect::Stream::Analysis
     iree::compiler::Dialect::Stream::IR
     iree::compiler::Dialect::Util::IR
   PUBLIC
diff --git a/compiler/src/iree/compiler/Preprocessing/Common/PadToIntrinsics.cpp b/compiler/src/iree/compiler/Preprocessing/Common/PadToIntrinsics.cpp
index a5e1a86..ba415b3 100644
--- a/compiler/src/iree/compiler/Preprocessing/Common/PadToIntrinsics.cpp
+++ b/compiler/src/iree/compiler/Preprocessing/Common/PadToIntrinsics.cpp
@@ -12,6 +12,7 @@
 #include "iree/compiler/Codegen/Utils/GPUUtils.h"
 #include "iree/compiler/Dialect/HAL/Analysis/DeviceAnalysis.h"
 #include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
+#include "iree/compiler/Dialect/Stream/Analysis/Affinity.h"
 #include "iree/compiler/Preprocessing/Common/Passes.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
@@ -25,7 +26,7 @@
 
 namespace mlir::iree_compiler::Preprocessing {
 
-#define GEN_PASS_DEF_PADTOINTRINSICS
+#define GEN_PASS_DEF_PADTOINTRINSICSPASS
 #include "iree/compiler/Preprocessing/Common/Passes.h.inc" // IWYU pragma: export
 
 namespace {
@@ -533,7 +534,7 @@
 }
 
 struct PadToIntrinsicsPass
-    : public impl::PadToIntrinsicsBase<PadToIntrinsicsPass> {
+    : public impl::PadToIntrinsicsPassBase<PadToIntrinsicsPass> {
   using Base::Base;
   void runOnOperation() override;
 };
@@ -544,10 +545,15 @@
   MLIRContext *context = &getContext();
   RewritePatternSet patterns(context);
 
-  auto funcOp = getOperation();
-  IREE::HAL::DeviceAnalysis deviceAnalysis(funcOp->getParentOp());
-  if (failed(deviceAnalysis.run()))
+  auto moduleOp = getOperation();
+  IREE::Stream::AffinityAnalysis affinityAnalysis(moduleOp);
+  if (failed(affinityAnalysis.run())) {
     return signalPassFailure();
+  }
+  IREE::HAL::DeviceAnalysis deviceAnalysis(moduleOp);
+  if (failed(deviceAnalysis.run())) {
+    return signalPassFailure();
+  }
 
   bool padConvOps = padTargetType == PadTargetType::ConvOp ||
                     padTargetType == PadTargetType::All;
@@ -555,37 +561,46 @@
                            padTargetType == PadTargetType::All;
   SmallVector<linalg::LinalgOp> targetConvOps;
   SmallVector<linalg::LinalgOp> targetContractOps;
-  funcOp.walk([&](linalg::LinalgOp linalgOp) {
-    if (isa<linalg::Conv2DNhwcHwcfOp>(linalgOp.getOperation()) && padConvOps) {
-      // Add convOps into worklist.
-      targetConvOps.push_back(linalgOp);
-    } else if (isa<linalg::BatchMatmulOp, linalg::MatmulOp,
-                   linalg::MatmulTransposeBOp>(linalgOp.getOperation()) &&
-               padContractionOps) {
-      // Add named contractionOps into worklist.
-      targetContractOps.push_back(linalgOp);
-    } else if (isa<linalg::GenericOp>(linalgOp.getOperation()) &&
-               linalg::isaContractionOpInterface(linalgOp) &&
-               padContractionOps) {
-      // Add named generic contractionOps into worklist.
-      targetContractOps.push_back(linalgOp);
-    }
-  });
+  for (auto funcOp : moduleOp.getOps<FunctionOpInterface>()) {
+    funcOp.walk([&](linalg::LinalgOp linalgOp) {
+      if (isa<linalg::Conv2DNhwcHwcfOp>(linalgOp.getOperation()) &&
+          padConvOps) {
+        targetConvOps.push_back(linalgOp);
+      } else if (isa<linalg::BatchMatmulOp, linalg::MatmulOp,
+                     linalg::MatmulTransposeBOp>(linalgOp.getOperation()) &&
+                 padContractionOps) {
+        targetContractOps.push_back(linalgOp);
+      } else if (isa<linalg::GenericOp>(linalgOp.getOperation()) &&
+                 linalg::isaContractionOpInterface(linalgOp) &&
+                 padContractionOps) {
+        targetContractOps.push_back(linalgOp);
+      }
+    });
+  }
 
   // Iterate through and pad ops in the worklists.
+  auto getRequiredExecutableTargetAttrs = [&](Operation *op) {
+    SetVector<IREE::HAL::ExecutableTargetAttr> executableTargetAttrs;
+    SmallVector<IREE::Stream::AffinityAttr> affinityAttrs;
+    if (affinityAnalysis.tryInferExecutionAffinity(op, affinityAttrs)) {
+      for (auto affinityAttr : affinityAttrs) {
+        deviceAnalysis.gatherRequiredExecutableTargets(affinityAttr, op,
+                                                       executableTargetAttrs);
+      }
+    }
+    return executableTargetAttrs;
+  };
   IRRewriter rewriter(context);
   for (auto convOp : targetConvOps) {
     rewriter.setInsertionPoint(convOp);
-    SetVector<IREE::HAL::ExecutableTargetAttr> executableTargets;
-    deviceAnalysis.gatherRequiredExecutableTargets(convOp, executableTargets);
-    padConvOp(rewriter, convOp, executableTargets.getArrayRef());
+    auto executableTargetAttrs = getRequiredExecutableTargetAttrs(convOp);
+    padConvOp(rewriter, convOp, executableTargetAttrs.getArrayRef());
   }
   for (auto contractOp : targetContractOps) {
     rewriter.setInsertionPoint(contractOp);
-    SetVector<IREE::HAL::ExecutableTargetAttr> executableTargets;
-    deviceAnalysis.gatherRequiredExecutableTargets(contractOp,
-                                                   executableTargets);
-    padContractionLikeOp(rewriter, contractOp, executableTargets.getArrayRef());
+    auto executableTargetAttrs = getRequiredExecutableTargetAttrs(contractOp);
+    padContractionLikeOp(rewriter, contractOp,
+                         executableTargetAttrs.getArrayRef());
   }
 }
 
diff --git a/compiler/src/iree/compiler/Preprocessing/Common/Passes.td b/compiler/src/iree/compiler/Preprocessing/Common/Passes.td
index edc1705..ca29a52 100644
--- a/compiler/src/iree/compiler/Preprocessing/Common/Passes.td
+++ b/compiler/src/iree/compiler/Preprocessing/Common/Passes.td
@@ -84,8 +84,8 @@
   ];
 }
 
-def PadToIntrinsics :
-    InterfacePass<"iree-preprocessing-pad-to-intrinsics", "mlir::FunctionOpInterface"> {
+def PadToIntrinsicsPass :
+    Pass<"iree-preprocessing-pad-to-intrinsics", "ModuleOp"> {
   let summary = "Pad linalg ops such that we can use target's intrinsics.";
   let dependentDialects = [
     "mlir::linalg::LinalgDialect",
diff --git a/compiler/src/iree/compiler/Preprocessing/Common/test/pad_to_intrinsics_mfma.mlir b/compiler/src/iree/compiler/Preprocessing/Common/test/pad_to_intrinsics_mfma.mlir
index 5761741..7d9da45 100644
--- a/compiler/src/iree/compiler/Preprocessing/Common/test/pad_to_intrinsics_mfma.mlir
+++ b/compiler/src/iree/compiler/Preprocessing/Common/test/pad_to_intrinsics_mfma.mlir
@@ -1,6 +1,6 @@
-// RUN: iree-opt --split-input-file %s --iree-gpu-test-target=gfx1100 --pass-pipeline="builtin.module(func.func(iree-preprocessing-pad-to-intrinsics,canonicalize))" | FileCheck %s
-// RUN: iree-opt --split-input-file %s --iree-gpu-test-target=gfx1100 --pass-pipeline="builtin.module(func.func(iree-preprocessing-pad-to-intrinsics{pad-target-type=conv},canonicalize))" | FileCheck %s -check-prefix=CONVOLUTION
-// RUN: iree-opt --split-input-file %s --iree-gpu-test-target=gfx1100 --pass-pipeline="builtin.module(func.func(iree-preprocessing-pad-to-intrinsics{pad-target-type=contraction},canonicalize))" | FileCheck %s -check-prefix=CONTRACT
+// RUN: iree-opt --split-input-file %s --iree-gpu-test-target=gfx1100 --pass-pipeline="builtin.module(iree-preprocessing-pad-to-intrinsics,func.func(canonicalize))" | FileCheck %s
+// RUN: iree-opt --split-input-file %s --iree-gpu-test-target=gfx1100 --pass-pipeline="builtin.module(iree-preprocessing-pad-to-intrinsics{pad-target-type=conv},func.func(canonicalize))" | FileCheck %s -check-prefix=CONVOLUTION
+// RUN: iree-opt --split-input-file %s --iree-gpu-test-target=gfx1100 --pass-pipeline="builtin.module(iree-preprocessing-pad-to-intrinsics{pad-target-type=contraction},func.func(canonicalize))" | FileCheck %s -check-prefix=CONTRACT
 
 
 // CHECK-LABEL: func.func @main0(
diff --git a/compiler/src/iree/compiler/Preprocessing/Common/test/pad_to_intrinsics_wmma.mlir b/compiler/src/iree/compiler/Preprocessing/Common/test/pad_to_intrinsics_wmma.mlir
index f7f8328..ece0283 100644
--- a/compiler/src/iree/compiler/Preprocessing/Common/test/pad_to_intrinsics_wmma.mlir
+++ b/compiler/src/iree/compiler/Preprocessing/Common/test/pad_to_intrinsics_wmma.mlir
@@ -1,6 +1,6 @@
-// RUN: iree-opt --split-input-file %s --iree-gpu-test-target=gfx1100 --pass-pipeline="builtin.module(func.func(iree-preprocessing-pad-to-intrinsics,canonicalize))" | FileCheck %s
-// RUN: iree-opt --split-input-file %s --iree-gpu-test-target=gfx1100 --pass-pipeline="builtin.module(func.func(iree-preprocessing-pad-to-intrinsics{pad-target-type=conv},canonicalize))" | FileCheck %s -check-prefix=CONVOLUTION
-// RUN: iree-opt --split-input-file %s --iree-gpu-test-target=gfx1100 --pass-pipeline="builtin.module(func.func(iree-preprocessing-pad-to-intrinsics{pad-target-type=contraction},canonicalize))" | FileCheck %s -check-prefix=CONTRACT
+// RUN: iree-opt --split-input-file %s --iree-gpu-test-target=gfx1100 --pass-pipeline="builtin.module(iree-preprocessing-pad-to-intrinsics,func.func(canonicalize))" | FileCheck %s
+// RUN: iree-opt --split-input-file %s --iree-gpu-test-target=gfx1100 --pass-pipeline="builtin.module(iree-preprocessing-pad-to-intrinsics{pad-target-type=conv},func.func(canonicalize))" | FileCheck %s -check-prefix=CONVOLUTION
+// RUN: iree-opt --split-input-file %s --iree-gpu-test-target=gfx1100 --pass-pipeline="builtin.module(iree-preprocessing-pad-to-intrinsics{pad-target-type=contraction},func.func(canonicalize))" | FileCheck %s -check-prefix=CONTRACT
 
 //       CHECK: func.func @matmul_static(
 //  CHECK-SAME:    %[[ARG0:.+]]: tensor<10x20xf16>,