Skip to content

Commit

Permalink
Clearing up TODOs.
Browse files Browse the repository at this point in the history
  • Loading branch information
@benvanik
benvanik committed Feb 1, 2025
1 parent c659564 commit 5be270d
Show file tree
Hide file tree
Showing 5 changed files with 82 additions and 33 deletions.
33 changes: 28 additions & 5 deletions compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1828,9 +1828,9 @@ LogicalResult TensorSplatOp::verify() {

LogicalResult TensorCloneOp::verify() {
if (failed(verifyOpDynamicDims(getOperation(), {getOperand()},
getArgumentDims())) ||
getOperandDims())) ||
failed(verifyOpDynamicDims(getOperation(), {getResult()},
getArgumentDims()))) {
getOperandDims()))) {
return failure();
}
return success();
Expand All @@ -1840,17 +1840,40 @@ LogicalResult TensorCloneOp::verify() {
// flow.tensor.barrier
//===----------------------------------------------------------------------===//

LogicalResult TensorBarrierOp::verify() { return success(); }
LogicalResult TensorBarrierOp::verify() {
if (failed(verifyOpDynamicDims(getOperation(), {getOperand()},
getOperandDims()))) {
return failure();
}
return success();
}

Value TensorBarrierOp::getTiedResult(unsigned resultIndex) {
return IREE::Util::TiedOpInterface::findTiedBaseValue(getOperand());
}

Value TensorBarrierOp::getTiedResultOperand(Value result) {
return getOperand();
}

::std::optional<unsigned>
TensorBarrierOp::getTiedResultOperandIndex(unsigned resultIndex) {
return {0}; // operand
}

SmallVector<int64_t> TensorBarrierOp::getTiedResultOperandIndices() {
return {0}; // operand
}

//===----------------------------------------------------------------------===//
// flow.tensor.transfer
//===----------------------------------------------------------------------===//

LogicalResult TensorTransferOp::verify() {
if (failed(verifyOpDynamicDims(getOperation(), {getOperand()},
getArgumentDims())) ||
getOperandDims())) ||
failed(verifyOpDynamicDims(getOperation(), {getResult()},
getArgumentDims()))) {
getOperandDims()))) {
return failure();
}
return success();
Expand Down
35 changes: 22 additions & 13 deletions compiler/src/iree/compiler/Dialect/Flow/IR/FlowOps.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1469,14 +1469,14 @@ def FLOW_TensorCloneOp : FLOW_PureOp<"tensor.clone", [

let arguments = (ins
FLOW_Tensor:$operand,
FLOW_ShapeDynamicDims:$argument_dims
FLOW_ShapeDynamicDims:$operand_dims
);
let results = (outs
FLOW_Tensor:$result
);

let assemblyFormat = [{
$operand `:` type($result) (`{` $argument_dims^ `}`)?
$operand `:` type($result) (`{` $operand_dims^ `}`)?
attr-dict-with-keyword
}];

Expand All @@ -1493,8 +1493,8 @@ def FLOW_TensorCloneOp : FLOW_PureOp<"tensor.clone", [
let extraClassDeclaration = [{
bool isHoistableLeafOp() { return false; }

ValueRange getOperandDynamicDims(unsigned idx) { return getArgumentDims(); }
ValueRange getResultDynamicDims(unsigned idx) { return getArgumentDims(); }
ValueRange getOperandDynamicDims(unsigned idx) { return getOperandDims(); }
ValueRange getResultDynamicDims(unsigned idx) { return getOperandDims(); }
}];

let hasVerifier = 1;
Expand All @@ -1506,23 +1506,32 @@ def FLOW_TensorBarrierOp : FLOW_PureOp<"tensor.barrier", [
AllTypesMatch<["operand", "result"]>,
DeclareOpInterfaceMethods<Util_HoistableOpInterface>,
Util_ShapeAwareOp,
DeclareOpInterfaceMethods<Util_TiedOpInterface, [
"getTiedResult",
"getTiedResultOperand",
"getTiedResultOperandIndex",
"getTiedResultOperandIndices",
]>,
]> {
let summary = [{indicates a value that must have a specific affinity}];
let description = [{
DO NOT SUBMIT
Prevents fusion and scheduling of a value across an affinity boundary.
May introduce copy-on-write behavior if the operand value is used as well as
the result and users should try to keep the operand to a single use by this
op.
}];

let arguments = (ins
FLOW_Tensor:$operand,
FLOW_ShapeDynamicDims:$argument_dims,
FLOW_ShapeDynamicDims:$operand_dims,
AnyAttr:$target
);
let results = (outs
FLOW_Tensor:$result
);

let assemblyFormat = [{
$operand `:` type($result) (`{` $argument_dims^ `}`)?
$operand `:` type($result) (`{` $operand_dims^ `}`)?
`on` $target
attr-dict-with-keyword
}];
Expand All @@ -1541,8 +1550,8 @@ def FLOW_TensorBarrierOp : FLOW_PureOp<"tensor.barrier", [
let extraClassDeclaration = [{
bool isHoistableLeafOp() { return false; }

ValueRange getOperandDynamicDims(unsigned idx) { return getArgumentDims(); }
ValueRange getResultDynamicDims(unsigned idx) { return getArgumentDims(); }
ValueRange getOperandDynamicDims(unsigned idx) { return getOperandDims(); }
ValueRange getResultDynamicDims(unsigned idx) { return getOperandDims(); }
}];

let hasVerifier = 1;
Expand All @@ -1565,15 +1574,15 @@ def FLOW_TensorTransferOp : FLOW_PureOp<"tensor.transfer", [

let arguments = (ins
FLOW_Tensor:$operand,
FLOW_ShapeDynamicDims:$argument_dims,
FLOW_ShapeDynamicDims:$operand_dims,
AnyAttr:$target
);
let results = (outs
FLOW_Tensor:$result
);

let assemblyFormat = [{
$operand `:` type($result) (`{` $argument_dims^ `}`)?
$operand `:` type($result) (`{` $operand_dims^ `}`)?
`to` $target
attr-dict-with-keyword
}];
Expand All @@ -1592,8 +1601,8 @@ def FLOW_TensorTransferOp : FLOW_PureOp<"tensor.transfer", [
let extraClassDeclaration = [{
bool isHoistableLeafOp() { return false; }

ValueRange getOperandDynamicDims(unsigned idx) { return getArgumentDims(); }
ValueRange getResultDynamicDims(unsigned idx) { return getArgumentDims(); }
ValueRange getOperandDynamicDims(unsigned idx) { return getOperandDims(); }
ValueRange getResultDynamicDims(unsigned idx) { return getOperandDims(); }
}];

let hasVerifier = 1;
Expand Down
4 changes: 2 additions & 2 deletions compiler/src/iree/compiler/Dialect/Stream/Conversion/FlowToStream/Patterns.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -229,8 +229,8 @@ struct ConvertTensorCloneOp
auto unknownType = rewriter.getType<IREE::Stream::ResourceType>();
auto cloneOp = rewriter.create<IREE::Stream::TensorCloneOp>(
op.getLoc(), unknownType, operand.resource, op.getOperand().getType(),
op.getArgumentDims(), operand.resourceSize, op.getResult().getType(),
flattenValues(adaptor.getArgumentDims()), operand.resourceSize,
op.getOperandDims(), operand.resourceSize, op.getResult().getType(),
flattenValues(adaptor.getOperandDims()), operand.resourceSize,
executionAffinityAttr);
rewriter.replaceOpWithMultiple(op, {{cloneOp, operand.resourceSize}});
return success();
Expand Down
8 changes: 4 additions & 4 deletions compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2306,10 +2306,10 @@ def Stream_AsyncBarrierOp : Stream_Op<"async.barrier", [
]> {
let summary = [{indicates a value that must have a specific affinity}];
let description = [{
DO NOT SUBMIT
tie to force in-place
do operand and result need same affinity?
no execution, so what does it mean?
Prevents fusion and scheduling of a value across an affinity boundary.
May introduce copy-on-write behavior if the operand value is used as well as
the result and users should try to keep the operand to a single use by this
op.
}];

let arguments = (ins
Expand Down
35 changes: 26 additions & 9 deletions compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleAllocation.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1349,6 +1349,28 @@ static Value findTiedYieldResult(Value seedValue) {
return {};
}

// Walks up the use-def chain to find an affinity the given local value is
// pinned to. May return nullptr if there's no assigned affinity and the
// enclosing execution region affinity should be used.
//
// TODO(benvanik): change this to use an affinity analysis on the escaping
// value instead. The local value may not have a transfer associated with it.
static IREE::Stream::AffinityAttr findLocalValueAffinity(Value value) {
while (value) {
if (auto transferOp = dyn_cast_if_present<IREE::Stream::AsyncTransferOp>(
value.getDefiningOp())) {
return transferOp.getResultAffinityAttr();
} else if (auto tiedOp =
llvm::dyn_cast_or_null<IREE::Util::TiedOpInterface>(
value.getDefiningOp())) {
value = tiedOp.getTiedResultOperand(value);
} else {
break;
}
}
return {};
}

// Returns a reversed list of subrange operations that lead from an initial
// resource down a sequence to |derivedValue|. The first element in the list
// will be the last subview of |derivedValue| and the last element will be the
Expand Down Expand Up @@ -1639,15 +1661,10 @@ allocateExecutionRegion(IREE::Stream::AsyncExecuteOp executeOp) {
continue;
}

// DO NOT SUBMIT
// have to get consumer affinity
// may be able to get from a transfer
// could check allocation compatibility with execute op and place there
IREE::Stream::AffinityAttr allocationAffinity;
if (auto transferOp = dyn_cast_if_present<IREE::Stream::AsyncTransferOp>(
yieldValue.getDefiningOp())) {
allocationAffinity = transferOp.getResultAffinityAttr();
} else {
// Find a pinned affinity for the value or inherit the execution region
// affinity.
auto allocationAffinity = findLocalValueAffinity(yieldValue);
if (!allocationAffinity) {
allocationAffinity = executeOp.getAffinityAttr();
}

Expand Down

0 comments on commit 5be270d

Please sign in to comment.