diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/MaterializeTargetDevices.cpp b/compiler/src/iree/compiler/Dialect/HAL/Transforms/MaterializeTargetDevices.cpp
index 395673c..67bf383 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/MaterializeTargetDevices.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/MaterializeTargetDevices.cpp
@@ -166,8 +166,9 @@
     }
 
     if (auto affinityOp = dyn_cast<IREE::Stream::AffinityOpInterface>(op)) {
-      if (!affinityOp.getAffinity())
-        affinityOp.setAffinity(affinityAttr);
+      if (!affinityOp.getAffinityAttr()) {
+        affinityOp.setAffinityAttr(affinityAttr);
+      }
     } else {
       if (!op.hasAttr(affinityName)) {
         op.setAttr(affinityName, affinityAttr);
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning.cpp b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning.cpp
index 5ed2ff8..93fcd37 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning.cpp
@@ -58,9 +58,9 @@
   for (auto *op : ops) {
     if (auto affinityOp = dyn_cast<IREE::Stream::AffinityOpInterface>(op)) {
       if (!IREE::Stream::AffinityAttr::areCompatible(
-              affinity, affinityOp.getAffinity())) {
+              affinity, affinityOp.getAffinityAttr())) {
         return op->emitError("op affinity ")
-               << affinityOp.getAffinity()
+               << affinityOp.getAffinityAttr()
                << " is not compatible with the partition affinity " << affinity;
       }
     }
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning/ReferencePartitioning.cpp b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning/ReferencePartitioning.cpp
index b86ff61..a4fff96 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning/ReferencePartitioning.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning/ReferencePartitioning.cpp
@@ -54,8 +54,8 @@
     DenseSet<Operation *> clonedOps;
     void insert(Operation *op) {
       if (auto affinityOp = dyn_cast<IREE::Stream::AffinityOpInterface>(op)) {
-        affinity = affinity ? affinity.joinAND(affinityOp.getAffinity())
-                            : affinityOp.getAffinity();
+        affinity = affinity ? affinity.joinAND(affinityOp.getAffinityAttr())
+                            : affinityOp.getAffinityAttr();
       }
       ops.insert(op);
     }
@@ -109,7 +109,7 @@
 
     IREE::Stream::AffinityAttr affinityAttr;
     if (auto affinityOp = dyn_cast<IREE::Stream::AffinityOpInterface>(op)) {
-      affinityAttr = affinityOp.getAffinity();
+      affinityAttr = affinityOp.getAffinityAttr();
     }
 
     LLVM_DEBUG({
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp
index 8e4d854..bdc5aaf 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp
@@ -25,13 +25,13 @@
 // size of operands must be queried from the input resource.
 static Value buildResultSizeOf(Location loc, Value tensorValue,
                                ValueRange dynamicDims,
+                               IREE::Stream::AffinityAttr affinityAttr,
                                ConversionPatternRewriter &rewriter) {
   // TODO(benvanik): see if we can stash this on the side to avoid expensive
   // materialization of a bunch of redundant IR.
   return rewriter.create<IREE::Stream::TensorSizeOfOp>(
       loc, rewriter.getIndexType(), TypeAttr::get(tensorValue.getType()),
-      dynamicDims,
-      IREE::Stream::AffinityAttr::lookup(tensorValue.getDefiningOp()));
+      dynamicDims, affinityAttr);
 }
 
 struct ConvertTensorConstantOp
@@ -123,13 +123,14 @@
     auto unknownType = rewriter.getType<IREE::Stream::ResourceType>();
     auto source =
         consumeTensorOperand(op.getLoc(), adaptor.getSource(), rewriter);
-    auto resultSize = buildResultSizeOf(op.getLoc(), op.getResult(),
-                                        op.getResultDims(), rewriter);
+    auto affinityAttr = IREE::Stream::AffinityAttr::lookup(op);
+    auto resultSize =
+        buildResultSizeOf(op.getLoc(), op.getResult(), op.getResultDims(),
+                          affinityAttr, rewriter);
     rewriter.replaceOpWithNewOp<IREE::Stream::TensorCloneOp>(
         op, unknownType, source.resource, op.getSource().getType(),
         op.getSourceDims(), source.resourceSize, op.getResult().getType(),
-        adaptor.getResultDims(), resultSize,
-        IREE::Stream::AffinityAttr::lookup(op));
+        adaptor.getResultDims(), resultSize, affinityAttr);
     return success();
   }
 };
