[all-commits] [llvm/llvm-project] 96e040: [mlir][ArmSVE] Add `-arm-sve-legalize-vector-stora...

[Benjamin Maxwell via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 96e040acee7c1728506ec49a5a229bfecd49f7db
      https://github.com/llvm/llvm-project/commit/96e040acee7c1728506ec49a5a229bfecd49f7db
  Author: Benjamin Maxwell <benjamin.maxwell at arm.com>
  Date:   2023-10-26 (Thu, 26 Oct 2023)

  Changed paths:
    M mlir/include/mlir/Dialect/ArmSVE/CMakeLists.txt
    A mlir/include/mlir/Dialect/ArmSVE/Transforms/CMakeLists.txt
    A mlir/include/mlir/Dialect/ArmSVE/Transforms/Passes.h
    A mlir/include/mlir/Dialect/ArmSVE/Transforms/Passes.td
    M mlir/include/mlir/InitAllPasses.h
    M mlir/lib/Dialect/ArmSVE/Transforms/CMakeLists.txt
    A mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp
    A mlir/test/Dialect/ArmSVE/legalize-vector-storage.mlir
    A mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/arrays-of-scalable-vectors.mlir

  Log Message:
  -----------
  [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)