-
Notifications
You must be signed in to change notification settings - Fork 618
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Improving linking support for ROCM and ukernels. (#19211)
To support externally-defined ukernels on ROCM the ROCMTarget has been brought in-line with LLVMCPU/CUDA by calling `linkBitcodeObjects`. To make authoring passes that include object references `#hal.executable.object` now allows any data type to be associated so long as it is serializable allowing for external resource attrs and other custom attributes that may serialize based on other information. To allow patterns to attach object references all ops within an executable variant can now declare a `hal.executable.objects` array that will be hoisted and merged into the top-level variant objects after our executable linking pass (before serialization where they are used).
- Loading branch information
Showing
13 changed files
with
191 additions
and
14 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
69 changes: 69 additions & 0 deletions
69
compiler/src/iree/compiler/Dialect/HAL/Transforms/HoistExecutableObjects.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,69 @@ | ||
| // 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/HAL/IR/HALOps.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_HOISTEXECUTABLEOBJECTSPASS | ||
| #include "iree/compiler/Dialect/HAL/Transforms/Passes.h.inc" | ||
|
|
||
| namespace { | ||
|
|
||
| //===----------------------------------------------------------------------===// | ||
| // --iree-hal-hoist-executable-objects | ||
| //===----------------------------------------------------------------------===// | ||
|
|
||
| struct HoistExecutableObjectsPass | ||
| : public IREE::HAL::impl::HoistExecutableObjectsPassBase< | ||
| HoistExecutableObjectsPass> { | ||
| void runOnOperation() override { | ||
| // Note that some executables may be external and not have any contents. | ||
| if (getOperation().isExternal()) { | ||
| return; | ||
| } | ||
|
|
||
| auto objectsAttrName = | ||
| StringAttr::get(&getContext(), "hal.executable.objects"); | ||
|
|
||
| // Seed with existing variant-level object attrs, if any present. | ||
| SetVector<Attribute> allObjectAttrs; | ||
| if (auto existingAttr = getOperation().getObjectsAttr()) { | ||
| allObjectAttrs.insert(existingAttr.begin(), existingAttr.end()); | ||
| } | ||
|
|
||
| // Move all op-level attributes into a unique set. Note that order can be | ||
| // important so we use an ordered set. | ||
| // | ||
| // We could do this first as a gather step in parallel if this walk gets too | ||
| // expensive. | ||
| bool foundAnyAttrs = false; | ||
| getOperation().getInnerModule().walk([&](Operation *op) { | ||
| auto objectsAttr = op->getAttrOfType<ArrayAttr>(objectsAttrName); | ||
| if (objectsAttr) { | ||
| allObjectAttrs.insert(objectsAttr.begin(), objectsAttr.end()); | ||
| op->removeAttr(objectsAttrName); | ||
| foundAnyAttrs = true; | ||
| } | ||
| }); | ||
|
|
||
| // Update the variant if any changes were made. | ||
| if (foundAnyAttrs) { | ||
| getOperation().setObjectsAttr( | ||
| ArrayAttr::get(&getContext(), allObjectAttrs.getArrayRef())); | ||
| } | ||
| } | ||
| }; | ||
|
|
||
| } // namespace | ||
|
|
||
| } // namespace mlir::iree_compiler::IREE::HAL |
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
61 changes: 61 additions & 0 deletions
61
compiler/src/iree/compiler/Dialect/HAL/Transforms/test/hoist_executable_objects.mlir
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,61 @@ | ||
| // RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(hal.executable(hal.executable.variant(iree-hal-hoist-executable-objects)))" %s | FileCheck %s | ||
|
|
||
| // Tests that attributes on top-level ops and nested ops are all detected, | ||
| // deduplicated, and moved to the variant. | ||
|
|
||
| // CHECK: hal.executable public @executable | ||
| hal.executable public @executable { | ||
| // CHECK: hal.executable.variant public @backend | ||
| // CHECK-SAME: objects([ | ||
| // CHECK-SAME: #hal.executable.object<{path = "existing_variant.obj"}>, | ||
| // CHECK-SAME: #hal.executable.object<{path = "extern_fn_common.obj"}>, | ||
| // CHECK-SAME: #hal.executable.object<{path = "extern_fn_a.obj"}>, | ||
| // CHECK-SAME: #hal.executable.object<{path = "extern_fn_b.obj"}>, | ||
| // CHECK-SAME: #hal.executable.object<{path = "nested_common.obj"}>, | ||
| // CHECK-SAME: #hal.executable.object<{path = "nested_a.obj"}>, | ||
| // CHECK-SAME: #hal.executable.object<{path = "nested_b.obj"}> | ||
| hal.executable.variant public @backend target(#hal.executable.target<"backend", "format">) objects([ | ||
| #hal.executable.object<{path = "existing_variant.obj"}> | ||
| ]) { | ||
| hal.executable.export @entry0 ordinal(0) layout(#hal.pipeline.layout<bindings = [ | ||
| #hal.pipeline.binding<storage_buffer> | ||
| ]>) | ||
| builtin.module { | ||
| // CHECK: func.func private @extern_fn_a | ||
| // CHECK-NOT: hal.executable.objects | ||
| func.func private @extern_fn_a() attributes { | ||
| hal.executable.objects = [ | ||
| #hal.executable.object<{path = "extern_fn_common.obj"}>, | ||
| #hal.executable.object<{path = "extern_fn_a.obj"}> | ||
| ] | ||
| } | ||
| // CHECK: func.func private @extern_fn_b | ||
| // CHECK-NOT: hal.executable.objects | ||
| func.func private @extern_fn_b() attributes { | ||
| hal.executable.objects = [ | ||
| #hal.executable.object<{path = "extern_fn_common.obj"}>, | ||
| #hal.executable.object<{path = "extern_fn_b.obj"}> | ||
| ] | ||
| } | ||
| func.func @entry0() { | ||
| // CHECK: call @extern_fn_a | ||
| // CHECK-NOT: hal.executable.objects | ||
| call @extern_fn_a() { | ||
| hal.executable.objects = [ | ||
| #hal.executable.object<{path = "nested_common.obj"}>, | ||
| #hal.executable.object<{path = "nested_a.obj"}> | ||
| ] | ||
| } : () -> () | ||
| call @extern_fn_b() { | ||
| // CHECK: call @extern_fn_b | ||
| // CHECK-NOT: hal.executable.objects | ||
| hal.executable.objects = [ | ||
| #hal.executable.object<{path = "nested_common.obj"}>, | ||
| #hal.executable.object<{path = "nested_b.obj"}> | ||
| ] | ||
| } : () -> () | ||
| return | ||
| } | ||
| } | ||
| } | ||
| } |
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