iree-org / iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.
http://iree.dev/
Apache License 2.0
2.47k stars 548 forks source link

[LLVMCPU][ArmSME] Rework how Arm streaming mode is set on dispatches #17646

Closed MacDue closed 1 day ago

MacDue commented 2 weeks ago

Previously, when the +sme feature flag was set Armv9 streaming SVE mode would be enabled for all dispatch regions lowered with the following experts:

This was not ideal as meant streaming mode could be added to dispatch regions that made no use of scalable vectors, where the (possibly) larger streaming vector length provides no benefit, and there may be a cost due to other overheads.

There was also a flag --iree-experimental-llvmcpu-arm-force-ssve which contrary to its name did not force streaming SVE mode. What this flag did do was disable tiling for 2D scalable ArmSME operations, then rely on something else later on setting the streaming mode (but it did not control it).

The patch aims to add clearer and more directed ways to enable streaming mode.

First, streaming mode is no longer set in any lowering experts (it's a fairly low-level concept, that does not need to be configured early in the pipeline). Second, the old --iree-experimental-llvmcpu-arm-force-ssve flag is removed.

Now to control tiling for ArmSME and using streaming mode there are two new flags.

iree-llvmcpu-disable-arm-sme-tiling:

This disables tiling for ArmSME (i.e. using 2D scalable tile sizes), even when the +sme feature flag is set. This results in operations instead being tiled for SVE or Neon (depending on the configuration).

iree-llvmcpu-force-arm-streaming:

This enables Arm streaming mode for any dispatch regions that contain scalable vectors. It ignores dispatches that don't contain scalable vectors as enabling streaming mode would provide no benefit.

ci-extra: build_test_all_arm64

MacDue commented 2 weeks ago

Depends on https://github.com/llvm/llvm-project/pull/94759

MacDue commented 4 days ago

You may not always want to use the ArmSME lowerings (2D scalable outer products, etc) when the +sme feature flag is present, so the --iree-llvmcpu-disable-arm-sme-tiling flag gives the option to disable tiling for them. Using the ArmSME lowerings is the default as it's what you'd want in most cases.

SSVE (that is SVE running in streaming mode) is something that needs some consideration for when it should be used. Currently, there's not a proper cost model that can decide when it is beneficial or not. So to allow testing it in the meantime the --iree-llvmcpu-force-arm-streaming enables it for any scalable dispatch regions that support it. This is more experimental so it is not enabled by default.

Combination-wise (assuming +sme,+sve with scalable vectorization enabled), it works out like this:

--iree-llvmcpu-disable-arm-sme-tiling no flag (default)
--iree-llvmcpu-force-arm-streaming Streaming SVE SME (force-arm-streaming is redundant in this case)
no flag (default) SVE SME
MaheshRavishankar commented 3 days ago

You may not always want to use the ArmSME lowerings (2D scalable outer products, etc) when the +sme feature flag is present, so the --iree-llvmcpu-disable-arm-sme-tiling flag gives the option to disable tiling for them. Using the ArmSME lowerings is the default as it's what you'd want in most cases.

SSVE (that is SVE running in streaming mode) is something that needs some consideration for when it should be used. Currently, there's not a proper cost model that can decide when it is beneficial or not. So to allow testing it in the meantime the --iree-llvmcpu-force-arm-streaming enables it for any scalable dispatch regions that support it. This is more experimental so it is not enabled by default.

Combination-wise (assuming +sme,+sve with scalable vectorization enabled), it works out like this:

--iree-llvmcpu-disable-arm-sme-tiling no flag (default) --iree-llvmcpu-force-arm-streaming Streaming SVE SME (force-arm-streaming is redundant in this case) no flag (default) SVE SME

Thanks @MacDue for the info. Shouldnt the no flag x no flag (the quadrant on the bottom right) be None, i.e. neither use SVE or SME. We dont want to accidently trigger SVE/SME because the /proc/cpuinfo has that. Or is this an explicit choice that if a hardware has SME IREEs CPU codegen will try to use it. That seems at odds with the level of support there is right now.

MacDue commented 3 days ago

None, i.e. neither use SVE or SME. We dont want to accidently trigger SVE/SME because the /proc/cpuinfo has that. Or is this an explicit choice that if a hardware has SME IREEs CPU codegen will try to use it.

It is an explicit choice that it will try to use SME/SVE if scalable vectorization is enabled (which that diagram assumes). Scalable vectorization is not on by default right now, so the user still has to opt-in for SME/SVE codegen.

(note the previous post says "assuming +sme,+sve with scalable vectorization enabled")

MacDue commented 2 days ago

@MaheshRavishankar We would really like to land this soon :)

MacDue commented 2 days ago

Could you land the PR? I don't have write access -- thanks!