Adding `#hal.device.alias` attribute for resolving device configs.
This allows for less verbose "I don't care, pick something for me"
attributes that are expanded into the full target devices and their
executable configurations. Resolution happens early in the process so
that any flags that may be influencing the resolved configurations are
captured and no longer required by the pipeline.
Tests and tooling could use these attributes in place of
`#hal.device.target` but would need to run the pass as part of their
pipeline in order to perform the expansion. Resolving in a pass
vs doing so inline also allows for signaling errors and passing in
scoped device target registries instead of relying on the globals that
are not available in API usage.
diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/HALAttrs.cpp b/compiler/src/iree/compiler/Dialect/HAL/IR/HALAttrs.cpp
index 7bab983..6d3c4a2 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/IR/HALAttrs.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/IR/HALAttrs.cpp
@@ -86,225 +86,6 @@
}
//===----------------------------------------------------------------------===//
-// #hal.device.target<*>
-//===----------------------------------------------------------------------===//
-
-// static
-DeviceTargetAttr DeviceTargetAttr::get(MLIRContext *context,
- StringRef deviceID) {
- // TODO(benvanik): query default configuration from the target backend.
- return get(context, StringAttr::get(context, deviceID),
- DictionaryAttr::get(context), {});
-}
-
-// static
-Attribute DeviceTargetAttr::parse(AsmParser &p, Type type) {
- StringAttr deviceIDAttr;
- DictionaryAttr configAttr;
- SmallVector<IREE::HAL::ExecutableTargetAttr> executableTargetAttrs;
- // `<"device-id"`
- if (failed(p.parseLess()) || failed(p.parseAttribute(deviceIDAttr))) {
- return {};
- }
- // `, `
- if (succeeded(p.parseOptionalComma())) {
- if (succeeded(p.parseOptionalLSquare())) {
- // `[targets, ...]` (optional)
- do {
- IREE::HAL::ExecutableTargetAttr executableTargetAttr;
- if (failed(p.parseAttribute(executableTargetAttr)))
- return {};
- executableTargetAttrs.push_back(executableTargetAttr);
- } while (succeeded(p.parseOptionalComma()));
- if (failed(p.parseRSquare()))
- return {};
- } else {
- // `{config dict}` (optional)
- if (failed(p.parseAttribute(configAttr)))
- return {};
- // `, [targets, ...]` (optional)
- if (succeeded(p.parseOptionalComma())) {
- if (failed(p.parseLSquare()))
- return {};
- do {
- IREE::HAL::ExecutableTargetAttr executableTargetAttr;
- if (failed(p.parseAttribute(executableTargetAttr)))
- return {};
- executableTargetAttrs.push_back(executableTargetAttr);
- } while (succeeded(p.parseOptionalComma()));
- if (failed(p.parseRSquare()))
- return {};
- }
- }
- }
- // `>`
- if (failed(p.parseGreater())) {
- return {};
- }
- return get(p.getContext(), deviceIDAttr, configAttr, executableTargetAttrs);
-}
-
-void DeviceTargetAttr::print(AsmPrinter &p) const {
- auto &os = p.getStream();
- os << "<";
- p.printAttribute(getDeviceID());
- auto configAttr = getConfiguration();
- if (configAttr && !configAttr.empty()) {
- os << ", ";
- p.printAttribute(configAttr);
- }
- auto executableTargetAttrs = getExecutableTargets();
- if (!executableTargetAttrs.empty()) {
- os << ", [";
- llvm::interleaveComma(executableTargetAttrs, os,
- [&](auto executableTargetAttr) {
- p.printAttribute(executableTargetAttr);
- });
- os << "]";
- }
- os << ">";
-}
-
-std::string DeviceTargetAttr::getSymbolNameFragment() {
- return sanitizeSymbolName(getDeviceID().getValue().lower());
-}
-
-bool DeviceTargetAttr::hasConfigurationAttr(StringRef name) {
- auto configAttr = getConfiguration();
- return configAttr && configAttr.get(name);
-}
-
-void DeviceTargetAttr::getExecutableTargets(
- SetVector<IREE::HAL::ExecutableTargetAttr> &resultAttrs) {
- for (auto attr : getExecutableTargets()) {
- resultAttrs.insert(attr);
- }
-}
-
-void IREE::HAL::DeviceTargetAttr::printStatusDescription(
- llvm::raw_ostream &os) const {
- cast<Attribute>().print(os, /*elideType=*/true);
-}
-
-// Produces a while-loop that enumerates each device available and tries to
-// match it against the target information. SCF is... not very wieldy, but this
-// is effectively:
-// ```
-// %device_count = hal.devices.count : index
-// %result:3 = scf.while(%i = 0, %match_ordinal = 0, %device = null) {
-// %is_null = util.cmp.eq %device, null : !hal.device
-// %in_bounds = arith.cmpi slt %i, %device_count : index
-// %continue_while = arith.andi %is_null, %in_bounds : i1
-// scf.condition(%continue_while) %i, %match_ordinal %device
-// : index, index, !hal.device
-// } do {
-// %device_i = hal.devices.get %i : !hal.device
-// %device_match = <<buildDeviceMatch>>(%device_i)
-// %ordinal_match = arith.cmpi eq %match_ordinal, %device_ordinal : index
-// %is_match = arith.andi %device_match, %ordinal_match : i1
-// %try_device = arith.select %is_match, %device_i, null : !hal.device
-// %next_i = arith.addi %i, %c1 : index
-// %match_adv = arith.select %device_match, %c1, %c0 : index
-// %next_match_ordinal = arith.addi %match_ordinal, %match_adv : index
-// scf.yield %next_i, %next_match_ordinal, %try_device
-// : index, index !hal.device
-// }
-// ```
-// Upon completion %result#1 contains the device (or null).
-// If the target had an ordinal specified we skip matches until a match with the
-// specified ordinal is reached.
-Value IREE::HAL::DeviceTargetAttr::buildDeviceEnumeration(
- Location loc, const IREE::HAL::TargetRegistry &targetRegistry,
- OpBuilder &builder) const {
- // Device configuration can control selection beyond just the match
- // expression.
- auto configAttr = getConfiguration();
- IntegerAttr deviceOrdinalAttr =
- configAttr ? configAttr.getAs<IntegerAttr>("ordinal") : IntegerAttr{};
-
- // Defers to the target backend to build the device match or does a simple
- // fallback for unregistered backends (usually for testing, but may be used
- // as a way to bypass validation for out-of-tree experiments).
- auto buildDeviceMatch = [&](Location loc, Value device,
- OpBuilder &builder) -> Value {
- // Ask the target backend to build the match expression. It may opt to
- // let the default handling take care of things.
- Value match;
- auto targetDevice = targetRegistry.getTargetDevice(getDeviceID());
- if (targetDevice)
- match = targetDevice->buildDeviceTargetMatch(loc, device, *this, builder);
- if (match)
- return match;
- return buildDeviceIDAndExecutableFormatsMatch(
- loc, device, getDeviceID(), getExecutableTargets(), builder);
- };
-
- // Enumerate all devices and match the first one found (if any).
- Type indexType = builder.getIndexType();
- Type deviceType = builder.getType<IREE::HAL::DeviceType>();
- Value c0 = builder.create<arith::ConstantIndexOp>(loc, 0);
- Value c1 = builder.create<arith::ConstantIndexOp>(loc, 1);
- Value nullDevice = builder.create<IREE::Util::NullOp>(loc, deviceType);
- Value deviceOrdinal = deviceOrdinalAttr
- ? builder.create<arith::ConstantIndexOp>(
- loc, deviceOrdinalAttr.getInt())
- : c0;
- Value deviceCount = builder.create<IREE::HAL::DevicesCountOp>(loc, indexType);
- auto whileOp = builder.create<scf::WhileOp>(
- loc,
- TypeRange{
- /*i=*/indexType,
- /*match_ordinal=*/indexType,
- /*device=*/deviceType,
- },
- ValueRange{
- /*i=*/c0,
- /*match_ordinal=*/c0,
- /*device=*/nullDevice,
- },
- [&](OpBuilder &beforeBuilder, Location loc, ValueRange operands) {
- Value isNull = beforeBuilder.create<IREE::Util::CmpEQOp>(
- loc, operands[/*device=*/2], nullDevice);
- Value inBounds = beforeBuilder.create<arith::CmpIOp>(
- loc, arith::CmpIPredicate::slt, operands[/*i=*/0], deviceCount);
- Value continueWhile =
- beforeBuilder.create<arith::AndIOp>(loc, isNull, inBounds);
- beforeBuilder.create<scf::ConditionOp>(loc, continueWhile, operands);
- },
- [&](OpBuilder &afterBuilder, Location loc, ValueRange operands) {
- // Check whether the device is a match.
- Value device = afterBuilder.create<IREE::HAL::DevicesGetOp>(
- loc, deviceType, operands[/*i=*/0]);
- Value isDeviceMatch = buildDeviceMatch(loc, device, afterBuilder);
-
- // Check whether whether this matching device ordinal is the requested
- // ordinal out of all matching devices.
- Value isOrdinalMatch = afterBuilder.create<arith::CmpIOp>(
- loc, arith::CmpIPredicate::eq, operands[/*match_ordinal=*/1],
- deviceOrdinal);
- Value nextMatchOrdinal = afterBuilder.create<arith::AddIOp>(
- loc, operands[/*match_ordinal=*/1],
- afterBuilder.create<arith::SelectOp>(loc, isDeviceMatch, c1, c0));
-
- // Break if the device and ordinal match, otherwise continue with null.
- Value isMatch = afterBuilder.create<arith::AndIOp>(loc, isDeviceMatch,
- isOrdinalMatch);
- Value tryDevice = afterBuilder.create<arith::SelectOp>(
- loc, isMatch, device, nullDevice);
-
- Value nextI =
- afterBuilder.create<arith::AddIOp>(loc, operands[/*i=*/0], c1);
- afterBuilder.create<scf::YieldOp>(
- loc, ValueRange{
- /*i=*/nextI,
- /*match_ordinal=*/nextMatchOrdinal,
- /*device=*/tryDevice,
- });
- });
- return whileOp.getResult(/*device=*/2);
-}
-
-//===----------------------------------------------------------------------===//
// #hal.executable.target<*>
//===----------------------------------------------------------------------===//
@@ -946,16 +727,25 @@
//===----------------------------------------------------------------------===//
// static
+DeviceSelectAttr DeviceSelectAttr::get(MLIRContext *context,
+ ArrayRef<Attribute> values) {
+ return DeviceSelectAttr::get(context, IREE::HAL::DeviceType::get(context),
+ ArrayAttr::get(context, values));
+}
+
+// static
LogicalResult
DeviceSelectAttr::verify(function_ref<mlir::InFlightDiagnostic()> emitError,
Type type, ArrayAttr devicesAttr) {
if (devicesAttr.empty())
return emitError() << "must have at least one device to select";
for (auto deviceAttr : devicesAttr) {
- if (!deviceAttr.isa<IREE::HAL::DeviceInitializationAttrInterface>()) {
- return emitError() << "can only select between #hal.device.target, "
- "#hal.device.ordinal, #hal.device.fallback, or "
- "other device initialization attributes";
+ if (!mlir::isa<IREE::HAL::DeviceAliasAttr>(deviceAttr) &&
+ !mlir::isa<IREE::HAL::DeviceInitializationAttrInterface>(deviceAttr)) {
+ return emitError() << "can only select between #hal.device.alias, "
+ "#hal.device.target, #hal.device.ordinal, "
+ "#hal.device.fallback, or other device "
+ "initialization attributes";
}
}
// TODO(benvanik): when !hal.device is parameterized we should check that the
diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/HALAttrs.td b/compiler/src/iree/compiler/Dialect/HAL/IR/HALAttrs.td
index cc7c523..2d10dc3 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/IR/HALAttrs.td
+++ b/compiler/src/iree/compiler/Dialect/HAL/IR/HALAttrs.td
@@ -477,65 +477,6 @@
"HAL binding array attribute">;
//===----------------------------------------------------------------------===//
-// #hal.device.target<*>
-//===----------------------------------------------------------------------===//
-
-def HAL_DeviceTargetAttr : AttrDef<HAL_Dialect, "DeviceTarget", [
- DeclareAttrInterfaceMethods<HAL_DeviceInitializationAttrInterface>,
-]> {
- let mnemonic = "device.target";
- let summary = [{generic device target specification}];
- let description = [{
- Specifies the properties of a target runtime device.
- Target devices are specified with a canonical identifier matching those used
- by the runtime (such as `cpu`, `vulkan`, etc). Target devices may support
- several target executable formats specified with `#hal.executable.target`.
- An optional configuration dictionary allows for overriding backend defaults.
-
- If used to initialize a device global returns the first device matching the
- target requirements or null if no devices match. An optional `ordinal`
- index may be provided that selects the N-th matching device and is used to
- select between multiple homogeneous devices.
-
- Example:
- ```mlir
- #hal.device.target<"llvm-cpu", {
- device_configuration = ...
- }, [
- #hal.executable.target<"llvm-cpu", "embedded-elf-arm_32">,
- #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64">,
- ]>
- ```
- }];
- let parameters = (ins
- AttrParameter<"StringAttr", "">:$deviceID,
- AttrParameter<"DictionaryAttr", "">:$configuration,
- ArrayRefParameter<"ExecutableTargetAttr", "">:$executable_targets
- );
- let builders = [
- AttrBuilder<(ins "StringRef":$deviceID)>,
- ];
-
- let extraClassDeclaration = [{
- Type getType() { return IREE::HAL::DeviceType::get(getContext()); }
-
- // Returns a symbol-compatible name that pseudo-uniquely identifies this
- // target. Callers must perform deduplication when required.
- std::string getSymbolNameFragment();
-
- // Returns true if there's an attribute with the given name in the
- // configuration dictionary.
- bool hasConfigurationAttr(StringRef name);
-
- // Returns zero or more executable targets that this device supports.
- void getExecutableTargets(
- SetVector<IREE::HAL::ExecutableTargetAttr> &resultAttrs);
- }];
-
- let hasCustomAssemblyFormat = 1;
-}
-
-//===----------------------------------------------------------------------===//
// #hal.executable.target<*>
//===----------------------------------------------------------------------===//
@@ -931,9 +872,6 @@
let builders = [
AttrBuilder<(ins
- "IREE::HAL::DeviceTargetAttr":$device
- )>,
- AttrBuilder<(ins
"ArrayRef<Attribute>":$values
)>,
];
diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/test/attributes.mlir b/compiler/src/iree/compiler/Dialect/HAL/IR/test/attributes.mlir
index 3a05d9c..ec0cbdd 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/IR/test/attributes.mlir
+++ b/compiler/src/iree/compiler/Dialect/HAL/IR/test/attributes.mlir
@@ -61,6 +61,20 @@
// -----
+// CHECK-LABEL: "device.aliases"
+"device.aliases"() {
+ // CHECK-SAME: alias_0 = #hal.device.alias<"a"> : !hal.device
+ alias_0 = #hal.device.alias<"a"> : !hal.device,
+ // CHECK-SAME: alias_1 = #hal.device.alias<"b", {}> : !hal.device
+ alias_1 = #hal.device.alias<"b", {}> : !hal.device,
+ // CHECK-SAME: alias_2 = #hal.device.alias<"c"[4]> : !hal.device
+ alias_2 = #hal.device.alias<"c"[4]> : !hal.device,
+ // CHECK-SAME: alias_3 = #hal.device.alias<"d", {config = 123 : index}>
+ alias_3 = #hal.device.alias<"d", {config = 123 : index}> : !hal.device
+} : () -> ()
+
+// -----
+
// CHECK-LABEL: "device.targets"
"device.targets"() {
// CHECK-SAME: target_0 = #hal.device.target<"a"> : !hal.device
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Dialect/HAL/Transforms/BUILD.bazel
index e5be69d..bdad427 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/BUILD.bazel
@@ -35,6 +35,7 @@
"PreprocessExecutables.cpp",
"PruneExecutables.cpp",
"RepeatDispatches.cpp",
+ "ResolveDeviceAliases.cpp",
"ResolveDevicePromises.cpp",
"ResolveExportOrdinals.cpp",
"SerializeExecutables.cpp",
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt
index c7b1207..72b0b74 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt
@@ -36,6 +36,7 @@
"PreprocessExecutables.cpp"
"PruneExecutables.cpp"
"RepeatDispatches.cpp"
+ "ResolveDeviceAliases.cpp"
"ResolveDevicePromises.cpp"
"ResolveExportOrdinals.cpp"
"SerializeExecutables.cpp"
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
index a58a51f..dc6ab9d 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
@@ -200,6 +200,8 @@
}
passManager.addPass(IREE::HAL::createMaterializeTargetDevicesPass());
passManager.addPass(IREE::HAL::createResolveDevicePromisesPass());
+ passManager.addPass(
+ IREE::HAL::createResolveDeviceAliasesPass({&targetRegistry}));
passManager.addPass(IREE::HAL::createVerifyDevicesPass({&targetRegistry}));
}
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.td b/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.td
index 8ff1cf7..aa896f2 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.td
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.td
@@ -91,6 +91,25 @@
];
}
+def ResolveDeviceAliasesPass :
+ Pass<"iree-hal-resolve-device-aliases", "mlir::ModuleOp"> {
+ let summary = "Resolves `#hal.device.alias` attributes to their expanded configurations.";
+ let description = [{
+ Resolves device aliases to the concrete targets using defaults, flags, and
+ registered device configurations.
+ }];
+ let options = [
+ Option<
+ "targetRegistry", "target-registry",
+ "llvm::cl::TargetRegistryRef", "",
+ "Target registry containing the list of available devices and backends."
+ >,
+ ];
+ let dependentDialects = [
+ "IREE::HAL::HALDialect",
+ ];
+}
+
def VerifyDevicesPass :
Pass<"iree-hal-verify-devices", "mlir::ModuleOp"> {
let summary = "Verifies that all devices can be targeted with the available compiler plugins.";
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/ResolveDeviceAliases.cpp b/compiler/src/iree/compiler/Dialect/HAL/Transforms/ResolveDeviceAliases.cpp
new file mode 100644
index 0000000..0108aee
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/ResolveDeviceAliases.cpp
@@ -0,0 +1,134 @@
+// 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 <memory>
+#include <utility>
+
+#include "iree/compiler/Dialect/HAL/IR/HALDialect.h"
+#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
+#include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h"
+#include "iree/compiler/Dialect/HAL/Transforms/Passes.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir::iree_compiler::IREE::HAL {
+
+#define GEN_PASS_DEF_RESOLVEDEVICEALIASESPASS
+#include "iree/compiler/Dialect/HAL/Transforms/Passes.h.inc"
+
+namespace {
+
+//===----------------------------------------------------------------------===//
+// --iree-hal-resolve-device-aliases
+//===----------------------------------------------------------------------===//
+
+static FailureOr<Attribute>
+resolveAliasAttr(Operation *forOp, IREE::HAL::DeviceAliasAttr aliasAttr,
+ const TargetRegistry &targetRegistry) {
+ // Lookup device in the registry.
+ auto targetDevice =
+ targetRegistry.getTargetDevice(aliasAttr.getDeviceID().getValue());
+ if (!targetDevice) {
+ auto diagnostic = forOp->emitError();
+ diagnostic << "unregistered device alias " << aliasAttr.getDeviceID()
+ << "; ensure it is linked into the compiler (available = [ ";
+ for (const auto &targetName : targetRegistry.getRegisteredTargetDevices()) {
+ diagnostic << "'" << targetName << "' ";
+ }
+ diagnostic << "])";
+ return diagnostic;
+ }
+
+ // Query the default device target.
+ auto defaultAttr =
+ targetDevice->getDefaultDeviceTarget(forOp->getContext(), targetRegistry);
+ assert(defaultAttr && "expected a default device target attr");
+
+ // Merge in any additional configuration from the alias attr.
+ if (aliasAttr.getOrdinal().has_value() ||
+ (aliasAttr.getConfiguration() && !aliasAttr.getConfiguration().empty())) {
+ NamedAttrList configAttrs;
+ if (auto defaultConfigAttr = defaultAttr.getConfiguration()) {
+ for (auto existingAttr : defaultConfigAttr) {
+ configAttrs.push_back(existingAttr);
+ }
+ }
+ if (auto overrideConfigAttr = aliasAttr.getConfiguration()) {
+ for (auto overrideAttr : overrideConfigAttr) {
+ configAttrs.set(overrideAttr.getName(), overrideAttr.getValue());
+ }
+ }
+ if (aliasAttr.getOrdinal().has_value()) {
+ configAttrs.set("ordinal",
+ IntegerAttr::get(IndexType::get(forOp->getContext()),
+ aliasAttr.getOrdinal().value()));
+ }
+ defaultAttr = IREE::HAL::DeviceTargetAttr::get(
+ forOp->getContext(), defaultAttr.getDeviceID(),
+ DictionaryAttr::get(forOp->getContext(), configAttrs),
+ defaultAttr.getExecutableTargets());
+ }
+
+ return defaultAttr;
+}
+
+static FailureOr<Attribute>
+resolveNestedAliasAttrs(Operation *forOp, Attribute attr,
+ const TargetRegistry &targetRegistry) {
+ if (auto aliasAttr = dyn_cast<IREE::HAL::DeviceAliasAttr>(attr)) {
+ return resolveAliasAttr(forOp, aliasAttr, targetRegistry);
+ } else if (auto selectAttr = dyn_cast<IREE::HAL::DeviceSelectAttr>(attr)) {
+ SmallVector<Attribute> resolvedAttrs;
+ bool didChange = false;
+ for (auto deviceAttr : selectAttr.getDevices()) {
+ auto resolvedAttr =
+ resolveNestedAliasAttrs(forOp, deviceAttr, targetRegistry);
+ if (failed(resolvedAttr)) {
+ return failure();
+ }
+ didChange = didChange || *resolvedAttr != deviceAttr;
+ resolvedAttrs.push_back(*resolvedAttr);
+ }
+ return didChange ? IREE::HAL::DeviceSelectAttr::get(attr.getContext(),
+ resolvedAttrs)
+ : attr;
+ } else {
+ return attr; // pass-through
+ }
+}
+
+struct ResolveDeviceAliasesPass
+ : public IREE::HAL::impl::ResolveDeviceAliasesPassBase<
+ ResolveDeviceAliasesPass> {
+ using IREE::HAL::impl::ResolveDeviceAliasesPassBase<
+ ResolveDeviceAliasesPass>::ResolveDeviceAliasesPassBase;
+ void runOnOperation() override {
+ // Walks all device globals and resolve any aliases found.
+ auto moduleOp = getOperation();
+ for (auto globalOp : moduleOp.getOps<IREE::Util::GlobalOpInterface>()) {
+ if (!isa<IREE::HAL::DeviceType>(globalOp.getGlobalType())) {
+ continue;
+ }
+ auto initialValue = globalOp.getGlobalInitialValue();
+ if (!initialValue) {
+ continue;
+ }
+ auto resolvedValue = resolveNestedAliasAttrs(globalOp, initialValue,
+ *targetRegistry.value);
+ if (failed(resolvedValue)) {
+ return signalPassFailure();
+ }
+ globalOp.setGlobalInitialValue(*resolvedValue);
+ }
+ }
+};
+
+} // namespace
+
+} // namespace mlir::iree_compiler::IREE::HAL
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/VerifyDevices.cpp b/compiler/src/iree/compiler/Dialect/HAL/Transforms/VerifyDevices.cpp
index 21f3754..e1ca624 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/VerifyDevices.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/VerifyDevices.cpp
@@ -50,7 +50,7 @@
auto diagnostic = deviceOp->emitError();
diagnostic << "unregistered target device "
<< deviceTargetAttr.getDeviceID()
- << "; ensure it is linked in to the compiler (available = [ ";
+ << "; ensure it is linked into the compiler (available = [ ";
for (const auto &targetName : targetRegistry.getRegisteredTargetDevices()) {
diagnostic << "'" << targetName << "' ";
}
@@ -65,7 +65,7 @@
auto diagnostic = deviceOp->emitError();
diagnostic << "unregistered target backend "
<< executableTargetAttr.getBackend()
- << "; ensure it is linked in to the compiler (available = [ ";
+ << "; ensure it is linked into the compiler (available = [ ";
for (const auto &targetName :
targetRegistry.getRegisteredTargetBackends()) {
diagnostic << "'" << targetName << "' ";
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/BUILD.bazel
index 3d1d096..812e2f9 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/BUILD.bazel
@@ -32,6 +32,7 @@
"preprocess_executables.mlir",
"prune_executables.mlir",
"repeat_dispatches.mlir",
+ "resolve_device_aliases.mlir",
"resolve_device_promises.mlir",
"resolve_export_ordinals.mlir",
"strip_executable_contents.mlir",
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/CMakeLists.txt
index 972947e..0fa3fa2 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/CMakeLists.txt
@@ -30,6 +30,7 @@
"preprocess_executables.mlir"
"prune_executables.mlir"
"repeat_dispatches.mlir"
+ "resolve_device_aliases.mlir"
"resolve_device_promises.mlir"
"resolve_export_ordinals.mlir"
"strip_executable_contents.mlir"
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/resolve_device_aliases.mlir b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/resolve_device_aliases.mlir
new file mode 100644
index 0000000..82a45cc
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/resolve_device_aliases.mlir
@@ -0,0 +1,41 @@
+// RUN: iree-opt --split-input-file --iree-hal-resolve-device-aliases %s --mlir-print-local-scope --verify-diagnostics | FileCheck %s
+
+// CHECK: util.global private @device
+// CHECK-SAME: #hal.device.target<"local"
+// CHECK-SAME: extra_config = 4 : index
+// CHECK-SAME: #hal.executable.target<"vmvx"
+util.global private @device = #hal.device.alias<"vmvx", {
+ extra_config = 4 : index
+}> : !hal.device
+
+// -----
+
+// CHECK: util.global private @device_ordinal
+// CHECK-SAME: #hal.device.target<"local"
+// CHECK-SAME: ordinal = 123 : index
+// CHECK-SAME: #hal.executable.target<"vmvx"
+util.global private @device_ordinal = #hal.device.alias<"vmvx"[123]> : !hal.device
+
+// -----
+
+// CHECK: util.global private @device_select
+// CHECK-SAME: #hal.device.select<[
+// CHECK-SAME: #hal.device.target<"local", {ordinal = 0 : index}
+// CHECK-SAME: #hal.device.target<"local", {ordinal = 1 : index}
+util.global private @device_select = #hal.device.select<[
+ #hal.device.alias<"vmvx"[0]> : !hal.device,
+ #hal.device.alias<"vmvx"[1]> : !hal.device
+]> : !hal.device
+
+// -----
+
+// expected-error@+1 {{unregistered device alias "__unregistered__"}}
+util.global private @device_unregistered = #hal.device.alias<"__unregistered__"> : !hal.device
+
+// -----
+
+// expected-error@+1 {{unregistered device alias "__unregistered__"}}
+util.global private @device_select_unregistered = #hal.device.select<[
+ #hal.device.alias<"vmvx"> : !hal.device,
+ #hal.device.alias<"__unregistered__"> : !hal.device
+]> : !hal.device