@@ -141,10 +142,12 @@
   matchAndRewrite(IREE::Flow::TensorAllocaOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type unknownType = IREE::Stream::ResourceType::get(getContext());
-    auto resultSize = buildResultSizeOf(op.getLoc(), op.getResult(),
-                                        op.getResultDims(), rewriter);
+    auto affinityAttr = IREE::Stream::AffinityAttr::lookup(op);
+    auto resultSize =
+        buildResultSizeOf(op.getLoc(), op.getResult(), op.getResultDims(),
+                          affinityAttr, rewriter);
     rewriter.replaceOpWithNewOp<IREE::Stream::AsyncAllocaOp>(
-        op, unknownType, resultSize, IREE::Stream::AffinityAttr::lookup(op));
+        op, unknownType, resultSize, affinityAttr);
     return success();
   }
 };
@@ -156,11 +159,13 @@
   matchAndRewrite(IREE::Flow::TensorEmptyOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type unknownType = IREE::Stream::ResourceType::get(getContext());
-    auto resultSize = buildResultSizeOf(op.getLoc(), op.getResult(),
-                                        op.getResultDims(), rewriter);
+    auto affinityAttr = IREE::Stream::AffinityAttr::lookup(op);
+    auto resultSize =
+        buildResultSizeOf(op.getLoc(), op.getResult(), op.getResultDims(),
+                          affinityAttr, rewriter);
     rewriter.replaceOpWithNewOp<IREE::Stream::TensorEmptyOp>(
         op, unknownType, op.getResult().getType(), adaptor.getResultDims(),
-        resultSize, IREE::Stream::AffinityAttr::lookup(op));
+        resultSize, affinityAttr);
     return success();
   }
 };
@@ -172,12 +177,13 @@
   matchAndRewrite(IREE::Flow::TensorSplatOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto unknownType = rewriter.getType<IREE::Stream::ResourceType>();
-    auto resultSize = buildResultSizeOf(op.getLoc(), op.getResult(),
-                                        op.getResultDims(), rewriter);
+    auto affinityAttr = IREE::Stream::AffinityAttr::lookup(op);
+    auto resultSize =
+        buildResultSizeOf(op.getLoc(), op.getResult(), op.getResultDims(),
+                          affinityAttr, rewriter);
     rewriter.replaceOpWithNewOp<IREE::Stream::TensorSplatOp>(
         op, unknownType, adaptor.getValue(), op.getResult().getType(),
-        adaptor.getResultDims(), resultSize,
-        IREE::Stream::AffinityAttr::lookup(op));
+        adaptor.getResultDims(), resultSize, affinityAttr);
     return success();
   }
 };
@@ -230,13 +236,15 @@
     auto unknownType = rewriter.getType<IREE::Stream::ResourceType>();
     auto source =
         consumeTensorOperand(op.getLoc(), adaptor.getSource(), rewriter);
-    auto resultSize = buildResultSizeOf(op.getLoc(), op.getResult(),
-                                        op.getResultDims(), rewriter);
+    auto affinityAttr = IREE::Stream::AffinityAttr::lookup(op);
+    auto resultSize =
+        buildResultSizeOf(op.getLoc(), op.getResult(), op.getResultDims(),
+                          affinityAttr, rewriter);
     rewriter.replaceOpWithNewOp<IREE::Stream::TensorSliceOp>(
         op, unknownType, source.resource, op.getSource().getType(),
         op.getSourceDims(), source.resourceSize, adaptor.getStartIndices(),
         adaptor.getLengths(), op.getResult().getType(), adaptor.getResultDims(),
-        resultSize, IREE::Stream::AffinityAttr::lookup(op));
+        resultSize, affinityAttr);
     return success();
   }
 };
