Enable SME2 Streaming SVE in ARM

[stevesuzuki-arm]

Enable SME2 Streaming SVE in ARM

This PR adds initial ARM SME2 streaming-mode support to Halide,
which allows us to compute with longer vector length SVE on targets with SME2.

A new sme_streaming(enable, var) scheduling directive provides the users
the option to control which loop is computed in streaming-mode.

The change introduces a new Target::SME2 feature with supplemental features Target::SME_SVLDDD, where DDD represents streaming vector length in bits (e.g. 128, 256, 512, ...). If Target::SME2 is enabled, exactly one of Target::SME_SVLDDD feature must be enabled as well.
natural_vector_size() now depends on whether in streaming-mode or not,
because streaming vector length may have a value different from non-streaming vector length.

In Halide lowering, a new LowerSMEStreamingTasks pass is added,
which extracts the loop with streaming-mode as internal closure function
so that we can attach the LLVM function attributes to transit to/from streaming-mode.

• aarch64_pstate_sm_body to emit smstart/smstop transition
• NoInline to prevent streaming closure from inlined to non-streaming function

In CodeGen, target_vscale() depends on whether streaming-mode or not
and it varies even in a Module, although it is constant within Function boundary.
In streaming-mode, vector type code-gen and intrinsic selection are
performed based on Target::sme_streaming_vector_bits() (streaming vscale).
In terms of coverage, it is almost the same as existing SVE2 code-gen
while SME2 specific instruction has not been enabled for now.

Additionally, the following changes are implemented:

• Auto-detect SME2 and SME_SVLDDD target features on host CPU
• Fall back from streaming SVE when vectorization factors are not feasible
• Gather/scatter in streaming mode is scalarized with warning
• Add runtime checks for streaming vscale mismatches with compile-time vscale

Checklist

• Tests added or updated (not required for docs, CI config, or typo fixes)
• Documentation updated (if public API changed)
• Python bindings updated (if public API changed)

[stevesuzuki-arm]

This PR is ready for review. I will touch on this in dev meeting if I have a chance.

[stevesuzuki-arm]

Based on the feedback in dev meeting, streaming_vector_bits has been replaced with Feature::SME_SVL

[codecov]

Codecov Report

❌ Patch coverage is 62.41379% with 109 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@d58798a). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
src/CodeGen_ARM.cpp	54.25%	34 Missing and 9 partials ⚠️
src/Target.cpp	51.68%	23 Missing and 20 partials ⚠️
src/LowerSMEStreamingTasks.cpp	81.69%	9 Missing and 4 partials ⚠️
src/IRPrinter.cpp	0.00%	2 Missing and 1 partial ⚠️
src/Lower.cpp	72.72%	1 Missing and 2 partials ⚠️
src/Profiling.cpp	0.00%	0 Missing and 2 partials ⚠️
src/DeviceInterface.cpp	0.00%	0 Missing and 1 partial ⚠️
src/Target.h	50.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #9126   +/-   ##
=======================================
  Coverage        ?   69.83%           
=======================================
  Files           ?      256           
  Lines           ?    77781           
  Branches        ?    18606           
=======================================
  Hits            ?    54317           
  Misses          ?    18007           
  Partials        ?     5457           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[zvookin]

The high order question is whether this should just be another top-level target feature for ARM processors. If so, no special loop annotation is required and vector_bits_N can be used to give the size of the SME2 unit.

It means one cannot generate both SME2 and NEON in the same pipeline, but having looked over the architecture spec I'm not convinced that is useful. It is slightly convenient in terms of bounds inference, but performance wise switching between the modes clears the entire vector state so imposing a function call boundary there is hardly a problem. (Same is true for SVE/SVE2.) My reading of the architecture spec is that fine grained switching between SME2 and one of the other vector extensions is not a great idea.

Also per being a singular resource, its interaction with parallelism requires care at the level outside of Halide generated code.

I expect doing it this way limits the processing that can be specified, but that would be true inside the loop labelled SME2 anyway. This may have been discussed with a question as to whether to fail compilation or to fallback to e.g. NEON. Really the initial use case is specialized kernels that are written specifically for the SME2 hardware anyway so failing compilation is fine.

[stevesuzuki-arm]

It is true that switching between streaming mode has some overhead. So, very frequent transitions (e.g. in inner-most loop) should be avoided in terms of performance.
IMO, when scheduling a long complex pipeline, it is not uncommon that one wants to run cleanly-vectorized part in streaming mode and others in non-streaming (e.g. scalar processing, calling runtime function, etc). In such a case, there are a couple (or several) of intermediates generated with compute_root() as boundary, which I think is an common situation. I'd assume streaming SVE is generally beneficial if the workload is simply vectorized and scaled even without matmul (caveat: depending on the u-implementation). I have some internal app which has this kind of mixed targeting, but the switching cost is negligible as compared to the performance gain. So, I think asking the user to separate the Halide module just for streaming/non-streaming boundary sounds inconvenient.
That is the idea why I expose it as an scheduling option for user.