[all-commits] [llvm/llvm-project] eaff02: [mlir][ArmSME] Switch to an attribute-based tile a...

Thu Nov 30 02:22:37 PST 2023

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: eaff02f28e194b497ec2d9a6bbecc33b54f6df27
      https://github.com/llvm/llvm-project/commit/eaff02f28e194b497ec2d9a6bbecc33b54f6df27
  Author: Benjamin Maxwell <benjamin.maxwell at arm.com>
  Date:   2023-11-30 (Thu, 30 Nov 2023)

  Changed paths:
    M mlir/include/mlir/Dialect/ArmSME/IR/ArmSME.h
    A mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEEnums.h
    M mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEIntrinsicOps.td
    M mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEOps.td
    M mlir/include/mlir/Dialect/ArmSME/IR/CMakeLists.txt
    M mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.h
    M mlir/include/mlir/Dialect/ArmSME/Utils/Utils.h
    M mlir/lib/Conversion/ArmSMEToLLVM/ArmSMEToLLVM.cpp
    M mlir/lib/Conversion/ArmSMEToLLVM/CMakeLists.txt
    M mlir/lib/Conversion/ArmSMEToSCF/ArmSMEToSCF.cpp
    M mlir/lib/Conversion/ArmSMEToSCF/CMakeLists.txt
    M mlir/lib/Conversion/VectorToArmSME/CMakeLists.txt
    M mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp
    M mlir/lib/Dialect/ArmSME/CMakeLists.txt
    M mlir/lib/Dialect/ArmSME/IR/ArmSME.cpp
    M mlir/lib/Dialect/ArmSME/IR/CMakeLists.txt
    A mlir/lib/Dialect/ArmSME/IR/Utils.cpp
    M mlir/lib/Dialect/ArmSME/Transforms/CMakeLists.txt
    M mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
    R mlir/lib/Dialect/ArmSME/Utils/CMakeLists.txt
    R mlir/lib/Dialect/ArmSME/Utils/Utils.cpp
    M mlir/test/Conversion/ArmSMEToSCF/arm-sme-to-scf.mlir
    M mlir/test/Dialect/ArmSME/arith-ops-to-sme.mlir
    R mlir/test/Dialect/ArmSME/arm-sme-to-llvm-casts.mlir
    M mlir/test/Dialect/ArmSME/arm-sme-to-llvm.mlir
    M mlir/test/Dialect/ArmSME/canonicalize.mlir
    M mlir/test/Dialect/ArmSME/cse.mlir
    M mlir/test/Dialect/ArmSME/invalid.mlir
    M mlir/test/Dialect/ArmSME/roundtrip.mlir
    A mlir/test/Dialect/ArmSME/tile-allocation-invalid.mlir
    M mlir/test/Dialect/ArmSME/tile-allocation.mlir
    M mlir/test/Dialect/ArmSME/tile-zero-masks.mlir
    M mlir/test/Dialect/ArmSME/vector-ops-to-llvm.mlir
    M mlir/test/Dialect/ArmSME/vector-ops-to-sme.mlir
    M mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/fill-2d.mlir
    M mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul-transpose-a.mlir
    M mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul.mlir
    M mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-load-vertical.mlir
    M mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-outerproduct-f32.mlir
    M mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-outerproduct-f64.mlir
    M mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-transfer-read-2d.mlir
    M mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-transfer-write-2d.mlir
    M mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-transpose.mlir
    M mlir/test/Integration/Dialect/Vector/CPU/ArmSME/tile_fill.mlir
    M mlir/test/Integration/Dialect/Vector/CPU/ArmSME/vector-load-store.mlir
    M mlir/test/Integration/Dialect/Vector/CPU/ArmSME/vector-ops.mlir
    M mlir/test/Target/LLVMIR/arm-sme-invalid.mlir
    M mlir/test/Target/LLVMIR/arm-sme.mlir

  Log Message:
  -----------
  [mlir][ArmSME] Switch to an attribute-based tile allocation scheme (#73253)

This reworks the ArmSME dialect to use attributes for tile allocation.
This has a number of advantages and corrects some issues with the
previous approach:

* Tile allocation can now be done ASAP (i.e. immediately after
`-convert-vector-to-arm-sme`)
* SSA form for control flow is now supported (e.g.`scf.for` loops that
yield tiles)
* ArmSME ops can be converted to intrinsics very late (i.e. after
lowering to control flow)
 * Tests are simplified by removing constants and casts
* Avoids correctness issues with representing LLVM `immargs` as MLIR
values
- The tile ID on the SME intrinsics is an `immarg` (so is required to be
a compile-time constant), `immargs` should be mapped to MLIR attributes
(this is already the case for intrinsics in the LLVM dialect)
- Using MLIR values for `immargs` can lead to invalid LLVM IR being
generated (and passes such as -cse making incorrect optimizations)

As part of this patch we bid farewell to the following operations:

```mlir
arm_sme.get_tile_id : i32
arm_sme.cast_tile_to_vector : i32 to vector<[4]x[4]xi32>
arm_sme.cast_vector_to_tile : vector<[4]x[4]xi32> to i32
```

These are now replaced with:
```mlir
// Allocates a new tile with (indeterminate) state:
arm_sme.get_tile : vector<[4]x[4]xi32>
// A placeholder operation for lowering ArmSME ops to intrinsics:
arm_sme.materialize_ssa_tile : vector<[4]x[4]xi32>
```

The new tile allocation works by operations implementing the
`ArmSMETileOpInterface`. This interface says that an operation needs to
be assigned a tile ID, and may conditionally allocate a new SME tile.

Operations allocate a new tile by implementing...
```c++
std::optional<arm_sme::ArmSMETileType> getAllocatedTileType()
```
...and returning what type of tile the op allocates (ZAB, ZAH, etc).

Operations that don't allocate a tile return `std::nullopt` (which is
the default behaviour).

Currently the following ops are defined as allocating:
```mlir
arm_sme.get_tile
arm_sme.zero
arm_sme.tile_load
arm_sme.outerproduct // (if no accumulator is specified)
```

Allocating operations become the roots for the tile allocation pass,
which currently just (naively) assigns all transitive uses of a root
operation the same tile ID. However, this is enough to handle current
use cases.

Once tile IDs have been allocated subsequent rewrites can forward the
tile IDs to any newly created operations.