@@ -676,6 +684,8 @@
   LogicalResult
   matchAndRewrite(IREE::Flow::DispatchOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
+    auto affinityAttr = IREE::Stream::AffinityAttr::lookup(op);
+
     // Zero is going to be used for each operand to start.
     auto zeroOffset = rewriter.create<arith::ConstantIndexOp>(op.getLoc(), 0);
 
@@ -723,7 +733,8 @@
         auto resultDynamicDims = IREE::Util::buildDynamicDimsForValue(
             op.getLoc(), result.value(), rewriter);
         resultSizes.push_back(buildResultSizeOf(op.getLoc(), result.value(),
-                                                resultDynamicDims, rewriter));
+                                                resultDynamicDims, affinityAttr,
+                                                rewriter));
         resultTypes.push_back(unknownType);
       }
     }
@@ -732,7 +743,7 @@
         op, resultTypes, adaptor.getWorkload(), adaptor.getEntryPointsAttr(),
         dispatchOperands, dispatchOperandSizes, dispatchOperandOffsets,
         dispatchOperandEnds, dispatchOperandLengths, resultSizes,
-        adaptor.getTiedOperandsAttr(), IREE::Stream::AffinityAttr::lookup(op));
+        adaptor.getTiedOperandsAttr(), affinityAttr);
     newOp->setDialectAttrs(op->getDialectAttrs());
     return success();
   }
@@ -778,6 +789,8 @@
   LogicalResult
   matchAndRewrite(IREE::Flow::CallOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
+    auto affinityAttr = IREE::Stream::AffinityAttr::lookup(op);
+
     // Zero is going to be used for each operand to start.
     auto zeroOffset = rewriter.create<arith::ConstantIndexOp>(op.getLoc(), 0);
 
@@ -825,7 +838,8 @@
         auto resultDynamicDims = IREE::Util::buildDynamicDimsForValue(
             op.getLoc(), result.value(), rewriter);
         resultSizes.push_back(buildResultSizeOf(op.getLoc(), result.value(),
-                                                resultDynamicDims, rewriter));
+                                                resultDynamicDims, affinityAttr,
+                                                rewriter));
         resultTypes.push_back(unknownType);
       }
     }
@@ -834,7 +848,7 @@
         op, resultTypes, adaptor.getCalleeAttr(), callOperands,
         callOperandSizes, callOperandOffsets, callOperandEnds,
         callOperandLengths, resultSizes, adaptor.getTiedOperandsAttr(),
-        IREE::Stream::AffinityAttr::lookup(op));
+        affinityAttr);
     newOp->setDialectAttrs(op->getDialectAttrs());
     return success();
   }
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Conversion/HALToStream/Patterns.cpp b/compiler/src/iree/compiler/Dialect/Stream/Conversion/HALToStream/Patterns.cpp
index 35eb31f..4323473 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Conversion/HALToStream/Patterns.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Conversion/HALToStream/Patterns.cpp
@@ -49,11 +49,7 @@
       }
     }
 
-    auto affinityAttr =
-        dyn_cast_if_present<IREE::Stream::AffinityAttr>(op.getAffinityAttr());
-    if (!affinityAttr) {
-      affinityAttr = IREE::Stream::AffinityAttr::lookup(op);
-    }
+    auto affinityAttr = IREE::Stream::AffinityAttr::lookup(op);
 
     // Import (buffer view to stream resource).
     auto resultType = rewriter.getType<IREE::Stream::ResourceType>(
@@ -138,11 +134,7 @@
       return rewriter.notifyMatchFailure(op, "unsupported HAL cast conversion");
     }
 
-    auto affinityAttr =
-        dyn_cast_if_present<IREE::Stream::AffinityAttr>(op.getAffinityAttr());
-    if (!affinityAttr) {
-      affinityAttr = IREE::Stream::AffinityAttr::lookup(op);
-    }
+    auto affinityAttr = IREE::Stream::AffinityAttr::lookup(op);
     auto source =
         consumeTensorOperand(op.getLoc(), adaptor.getSource(), rewriter);
 
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamInterfaces.td b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamInterfaces.td
index f34003d..f17b753 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamInterfaces.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamInterfaces.td
@@ -145,9 +145,10 @@
         Returns the stream affinity for the op, indicating where it should run.
       }],
       /*retTy=*/"IREE::Stream::AffinityAttr",
