Adding affinity analysis.
This performs whole-program analysis to enable the querying of the
ideal affinity for globals, execution ops, and resources. It can run at
most phases of compilation (including on linalg/flow IR) though it's
primarily used by the stream dialect passes such as conversion.
The `AnnotateAffinitiesPass` has been added to aid debugging and the
compiler `iree-stream-annotate-input-affinities` flag can be used to
turn it on - it has no impact on the program generated but can be useful
if affinity analysis fails during conversion.
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 ®ion : 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 ®istry) {
- 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>,
