[mlir][ArmSME] Extend streaming-mode pass to support enabling ZA
authorCullen Rhodes <cullen.rhodes@arm.com>
Fri, 16 Jun 2023 09:14:53 +0000 (09:14 +0000)
committerCullen Rhodes <cullen.rhodes@arm.com>
Fri, 16 Jun 2023 09:26:42 +0000 (09:26 +0000)
commite947e760585c2a457775d8723e1503f7d5b56984
tree186ca7038d26a8ce91845f2bf68fa381fe7b063d
parent7a12fbc05f4f53a49213adc20cbab34425cfbfa7
[mlir][ArmSME] Extend streaming-mode pass to support enabling ZA

This patch extends the 'enable-arm-streaming' pass with a new option to
enable the ZA storage array by adding the 'arm_za' attribute to
'func.func' ops.

A later patch will insert `llvm.aarch64.sme.za.enable` at the beginning
of 'func.func' ops and `llvm.aarch64.sme.za.disable` before
`func.return` statements when lowering to LLVM dialect.

Currently the pass only supports enabling ZA with streaming-mode on but
the SME LDR, STR and ZERO instructions can access ZA when not in
streaming-mode (section B1.1.1, IDGNQM [1]), so it may be worth making
these options independent in the future.

N.B. This patch is generally useful in the context of SME enablement in
MLIR, but it will help enable writing an integration test for rewrite
pattern that lowers `vector.transfer_write` -> `zero {za}` (D152508).

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

Reviewed By: awarzynski, dcaballe

Differential Revision: https://reviews.llvm.org/D152695
mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.h
mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td
mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir