Attaching AffinityOpInterface to common early-phase ops.
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamInterfaces.td b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamInterfaces.td
index d1538f4..343d08a 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamInterfaces.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamInterfaces.td
@@ -128,6 +128,20 @@
let methods = [
InterfaceMethod<
/*desc=*/[{
+ Returns whether the op requires an affinity to be assigned.
+ Some ops may represent either host and device operations depending on
+ their operands/results and only sometimes require an affinity.
+ }],
+ /*retTy=*/"bool",
+ /*methodName=*/"requiresAffinity",
+ /*args=*/(ins),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return true;
+ }]
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
Returns the stream affinity for the op, indicating where it should run.
}],
/*retTy=*/"IREE::Stream::AffinityAttr",
diff --git a/compiler/src/iree/compiler/ExternalInterfaces/BUILD.bazel b/compiler/src/iree/compiler/ExternalInterfaces/BUILD.bazel
index 811423d..7bbd7f5 100644
--- a/compiler/src/iree/compiler/ExternalInterfaces/BUILD.bazel
+++ b/compiler/src/iree/compiler/ExternalInterfaces/BUILD.bazel
@@ -17,17 +17,20 @@
srcs = [
"FlowExternalModels.cpp",
"Interfaces.cpp",
+ "StreamExternalModels.cpp",
"UtilExternalModels.cpp",
],
hdrs = [
"FlowExternalModels.h",
"Interfaces.h",
+ "StreamExternalModels.h",
"UtilExternalModels.h",
],
deps = [
"//compiler/src/iree/compiler/Dialect/Encoding/IR",
"//compiler/src/iree/compiler/Dialect/Flow/IR",
"//compiler/src/iree/compiler/Dialect/LinalgExt/IR",
+ "//compiler/src/iree/compiler/Dialect/Stream/IR",
"//compiler/src/iree/compiler/Dialect/Util/IR",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:IR",
diff --git a/compiler/src/iree/compiler/ExternalInterfaces/CMakeLists.txt b/compiler/src/iree/compiler/ExternalInterfaces/CMakeLists.txt
index 9fa28ee..4e2f29a 100644
--- a/compiler/src/iree/compiler/ExternalInterfaces/CMakeLists.txt
+++ b/compiler/src/iree/compiler/ExternalInterfaces/CMakeLists.txt
@@ -16,10 +16,12 @@
HDRS
"FlowExternalModels.h"
"Interfaces.h"
+ "StreamExternalModels.h"
"UtilExternalModels.h"
SRCS
"FlowExternalModels.cpp"
"Interfaces.cpp"
+ "StreamExternalModels.cpp"
"UtilExternalModels.cpp"
DEPS
MLIRArithDialect
@@ -33,6 +35,7 @@
iree::compiler::Dialect::Encoding::IR
iree::compiler::Dialect::Flow::IR
iree::compiler::Dialect::LinalgExt::IR
+ iree::compiler::Dialect::Stream::IR
iree::compiler::Dialect::Util::IR
PUBLIC
)
diff --git a/compiler/src/iree/compiler/ExternalInterfaces/Interfaces.cpp b/compiler/src/iree/compiler/ExternalInterfaces/Interfaces.cpp
index c761ded..9ebefa3 100644
--- a/compiler/src/iree/compiler/ExternalInterfaces/Interfaces.cpp
+++ b/compiler/src/iree/compiler/ExternalInterfaces/Interfaces.cpp
@@ -7,12 +7,14 @@
#include "iree/compiler/ExternalInterfaces/Interfaces.h"
#include "iree/compiler/ExternalInterfaces/FlowExternalModels.h"
+#include "iree/compiler/ExternalInterfaces/StreamExternalModels.h"
#include "iree/compiler/ExternalInterfaces/UtilExternalModels.h"
namespace mlir::iree_compiler {
void registerExternalInterfaces(DialectRegistry ®istry) {
registerFlowExternalModels(registry);
+ registerStreamExternalModels(registry);
registerUtilExternalModels(registry);
}
diff --git a/compiler/src/iree/compiler/ExternalInterfaces/StreamExternalModels.cpp b/compiler/src/iree/compiler/ExternalInterfaces/StreamExternalModels.cpp
new file mode 100644
index 0000000..e3ba257
--- /dev/null
+++ b/compiler/src/iree/compiler/ExternalInterfaces/StreamExternalModels.cpp
@@ -0,0 +1,83 @@
+// 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/ExternalInterfaces/StreamExternalModels.h"
+
+#include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
+#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
+#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
+
+namespace mlir::iree_compiler {
+
+namespace {
+
+template <typename OpT>
+struct AffinityOpAttrExternalModel
+ : public IREE::Stream::AffinityOpInterface::ExternalModel<
+ AffinityOpAttrExternalModel<OpT>, OpT> {
+ static void add(MLIRContext *context) {
+ OpT::template attachInterface<AffinityOpAttrExternalModel<OpT>>(*context);
+ }
+
+ // Most structural ops don't require affinities and after placement we don't
+ // use the affinities even if the ops still exist.
+ bool requiresAffinity(Operation *op) const { return false; }
+
+ IREE::Stream::AffinityAttr getAffinity(Operation *op) const {
+ return op->getAttrOfType<IREE::Stream::AffinityAttr>("stream.affinity");
+ }
+
+ void setAffinity(Operation *op, IREE::Stream::AffinityAttr value) const {
+ if (value)
+ op->setAttr("stream.affinity", value);
+ else
+ op->removeAttr("stream.affinity");
+ }
+};
+
+template <typename OpT>
+struct GlobalOpAffinityAttrExternalModel
+ : public IREE::Stream::AffinityOpInterface::ExternalModel<
+ GlobalOpAffinityAttrExternalModel<OpT>, OpT> {
+ static void add(MLIRContext *context) {
+ OpT::template attachInterface<GlobalOpAffinityAttrExternalModel<OpT>>(
+ *context);
+ }
+
+ // Affinity only required for globals that hold resources that require
+ // placement.
+ bool requiresAffinity(Operation *op) const {
+ auto globalType = cast<IREE::Util::GlobalOpInterface>(op).getGlobalType();
+ return isa<TensorType>(globalType);
+ }
+
+ IREE::Stream::AffinityAttr getAffinity(Operation *op) const {
+ return op->getAttrOfType<IREE::Stream::AffinityAttr>("stream.affinity");
+ }
+
+ void setAffinity(Operation *op, IREE::Stream::AffinityAttr value) const {
+ if (value)
+ op->setAttr("stream.affinity", value);
+ else
+ op->removeAttr("stream.affinity");
+ }
+};
+
+} // namespace
+
+void registerStreamExternalModels(DialectRegistry ®istry) {
+ // Must ensure that any dependent dialects are registered.
+ registry.insert<IREE::Util::UtilDialect>();
+
+ registry.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);
+ });
+}
+
+} // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/ExternalInterfaces/StreamExternalModels.h b/compiler/src/iree/compiler/ExternalInterfaces/StreamExternalModels.h
new file mode 100644
index 0000000..ba29882
--- /dev/null
+++ b/compiler/src/iree/compiler/ExternalInterfaces/StreamExternalModels.h
@@ -0,0 +1,20 @@
+// 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_EXTERNALINTERFACES_STREAMEXTERNALMODELS_H_
+#define IREE_COMPILER_EXTERNALINTERFACES_STREAMEXTERNALMODELS_H_
+
+namespace mlir {
+class DialectRegistry;
+} // namespace mlir
+
+namespace mlir::iree_compiler {
+
+void registerStreamExternalModels(DialectRegistry ®istry);
+
+} // namespace mlir::iree_compiler
+
+#endif // IREE_COMPILER_EXTERNALINTERFACES_STREAMEXTERNALMODELS_H_
