Handling AffinityOpInterface on stream.async.transfer.
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 ®istry) { - // 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);