-
Notifications
You must be signed in to change notification settings - Fork 10.8k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[mlir][ArmSVE] Add
-arm-sve-legalize-vector-storage
pass (#68794)
This patch adds a pass that ensures that loads, stores, and allocations of SVE vector types will be legal in the LLVM backend. It does this at the memref level, so this pass must be applied before lowering all the way to LLVM. This pass currently fixes two issues. ## Loading and storing predicate types It is only legal to load/store predicate types equal to (or greater than) a full predicate register, which in MLIR is `vector<[16]xi1>`. Smaller predicate types (`vector<[1|2|4|8]xi1>`) must be converted to/from a full predicate type (referred to as a `svbool`) before and after storing and loading respectively. This pass does this by widening allocations and inserting conversion intrinsics. For example: ```mlir %alloca = memref.alloca() : memref<vector<[4]xi1>> %mask = vector.constant_mask [4] : vector<[4]xi1> memref.store %mask, %alloca[] : memref<vector<[4]xi1>> %reload = memref.load %alloca[] : memref<vector<[4]xi1>> ``` Becomes: ```mlir %alloca = memref.alloca() {alignment = 1 : i64} : memref<vector<[16]xi1>> %mask = vector.constant_mask [4] : vector<[4]xi1> %svbool = arm_sve.convert_to_svbool %mask : vector<[4]xi1> memref.store %svbool, %alloca[] : memref<vector<[16]xi1>> %reload_svbool = memref.load %alloca[] : memref<vector<[16]xi1>> %reload = arm_sve.convert_from_svbool %reload_svbool : vector<[4]xi1> ``` ## Relax alignments for SVE vector allocas The storage for SVE vector types only needs to have an alignment that matches the element type (for example 4 byte alignment for `f32`s). However, the LLVM backend currently defaults to aligning to `base size x element size` bytes. For non-legal vector types like `vector<[8]xf32>` this results in 8 x 4 = 32-byte alignment, but the backend only supports up to 16-byte alignment for SVE vectors on the stack. Explicitly setting a smaller alignment prevents this issue. Depends on: #68586 and #68695 (for testing)
- Loading branch information
Showing
9 changed files
with
772 additions
and
0 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1 +1,2 @@ | ||
add_subdirectory(IR) | ||
add_subdirectory(Transforms) |
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,5 @@ | ||
set(LLVM_TARGET_DEFINITIONS Passes.td) | ||
mlir_tablegen(Passes.h.inc -gen-pass-decls -name ArmSVE) | ||
add_public_tablegen_target(MLIRArmSVEPassIncGen) | ||
|
||
add_mlir_doc(Passes ArmSVEPasses ./ -gen-pass-doc) |
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,36 @@ | ||
//===- Passes.h - Pass Entrypoints ------------------------------*- C++ -*-===// | ||
// | ||
// Part of the LLVM Project, 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 | ||
// | ||
//===----------------------------------------------------------------------===// | ||
|
||
#ifndef MLIR_DIALECT_ARMSVE_TRANSFORMS_PASSES_H | ||
#define MLIR_DIALECT_ARMSVE_TRANSFORMS_PASSES_H | ||
|
||
#include "mlir/Conversion/LLVMCommon/TypeConverter.h" | ||
#include "mlir/Pass/Pass.h" | ||
|
||
namespace mlir::arm_sve { | ||
|
||
#define GEN_PASS_DECL | ||
#include "mlir/Dialect/ArmSVE/Transforms/Passes.h.inc" | ||
|
||
/// Pass to legalize Arm SVE vector storage. | ||
std::unique_ptr<Pass> createLegalizeVectorStoragePass(); | ||
|
||
/// Collect a set of patterns to legalize Arm SVE vector storage. | ||
void populateLegalizeVectorStoragePatterns(RewritePatternSet &patterns); | ||
|
||
//===----------------------------------------------------------------------===// | ||
// Registration | ||
//===----------------------------------------------------------------------===// | ||
|
||
/// Generate the code for registering passes. | ||
#define GEN_PASS_REGISTRATION | ||
#include "mlir/Dialect/ArmSVE/Transforms/Passes.h.inc" | ||
|
||
} // namespace mlir::arm_sve | ||
|
||
#endif // MLIR_DIALECT_ARMSVE_TRANSFORMS_PASSES_H |
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,68 @@ | ||
//===-- Passes.td - ArmSVE pass definition file ------------*- tablegen -*-===// | ||
// | ||
// Part of the LLVM Project, 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 | ||
// | ||
//===----------------------------------------------------------------------===// | ||
|
||
#ifndef MLIR_DIALECT_ARMSVE_TRANSFORMS_PASSES_TD | ||
#define MLIR_DIALECT_ARMSVE_TRANSFORMS_PASSES_TD | ||
|
||
include "mlir/Pass/PassBase.td" | ||
|
||
def LegalizeVectorStorage | ||
: Pass<"arm-sve-legalize-vector-storage", "mlir::func::FuncOp"> { | ||
let summary = "Ensures stores of SVE vector types will be legal"; | ||
let description = [{ | ||
This pass ensures that loads, stores, and allocations of SVE vector types | ||
will be legal in the LLVM backend. It does this at the memref level, so this | ||
pass must be applied before lowering all the way to LLVM. | ||
|
||
This pass currently addresses two issues. | ||
|
||
## Loading and storing predicate types | ||
|
||
It is only legal to load/store predicate types equal to (or greater than) a | ||
full predicate register, which in MLIR is `vector<[16]xi1>`. Smaller | ||
predicate types (`vector<[1|2|4|8]xi1>`) must be converted to/from a full | ||
predicate type (referred to as a `svbool`) before and after storing and | ||
loading respectively. This pass does this by widening allocations and | ||
inserting conversion intrinsics. Note: Non-powers-of-two masks (e.g. | ||
`vector<[7]xi1>`), which are not SVE predicates, are ignored. | ||
|
||
For example: | ||
|
||
```mlir | ||
%alloca = memref.alloca() : memref<vector<[4]xi1>> | ||
%mask = vector.constant_mask [4] : vector<[4]xi1> | ||
memref.store %mask, %alloca[] : memref<vector<[4]xi1>> | ||
%reload = memref.load %alloca[] : memref<vector<[4]xi1>> | ||
``` | ||
Becomes: | ||
```mlir | ||
%alloca = memref.alloca() {alignment = 1 : i64} : memref<vector<[16]xi1>> | ||
%mask = vector.constant_mask [4] : vector<[4]xi1> | ||
%svbool = arm_sve.convert_to_svbool %mask : vector<[4]xi1> | ||
memref.store %svbool, %alloca[] : memref<vector<[16]xi1>> | ||
%reload_svbool = memref.load %alloca[] : memref<vector<[16]xi1>> | ||
%reload = arm_sve.convert_from_svbool %reload_svbool : vector<[4]xi1> | ||
``` | ||
|
||
## Relax alignments for SVE vector allocas | ||
|
||
The storage for SVE vector types only needs to have an alignment that | ||
matches the element type (for example 4 byte alignment for `f32`s). However, | ||
the LLVM backend currently defaults to aligning to `base size` x | ||
`element size` bytes. For non-legal vector types like `vector<[8]xf32>` this | ||
results in 8 x 4 = 32-byte alignment, but the backend only supports up to | ||
16-byte alignment for SVE vectors on the stack. Explicitly setting a smaller | ||
alignment prevents this issue. | ||
}]; | ||
let constructor = "mlir::arm_sve::createLegalizeVectorStoragePass()"; | ||
let dependentDialects = ["func::FuncDialect", | ||
"memref::MemRefDialect", "vector::VectorDialect", | ||
"arm_sve::ArmSVEDialect"]; | ||
} | ||
|
||
#endif // MLIR_DIALECT_ARMSVE_TRANSFORMS_PASSES_TD |
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
Oops, something went wrong.