-
Notifications
You must be signed in to change notification settings - Fork 576
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Codegen][GPU] Add pass to unroll to native mma widths (#18101)
Left as TODO is to extend this pass to also unroll consumers/producers. To eventually connect everything we should add a pattern that unrolls iter args of `scf.for` loops. Note that the unrolling pattern is already tested in its own test, so simply updating the pipeline test is fine here. Once extending this pass to do more unrolling a test can be added.
- Loading branch information
Showing
9 changed files
with
175 additions
and
97 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
58 changes: 58 additions & 0 deletions
58
compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/UnrollToIntrinsics.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,58 @@ | ||
| // 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/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h" | ||
| #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h" | ||
| #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h" | ||
| #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUInterfaces.h" | ||
| #include "iree/compiler/Codegen/Dialect/GPU/Transforms/Passes.h" | ||
| #include "iree/compiler/Codegen/Dialect/GPU/Transforms/Transforms.h" | ||
| #include "mlir/Dialect/Linalg/IR/Linalg.h" | ||
| #include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h" | ||
| #include "mlir/Dialect/Linalg/Transforms/Transforms.h" | ||
| #include "mlir/Dialect/Tensor/IR/Tensor.h" | ||
| #include "mlir/IR/OpDefinition.h" | ||
| #include "mlir/IR/PatternMatch.h" | ||
| #include "mlir/Interfaces/FunctionInterfaces.h" | ||
| #include "mlir/Transforms/GreedyPatternRewriteDriver.h" | ||
|
|
||
| namespace mlir::iree_compiler::IREE::GPU { | ||
|
|
||
| #define GEN_PASS_DEF_UNROLLTOINTRINSICSPASS | ||
| #include "iree/compiler/Codegen/Dialect/GPU/Transforms/Passes.h.inc" | ||
|
|
||
| namespace { | ||
| struct UnrollToIntrinsicsPass final | ||
| : impl::UnrollToIntrinsicsPassBase<UnrollToIntrinsicsPass> { | ||
| void runOnOperation() override; | ||
| }; | ||
| } // namespace | ||
|
|
||
| void UnrollToIntrinsicsPass::runOnOperation() { | ||
| MLIRContext *context = &getContext(); | ||
|
|
||
| { | ||
| RewritePatternSet patterns(context); | ||
| GPU::populateIREEGPUVectorUnrollPatterns(patterns); | ||
| if (failed(applyPatternsAndFoldGreedily(getOperation(), | ||
| std::move(patterns)))) { | ||
| return signalPassFailure(); | ||
| } | ||
| } | ||
|
|
||
| // Post unrolling unit dim folding patterns in preparation for later | ||
| // lowerings. | ||
| { | ||
| RewritePatternSet patterns(context); | ||
| GPU::populateIREEGPUDropUnitDimsPatterns(patterns); | ||
| if (failed(applyPatternsAndFoldGreedily(getOperation(), | ||
| std::move(patterns)))) { | ||
| return signalPassFailure(); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| } // namespace mlir::iree_compiler::IREE::GPU |
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