[mlir][ArmSME] Implement tile allocation
authorCullen Rhodes <cullen.rhodes@arm.com>
Tue, 18 Jul 2023 08:29:53 +0000 (08:29 +0000)
committerCullen Rhodes <cullen.rhodes@arm.com>
Tue, 18 Jul 2023 08:46:40 +0000 (08:46 +0000)
commitfb54fec7263ab420e9e02e1ea2b66403cf655a33
tree2f4d965725c47da93b766a6a6f02c5ea1d2c3805
parent0db5d8e12384d780d2ded6b972233dc8b1632560
[mlir][ArmSME] Implement tile allocation

This patch adds a pass '-allocate-sme-tiles' to the ArmSME dialect that
implements allocation of SME ZA tiles.

It does this at the 'func.func' op level by replacing
'arm_sme.get_tile_id' ops with 'arith.constant' ops that represent the
tile number. The tiles in use in a given function are tracked by an
integer function attribute 'arm_sme.tiles_in_use' that is a 16-bit tile
mask with a bit for each 128-bit element tile (ZA0.Q-ZA15.Q), the
smallest ZA tile granule. This is initialized on the first
'arm_sme.get_tile_id' rewrite and updated on each subsequent rewrite.
Mixing of different element tile types is supported.

Section B2.3.2 of the SME spec [1] describes how the 128-bit element
tiles overlap with other element tiles.

Depends on D154941

[1] https://developer.arm.com/documentation/ddi0616/aa

Reviewed By: awarzynski

Differential Revision: https://reviews.llvm.org/D154955
mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.h
mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td
mlir/lib/Dialect/ArmSME/Transforms/CMakeLists.txt
mlir/lib/Dialect/ArmSME/Transforms/LegalizeForLLVMExport.cpp
mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp [new file with mode: 0644]
mlir/test/Dialect/ArmSME/tile-allocation.mlir [new file with mode: 0644]
mlir/test/Integration/Dialect/Vector/CPU/ArmSME/vector-ops.mlir