-      /*methodName=*/"getAffinity",
+      /*methodName=*/"getAffinityAttr",
       /*args=*/(ins),
-      /*methodBody=*/[{
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
         return dyn_cast_or_null<IREE::Stream::AffinityAttr>($_self->getAttr("affinity"));
       }]
     >,
@@ -156,9 +157,10 @@
         Sets the stream affinity for the op, indicating where it should run.
       }],
       /*retTy=*/"void",
-      /*methodName=*/"setAffinity",
+      /*methodName=*/"setAffinityAttr",
       /*args=*/(ins "IREE::Stream::AffinityAttr":$value),
-      /*methodBody=*/[{
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
         if (value) $_self->setAttr("affinity", value);
         else $_self->removeAttr("affinity");
       }]
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
index d9da282..698c7b9 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
@@ -2019,6 +2019,17 @@
   return success();
 }
 
+IREE::Stream::AffinityAttr AsyncTransferOp::getAffinityAttr() {
+  return getResultAffinityAttr();
+}
+
+void AsyncTransferOp::setAffinityAttr(IREE::Stream::AffinityAttr value) {
+  if (value)
+    setResultAffinityAttr(value);
+  else
+    removeResultAffinityAttr();
+}
+
 void AsyncTransferOp::getAsyncAccessRanges(
     SmallVectorImpl<AsyncAccessRange> &ranges) {
   ranges.push_back({ResourceAccessBitfield::Read, getSource(), Value{},
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
index 99e793a..dbe5207 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
@@ -86,10 +86,7 @@
 let opDocGroup = OpGroupResourceOps in {
 
 def Stream_ResourceAllocOp : Stream_Op<"resource.alloc", [
-  DeclareOpInterfaceMethods<Stream_AffinityOp, [
-    "getAffinity",
-    "setAffinity",
-  ]>,
+  Stream_AffinityOp,
   Util_SizeAwareOp,
   AlwaysSpeculatable,
   MemoryEffects<[MemAlloc]>,
@@ -148,10 +145,7 @@
 }
 
 def Stream_ResourceAllocaOp : Stream_Op<"resource.alloca", [
-  DeclareOpInterfaceMethods<Stream_AffinityOp, [
-    "getAffinity",
-    "setAffinity",
-  ]>,
+  Stream_AffinityOp,
   Stream_TimelineOp,
   Util_SizeAwareOp,
   AlwaysSpeculatable,
@@ -209,10 +203,7 @@
 }
 
 def Stream_ResourceDeallocaOp : Stream_Op<"resource.dealloca", [
-  DeclareOpInterfaceMethods<Stream_AffinityOp, [
-    "getAffinity",
-    "setAffinity",
-  ]>,
+  Stream_AffinityOp,
   Stream_TimelineOp,
   Util_SizeAwareOp,
   MemoryEffects<[MemFree]>,
@@ -645,10 +636,7 @@
 def Stream_ParameterLoadOp : Stream_PureOp<"parameter.load", [
   AttrSizedOperandSegments,
   AllTypesMatch<["results"]>,
-  DeclareOpInterfaceMethods<Stream_AffinityOp, [
-    "getAffinity",
-    "setAffinity",
-  ]>,
+  Stream_AffinityOp,
   Stream_CmdPhaseOp,
   Stream_TimelineOp,
   Util_SizeAwareOp,
@@ -702,10 +690,7 @@
 }
 
 def Stream_ParameterReadOp : Stream_Op<"parameter.read", [
-  DeclareOpInterfaceMethods<Stream_AffinityOp, [
-    "getAffinity",
-    "setAffinity",
-  ]>,
+  Stream_AffinityOp,
   Stream_CmdPhaseOp,
   Stream_TimelineOp,
   Util_SizeAwareOp,
@@ -757,10 +742,7 @@
 }
 
 def Stream_ParameterWriteOp : Stream_Op<"parameter.write", [
-  DeclareOpInterfaceMethods<Stream_AffinityOp, [
-    "getAffinity",
-    "setAffinity",
-  ]>,
+  Stream_AffinityOp,
   Stream_CmdPhaseOp,
   Stream_TimelineOp,
   Util_SizeAwareOp,
@@ -813,10 +795,7 @@
 
 def Stream_ParameterGatherOp : Stream_Op<"parameter.gather", [
   AttrSizedOperandSegments,
-  DeclareOpInterfaceMethods<Stream_AffinityOp, [
-    "getAffinity",
-    "setAffinity",
-  ]>,
+  Stream_AffinityOp,
   Stream_CmdPhaseOp,
   Stream_TimelineOp,
   Util_SizeAwareOp,
@@ -872,10 +851,7 @@
 
 def Stream_ParameterScatterOp : Stream_Op<"parameter.scatter", [
   AttrSizedOperandSegments,
-  DeclareOpInterfaceMethods<Stream_AffinityOp, [
-    "getAffinity",
-    "setAffinity",
-  ]>,
+  Stream_AffinityOp,
   Stream_CmdPhaseOp,
   Stream_TimelineOp,
   Util_SizeAwareOp,
@@ -982,10 +958,7 @@
 }
 
 def Stream_FileReadOp : Stream_Op<"file.read", [
-  DeclareOpInterfaceMethods<Stream_AffinityOp, [
-    "getAffinity",
-    "setAffinity",
-  ]>,
+  Stream_AffinityOp,
   Stream_CmdPhaseOp,
   Stream_TimelineOp,
   Util_SizeAwareOp,
@@ -1040,10 +1013,7 @@
 }
 
 def Stream_FileWriteOp : Stream_Op<"file.write", [
-  DeclareOpInterfaceMethods<Stream_AffinityOp, [
-    "getAffinity",
-    "setAffinity",
-  ]>,
+  Stream_AffinityOp,
   Stream_CmdPhaseOp,
   Stream_TimelineOp,
   Util_SizeAwareOp,
@@ -1783,10 +1753,7 @@
 let opDocGroup = OpGroupAsyncOps in {
 
 def Stream_AsyncAllocaOp : Stream_Op<"async.alloca", [
-  DeclareOpInterfaceMethods<Stream_AffinityOp, [
-    "getAffinity",
-    "setAffinity",
-  ]>,
+  Stream_AffinityOp,
   Stream_AsyncPhaseOp,
   DeclareOpInterfaceMethods<Stream_StreamableOp, [
     "isMetadata",
@@ -2250,7 +2217,10 @@
 }
 
 def Stream_AsyncTransferOp : Stream_Op<"async.transfer", [
-  Stream_AffinityOp,
+  DeclareOpInterfaceMethods<Stream_AffinityOp, [
+    "getAffinityAttr",
+    "setAffinityAttr",
+  ]>,
   Stream_AsyncPhaseOp,
   Stream_StreamableOp,
   DeclareOpInterfaceMethods<Stream_AsyncAccessOp, [
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.cpp
index 7ce79c9..19c2410 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamTypes.cpp
@@ -274,7 +274,7 @@
       return attr;
     // See if the affinity specified provides a resource configuration.
     if (auto affinityOp = llvm::dyn_cast<AffinityOpInterface>(op)) {
-      auto affinityAttr = affinityOp.getAffinity();
+      auto affinityAttr = affinityOp.getAffinityAttr();
       if (affinityAttr) {
         auto attr = affinityAttr.getResourceConfigAttr();
         if (attr)
@@ -339,7 +339,7 @@
   auto attrId = StringAttr::get(op->getContext(), "stream.affinity");
   while (op) {
     if (auto affinityOp = llvm::dyn_cast<AffinityOpInterface>(op)) {
-      auto affinity = affinityOp.getAffinity();
+      auto affinity = affinityOp.getAffinityAttr();
       if (affinity)
         return affinity;
     }
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ConvertToStream.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ConvertToStream.cpp
index aa5cb25..b0b66ac 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ConvertToStream.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ConvertToStream.cpp
@@ -46,6 +46,7 @@
 static Value buildTensorImportOp(Location loc, Value sourceTensor,
                                  Type targetType,
                                  SmallPtrSetImpl<Operation *> &consumingOps,
+                                 IREE::Stream::AffinityAttr affinityAttr,
                                  OpBuilder &builder) {
   // Gather dynamic dimensions from the input value.
   auto dynamicDims =
@@ -56,8 +57,7 @@
   // a transfer operation that may need to reformat the tensor.
   auto encodingAttr = TypeAttr::get(sourceTensor.getType());
   Value resultSize = builder.create<IREE::Stream::TensorSizeOfOp>(
-      loc, builder.getIndexType(), encodingAttr, dynamicDims,
-      /*affinity=*/nullptr);
+      loc, builder.getIndexType(), encodingAttr, dynamicDims, affinityAttr);
 
   // Associate the external SSA value, encoding, and shape information with the
   // stream resource. When lowering we'll then have all the metadata required
@@ -66,7 +66,7 @@
       IREE::Stream::Lifetime::External);
   auto importOp = builder.create<IREE::Stream::TensorImportOp>(
       loc, externalType, sourceTensor, encodingAttr, dynamicDims, resultSize,
-      /*affinity=*/nullptr);
+      affinityAttr);
   consumingOps.insert(importOp);
 
   // If needed insert a transfer to the target lifetime.
@@ -75,8 +75,8 @@
     result = builder
                  .create<IREE::Stream::AsyncTransferOp>(
                      loc, targetType, result, resultSize, resultSize,
-                     /*source_affinity=*/nullptr,
-                     /*result_affinity=*/nullptr)
+                     /*source_affinity=*/affinityAttr,
+                     /*result_affinity=*/affinityAttr)
                  .getResult();
   }
 
@@ -90,6 +90,7 @@
 // external tensor value.
 static Value buildTensorExportOp(Location loc, Value sourceValue,
                                  TensorType targetType, ValueRange dynamicDims,
+                                 IREE::Stream::AffinityAttr affinityAttr,
                                  OpBuilder &builder) {
   auto source = consumeTensorOperand(loc, sourceValue, builder);
 
@@ -101,14 +102,13 @@
         loc, externalType, source.resource, source.resourceSize,
         source.resourceSize,
         /*source_affinity=*/nullptr,
-        /*result_affinity=*/nullptr);
+        /*result_affinity=*/affinityAttr);
   }
 
   // Associate the stream resource and external encoding and shape information.
   auto newOp = builder.create<IREE::Stream::TensorExportOp>(
       loc, targetType, source.resource, TypeAttr::get(targetType), dynamicDims,
-      source.resourceSize,
-      /*affinity=*/nullptr);
+      source.resourceSize, affinityAttr);
   return newOp.getResult();
 }
 
@@ -141,6 +141,8 @@
     if (!doesOperationNeedWrapping(op))
       return failure();
 
+    auto affinityAttr = IREE::Stream::AffinityAttr::lookup(op);
+
     // Export resources into tensor operands for the op to consume.
     SmallVector<Value> newOperands;
     newOperands.reserve(op->getNumOperands());
@@ -156,8 +158,9 @@
 
       auto dynamicDims = IREE::Util::buildDynamicDimsForValue(
           op->getLoc(), oldOperand, rewriter);
-      newOperands.push_back(buildTensorExportOp(
-          op->getLoc(), newOperand, tensorType, dynamicDims, rewriter));
+      newOperands.push_back(buildTensorExportOp(op->getLoc(), newOperand,
+                                                tensorType, dynamicDims,
+                                                affinityAttr, rewriter));
     }
     rewriter.modifyOpInPlace(op, [&]() { op->setOperands(newOperands); });
 
@@ -173,7 +176,7 @@
       SmallPtrSet<Operation *, 4> consumingOps;
       auto importedValue = buildTensorImportOp(
           op->getLoc(), result, rewriter.getType<IREE::Stream::ResourceType>(),
-          consumingOps, rewriter);
+          consumingOps, affinityAttr, rewriter);
       result.replaceAllUsesExcept(importedValue, consumingOps);
     }
 
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/MaterializeCopyOnWrite.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/MaterializeCopyOnWrite.cpp
index 34c0ef8..5b7d3a9 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/MaterializeCopyOnWrite.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/MaterializeCopyOnWrite.cpp
@@ -103,7 +103,7 @@
   IREE::Stream::AffinityAttr affinity;
   if (auto affinityOp =
           dyn_cast<IREE::Stream::AffinityOpInterface>(tiedOp.getOperation())) {
-    affinity = affinityOp.getAffinity();
+    affinity = affinityOp.getAffinityAttr();
   }
 
   // Clones each operand that is tied to a result and it may be required.
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/RefineUsage.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/RefineUsage.cpp
index bc73616..02c2bb0 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/RefineUsage.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/RefineUsage.cpp
@@ -65,7 +65,7 @@
 // Returns either the affinity of |op| or nullptr.
 static IREE::Stream::AffinityAttr getOpAffinity(Operation *op) {
   if (auto affinityOp = dyn_cast<IREE::Stream::AffinityOpInterface>(op)) {
-    return affinityOp.getAffinity();
+    return affinityOp.getAffinityAttr();
   }
   return {};
 }
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleExecution.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleExecution.cpp
index 9c5e3d4..c850c3b 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleExecution.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleExecution.cpp
@@ -152,8 +152,8 @@
     // want to preserve those as long as possible.
     if (auto affinityOp =
             dyn_cast<IREE::Stream::AffinityOpInterface>(clonedOp)) {
-      if (affinityOp.getAffinity() == partition->affinity) {
-        affinityOp.setAffinity(nullptr);
+      if (affinityOp.getAffinityAttr() == partition->affinity) {
+        affinityOp.setAffinityAttr(nullptr);
       }
     }
 
diff --git a/compiler/src/iree/compiler/ExternalInterfaces/BUILD.bazel b/compiler/src/iree/compiler/ExternalInterfaces/BUILD.bazel
index 7bbd7f5..66643e6 100644
--- a/compiler/src/iree/compiler/ExternalInterfaces/BUILD.bazel
+++ b/compiler/src/iree/compiler/ExternalInterfaces/BUILD.bazel
@@ -29,6 +29,7 @@
     deps = [
         "//compiler/src/iree/compiler/Dialect/Encoding/IR",
         "//compiler/src/iree/compiler/Dialect/Flow/IR",
+        "//compiler/src/iree/compiler/Dialect/HAL/IR",
         "//compiler/src/iree/compiler/Dialect/LinalgExt/IR",
         "//compiler/src/iree/compiler/Dialect/Stream/IR",
         "//compiler/src/iree/compiler/Dialect/Util/IR",
diff --git a/compiler/src/iree/compiler/ExternalInterfaces/CMakeLists.txt b/compiler/src/iree/compiler/ExternalInterfaces/CMakeLists.txt
index 4e2f29a..a63fca3 100644
--- a/compiler/src/iree/compiler/ExternalInterfaces/CMakeLists.txt
+++ b/compiler/src/iree/compiler/ExternalInterfaces/CMakeLists.txt
@@ -34,6 +34,7 @@
     MLIRValueBoundsOpInterface
     iree::compiler::Dialect::Encoding::IR
     iree::compiler::Dialect::Flow::IR
+    iree::compiler::Dialect::HAL::IR
     iree::compiler::Dialect::LinalgExt::IR
     iree::compiler::Dialect::Stream::IR
     iree::compiler::Dialect::Util::IR
diff --git a/compiler/src/iree/compiler/ExternalInterfaces/StreamExternalModels.cpp b/compiler/src/iree/compiler/ExternalInterfaces/StreamExternalModels.cpp
index e3ba257..b82d599 100644
--- a/compiler/src/iree/compiler/ExternalInterfaces/StreamExternalModels.cpp
+++ b/compiler/src/iree/compiler/ExternalInterfaces/StreamExternalModels.cpp
@@ -6,6 +6,10 @@
 
 #include "iree/compiler/ExternalInterfaces/StreamExternalModels.h"
 
+#include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
+#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
+#include "iree/compiler/Dialect/HAL/IR/HALDialect.h"
+#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
 #include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
 #include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
 #include "iree/compiler/Dialect/Util/IR/UtilOps.h"
@@ -14,27 +18,47 @@
 
 namespace {
 
-template <typename OpT>
-struct AffinityOpAttrExternalModel
+struct FlowTransferTargetAffinityAttrExternalModel
     : public IREE::Stream::AffinityOpInterface::ExternalModel<
-          AffinityOpAttrExternalModel<OpT>, OpT> {
+          FlowTransferTargetAffinityAttrExternalModel,
+          IREE::Flow::TensorTransferOp> {
   static void add(MLIRContext *context) {
-    OpT::template attachInterface<AffinityOpAttrExternalModel<OpT>>(*context);
+    IREE::Flow::TensorTransferOp::attachInterface<
+        FlowTransferTargetAffinityAttrExternalModel>(*context);
   }
 
-  // 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 true; }
+
+  IREE::Stream::AffinityAttr getAffinity(Operation *op) const {
+    return op->getAttrOfType<IREE::Stream::AffinityAttr>("target");
+  }
+
+  void setAffinity(Operation *op, IREE::Stream::AffinityAttr value) const {
+    op->setAttr("target", value);
+  }
+};
+
+template <typename OpT>
+struct HALTensorAffinityAttrExternalModel
+    : public IREE::Stream::AffinityOpInterface::ExternalModel<
+          HALTensorAffinityAttrExternalModel<OpT>, OpT> {
+  static void add(MLIRContext *context) {
+    OpT::template attachInterface<HALTensorAffinityAttrExternalModel<OpT>>(
+        *context);
+  }
+
   bool requiresAffinity(Operation *op) const { return false; }
 
   IREE::Stream::AffinityAttr getAffinity(Operation *op) const {
-    return op->getAttrOfType<IREE::Stream::AffinityAttr>("stream.affinity");
+    return op->getAttrOfType<IREE::Stream::AffinityAttr>("affinity");
   }
 
   void setAffinity(Operation *op, IREE::Stream::AffinityAttr value) const {
     if (value)
-      op->setAttr("stream.affinity", value);
-    else
-      op->removeAttr("stream.affinity");
+      op->setAttr("affinity", value);
+    } else {
+      op->removeAttr("affinity");
+    }
   }
 };
 
@@ -61,17 +85,58 @@
   void setAffinity(Operation *op, IREE::Stream::AffinityAttr value) const {
     if (value)
       op->setAttr("stream.affinity", value);
-    else
+    } else {
       op->removeAttr("stream.affinity");
+    }
+  }
+};
+
+template <typename OpT>
+struct AffinityOpAttrExternalModel
+    : public IREE::Stream::AffinityOpInterface::ExternalModel<
+          AffinityOpAttrExternalModel<OpT, kRequiresAffinity>, OpT> {
+  static void add(MLIRContext *context) {
+    OpT::template attachInterface<
+        AffinityOpAttrExternalModel<OpT, kRequiresAffinity>>(*context);
+  }
+
+  // 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; }
+
+  IREE::Stream::AffinityAttr getAffinity(Operation *op) const {
+    return op->getAttrOfType<IREE::Stream::AffinityAttr>("stream.affinity");
+  }
+
+  void setAffinity(Operation *op, IREE::Stream::AffinityAttr value) const {
+    if (value)
+      op->setAttr("stream.affinity", value);
+    } else {
+      op->removeAttr("stream.affinity");
+    }
   }
 };
 
 } // namespace
 
 void registerStreamExternalModels(DialectRegistry &registry) {
-  // Must ensure that any dependent dialects are registered.
-  registry.insert<IREE::Util::UtilDialect>();
+  registry.insert<IREE::Flow::FlowDialect>();
+  registry.addExtension(
+      +[](MLIRContext *context, IREE::Flow::FlowDialect *dialect) {
+        FlowTransferTargetAffinityAttrExternalModel::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);