-
Notifications
You must be signed in to change notification settings - Fork 653
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
New AssignTargetDevices pass to replace the legacy one.
The legacy pass has been moved aside so that the old flags still work but will be removed in the future.
- Loading branch information
Showing
52 changed files
with
903 additions
and
249 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
114 changes: 114 additions & 0 deletions
114
compiler/src/iree/compiler/Dialect/HAL/Transforms/AssignLegacyTargetDevices.cpp
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,114 @@ | ||
| // Copyright 2021 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/TargetBackend.h" | ||
| #include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h" | ||
| #include "iree/compiler/Dialect/HAL/Transforms/Passes.h" | ||
| #include "llvm/ADT/STLExtras.h" | ||
| #include "llvm/Support/raw_ostream.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_ASSIGNLEGACYTARGETDEVICESPASS | ||
| #include "iree/compiler/Dialect/HAL/Transforms/Passes.h.inc" | ||
|
|
||
| namespace { | ||
|
|
||
| //===----------------------------------------------------------------------===// | ||
| // --iree-hal-assign-legacy-target-devices | ||
| //===----------------------------------------------------------------------===// | ||
|
|
||
| struct AssignLegacyTargetDevicesPass | ||
| : public IREE::HAL::impl::AssignLegacyTargetDevicesPassBase< | ||
| AssignLegacyTargetDevicesPass> { | ||
| using IREE::HAL::impl::AssignLegacyTargetDevicesPassBase< | ||
| AssignLegacyTargetDevicesPass>::AssignLegacyTargetDevicesPassBase; | ||
|
|
||
| void runOnOperation() override { | ||
| auto moduleOp = getOperation(); | ||
|
|
||
| // If no targets are specified we can't do anything - another pass earlier | ||
| // in the pipeline will have had to add the targets. | ||
| if (targetBackends.empty()) | ||
| return; | ||
|
|
||
| // Check to see if targets are already specified and if so then no-op the | ||
| // pass so that we don't mess with whatever the user intended. | ||
| auto existingTargetsAttr = | ||
| moduleOp->getAttrOfType<ArrayAttr>("hal.device.targets"); | ||
| if (existingTargetsAttr) | ||
| return; | ||
|
|
||
| // If there are any device globals declared then bail as it means the user | ||
| // has already materialized the devices they want. | ||
| for (auto globalOp : moduleOp.getOps<IREE::Util::GlobalOpInterface>()) { | ||
| if (isa<IREE::HAL::DeviceType>(globalOp.getGlobalType())) | ||
| return; | ||
| } | ||
|
|
||
| llvm::SmallDenseSet<Attribute> targetAttrSet; | ||
| SmallVector<Attribute> targetAttrs; | ||
| for (const auto &targetBackendName : targetBackends) { | ||
| auto targetBackend = targetRegistry->getTargetBackend(targetBackendName); | ||
| if (!targetBackend) { | ||
| auto diagnostic = emitError(moduleOp.getLoc()) | ||
| << "target backend '" << targetBackendName | ||
| << "' not registered; registered backends: ["; | ||
| llvm::interleaveComma(targetRegistry->getRegisteredTargetBackends(), | ||
| diagnostic); | ||
| diagnostic << "]"; | ||
| return signalPassFailure(); | ||
| } | ||
| auto targetDeviceName = targetBackend->getLegacyDefaultDeviceID(); | ||
| auto targetDevice = targetRegistry->getTargetDevice(targetDeviceName); | ||
| if (!targetDevice) { | ||
| auto diagnostic = emitError(moduleOp.getLoc()) | ||
| << "target device '" << targetDeviceName | ||
| << "' not registered; registered devices: ["; | ||
| llvm::interleaveComma(targetRegistry->getRegisteredTargetDevices(), | ||
| diagnostic); | ||
| diagnostic << "]"; | ||
| return signalPassFailure(); | ||
| } | ||
|
|
||
| // Ask the target backend for its default device specification attribute. | ||
| auto targetAttr = targetDevice->getDefaultDeviceTarget( | ||
| moduleOp.getContext(), *targetRegistry.value); | ||
| if (!targetAttr) { | ||
| emitError(moduleOp.getLoc()) << "no default device targets available"; | ||
| return signalPassFailure(); | ||
| } | ||
| if (!targetAttrSet.contains(targetAttr)) { | ||
| targetAttrSet.insert(targetAttr); | ||
| targetAttrs.push_back(targetAttr); | ||
| } | ||
| } | ||
|
|
||
| Attribute targetsAttr; | ||
| if (targetAttrs.size() == 1) { | ||
| targetsAttr = targetAttrs.front(); | ||
| } else { | ||
| targetsAttr = | ||
| IREE::HAL::DeviceSelectAttr::get(moduleOp.getContext(), targetAttrs); | ||
| } | ||
| moduleOp->setAttr("hal.device.targets", | ||
| ArrayAttr::get(moduleOp.getContext(), targetsAttr)); | ||
| } | ||
| }; | ||
|
|
||
| } // namespace | ||
|
|
||
| } // namespace mlir::iree_compiler::IREE::HAL |
Oops, something went wrong.