-
Notifications
You must be signed in to change notification settings - Fork 3
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[compiler] Integrate xDSL kernel compilation into IREE
This PR contains the initial version integrating xDSL with IREE. The steps performed are the following: * Split a kernel function into two: - A pure MLIR dialects module that will be compiled by xDSL - A trampoline containing HAL operations that calls the latter * Compile the pure MLIR functions with xDSL to an object file * Compile the trampoline using IREE's lowerings and LLVM to an object file * Combine all object files into a single object file * Output a static library the same way the LLVMCPU target does This introduces two new build dependencies: * xdsl-opt for kernel compilation * The pulp-toolchain to assemble the output from xDSL
- Loading branch information
Showing
19 changed files
with
734 additions
and
168 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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,28 @@ | ||
| #include "Passes.h" | ||
|
|
||
| #include "mlir/Dialect/MemRef/IR/MemRef.h" | ||
|
|
||
| namespace quidditch { | ||
| #define GEN_PASS_DEF_FILTERFORXDSLPASS | ||
| #include "Quidditch/Passes.h.inc" | ||
| } // namespace quidditch | ||
|
|
||
| using namespace mlir; | ||
|
|
||
| namespace { | ||
| class FilterForxDSL | ||
| : public quidditch::impl::FilterForxDSLPassBase<FilterForxDSL> { | ||
| public: | ||
| using Base::Base; | ||
|
|
||
| protected: | ||
| void runOnOperation() override; | ||
| }; | ||
| } // namespace | ||
|
|
||
| void FilterForxDSL::runOnOperation() { | ||
| getOperation()->walk([&](Operation* op) { | ||
| if (isa<memref::AssumeAlignmentOp>(op)) | ||
| op->erase(); | ||
| }); | ||
| } |
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,98 @@ | ||
| #include "Passes.h" | ||
|
|
||
| #include <iree/compiler/Dialect/HAL/IR/HALDialect.h> | ||
| #include <mlir/Analysis/SliceAnalysis.h> | ||
| #include <mlir/Dialect/Func/IR/FuncOps.h> | ||
| #include <mlir/IR/IRMapping.h> | ||
|
|
||
| namespace quidditch { | ||
| #define GEN_PASS_DEF_HOISTHALOPSTOFUNCPASS | ||
| #include "Quidditch/Passes.h.inc" | ||
| } // namespace quidditch | ||
|
|
||
| using namespace mlir; | ||
|
|
||
| namespace { | ||
| class HoistHALOpsToFunc | ||
| : public quidditch::impl::HoistHALOpsToFuncPassBase<HoistHALOpsToFunc> { | ||
| public: | ||
| using Base::Base; | ||
|
|
||
| protected: | ||
| void runOnOperation() override; | ||
| }; | ||
| } // namespace | ||
|
|
||
| void HoistHALOpsToFunc::runOnOperation() { | ||
| OpBuilder builder(&getContext()); | ||
| builder.setInsertionPointToEnd(getOperation().getBody()); | ||
|
|
||
| Dialect *dialect = | ||
| getContext() | ||
| .getLoadedDialect<mlir::iree_compiler::IREE::HAL::HALDialect>(); | ||
| if (!dialect) | ||
| return; | ||
|
|
||
| FunctionType emptyFunctionType = builder.getFunctionType({}, {}); | ||
| for (auto func : llvm::to_vector(getOperation().getOps<func::FuncOp>())) { | ||
| // IREE functions all have empty function types with public visibility. Skip | ||
| // over any other functions. | ||
| if (func.getFunctionType() != emptyFunctionType || !func.isPublic()) | ||
| continue; | ||
|
|
||
| func->setAttr("xdsl_generated", builder.getUnitAttr()); | ||
|
|
||
| // Find all HAL operations that need to be hoisted and any other operations | ||
| // they depend on. | ||
| SmallVector<Operation *> halOperations; | ||
| SetVector<Operation *> toClone; | ||
| func.getFunctionBody().walk([&](Operation *operation) { | ||
| if (operation->getDialect() != dialect) | ||
| return; | ||
|
|
||
| halOperations.push_back(operation); | ||
| for (Value result : operation->getResults()) { | ||
| unsigned int index = func.getNumArguments(); | ||
| func.insertArgument(index, result.getType(), | ||
| builder.getDictionaryAttr({}), result.getLoc()); | ||
| result.replaceAllUsesWith(func.getArgument(index)); | ||
| } | ||
|
|
||
| // Include all operations that the HAL operation transitively depends on. | ||
| BackwardSliceOptions options; | ||
| options.inclusive = true; | ||
| // This inserts into the SetVector in topological order. As we are going | ||
| // through HAL operations in-order, it is guaranteed that any operation | ||
| // already contained in the set appears prior to the HAL operation. | ||
| mlir::getBackwardSlice(operation, &toClone, options); | ||
| }); | ||
|
|
||
| auto wrapper = builder.create<func::FuncOp>( | ||
| builder.getUnknownLoc(), (func.getName() + "$iree_to_xdsl").str(), | ||
| emptyFunctionType); | ||
|
|
||
| OpBuilder::InsertionGuard guard{builder}; | ||
| builder.setInsertionPointToStart(wrapper.addEntryBlock()); | ||
| // Create the clones and persist mapping between clones. This makes sure | ||
| // all operands are remapped. | ||
| IRMapping mapping; | ||
| for (Operation *op : toClone) { | ||
| if (!mlir::isPure(op)) { | ||
| op->emitError("Pass does not expect HAL operations to depend on " | ||
| "impure operations"); | ||
| signalPassFailure(); | ||
| return; | ||
| } | ||
|
|
||
| builder.insert(op->clone(mapping)); | ||
| } | ||
|
|
||
| SmallVector<Value> arguments; | ||
| for (Operation *op : halOperations) | ||
| for (Value value : op->getResults()) | ||
| arguments.push_back(mapping.lookup(value)); | ||
|
|
||
| builder.create<func::CallOp>(wrapper.getLoc(), func, arguments); | ||
| builder.create<func::ReturnOp>(wrapper.getLoc()); | ||
| } | ||
| } |
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,10 @@ | ||
|
|
||
| #pragma once | ||
|
|
||
| #include <mlir/IR/BuiltinOps.h> | ||
| #include <mlir/Pass/Pass.h> | ||
|
|
||
| namespace quidditch { | ||
| #define GEN_PASS_DECL | ||
| #include "Quidditch/Passes.h.inc" | ||
| } // namespace quidditch |
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,31 @@ | ||
| #ifndef QUIDDITCH_PASSES | ||
| #define QUIDDITCH_PASSES | ||
|
|
||
| include "mlir/Pass/PassBase.td" | ||
|
|
||
| def HoistHALOpsToFuncPass | ||
| : Pass<"quidditch-hoist-hal-ops-to-func", "mlir::ModuleOp"> { | ||
| let description = [{ | ||
| Pass performing HAL operation outlining. | ||
| For every function in the module this creates another function called "<old-name>$iree_to_xDSL" and moves all HAL | ||
| operations to that function before calling the original function. | ||
|
|
||
| This allows to pass the original function to xDSL for kernel compilation while compiling the new generated function | ||
| can be lowered by IREE. | ||
| This pass can be removed if xDSL were to learn how to lower IREE HAL operations to RISC-V. | ||
|
|
||
| The original function is additionally tagged with a "xdsl_generated" unit attribute. | ||
| }]; | ||
| } | ||
|
|
||
| def FilterForxDSLPass | ||
| : Pass<"quidditch-filter-for-xdsl"> { | ||
| let description = [{ | ||
| Pass rewriting operations to workaround xDSL operations. | ||
|
|
||
| Right now these are: | ||
| * Remove `memref.assume_alignment`. | ||
| }]; | ||
| } | ||
|
|
||
| #endif |
Oops, something went wrong.