Enable SME2 Streaming SVE in ARM#9126
Conversation
Added: - Target::SME2 definition - streaming_vector_bits in Target for SME2 - Auto-detect SME2 and streaming_vector_bits - sme_streaming() scheduling directive in Func and Pipeline - DeviceAPI::Host_SMEStreaming in IR "For" - LowerSMEStreamingTasks pass to extract streaming closure - Attribute in LoweredFunc for streaming closure - LLVM Function attribute to control streaming mode - NoInline to prevent streaming closure from inlined - "aarch64_pstate_sm_body" to emit smstart/smstop transition - Disable gather/scatter in SME streaming mode Tests: - Add correctness/sme_streaming - Run simd_op_check_sve2 in SME streaming mode - Add test to assert runtime streaming vscale
|
This PR is ready for review. I will touch on this in dev meeting if I have a chance. |
Reason:
While vector_bits is used across multiple target architectures,
streaming_vector_bits is aarch64 specific. So we choose to
use Target::Feature rather than a new member for arbitrary bits.
- Removed Target::streaming_vector_bits member variable
- Added Feature::SME_SVL{128,256,512,1024,2048}
Revert the changes in halide_error_vscale_invalid to avoid potential runtime breaking changes.
Because streaming_vector_bits member variable has been removed.
|
Based on the feedback in dev meeting, |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #9126 +/- ##
=======================================
Coverage ? 69.31%
=======================================
Files ? 255
Lines ? 78468
Branches ? 18781
=======================================
Hits ? 54389
Misses ? 18554
Partials ? 5525 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
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 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. |
|
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. |
| D3D12Compute, | ||
| Vulkan, | ||
| WebGPU, | ||
| Host_SMEStreaming, |
There was a problem hiding this comment.
Could this just be "SME"? Or "SMEStreaming"? I'm not sure what Host_ buys you.
There was a problem hiding this comment.
The nuance I intended with "Host" is that it is almost like running on host CPU because from SW point of view:
- Instructions are streamed from host, which looks similar to issueing other CPU instructions
- No device buffer
- No device runtime
I think the word "streaming" helps to emphasize the execution in streaming mode as there are other SME specific features/instructions legal even in non-streaming mode.
So, I'm fine with "SMEStreaming" or "Host_SMEStreaming". Do you prefer to remove "Host"?
| * When a loop is marked with sme_streaming(true), that loop including its inner loops | ||
| * are executed in Streaming mode. Marking with sme_streaming(false) prevents the loop | ||
| * from being executed in Streaming mode. */ | ||
| Func &sme_streaming(bool enable, const VarOrRVar &x = Var::outermost()); |
There was a problem hiding this comment.
I'm not sure I understand the function of the bool parameter here. Most of the other similar methods just mark a loop as something, and if you don't want that you don't call the method (or you call some other method like unroll(). Why is this not just the same as the hexagon method?
There was a problem hiding this comment.
test_2_stages_consumer_streaming_at() in test/correctness/sme_streaming.cpp is an example where false is set, which deals with not-ideal situation where producer tile needs to be computed in non-streaming due to some reason.
We may remove the bool parameter if we don't support this case in favor of simplicity. So, I think it is a design choice.
There was a problem hiding this comment.
Is this schedule:
g.compute_root().sme_streaming(true, x).split(x, xo, xi, 256);
// explicitly set false, otherwise streaming is enabled
f.compute_at(g, xo).sme_streaming(false);
equivalent to this?
g.compute_root().split(x, xo, xi, 256).sme_streaming(true, xi);
f.compute_at(g, xo);
I guess maybe they differ in how the generated function calls are structured? The sense I'm getting is that unlike other offload engines, with SME you can have host loop inside a device loop - you can "come back" from the device temporarily for some code, and this can be useful. Is that correct?
If so, I think we just want a .host(VarOrRVar) scheduling call that sets deviceAPI to Host for that loop. (Open to other opinions for the name).
There was a problem hiding this comment.
If we look at only the vectorization of arithmetics to compute output, they are equivalent. On the other hand, function structures are different.
produce g:
for x.xo<SMEStreaming>:
produce f:
for __outermost in [0, 0]: <== marked as non-streaming! (or .host)
for x:
f(...) = ...
consume f:
for x.xi in [0, 255]<SMEStreaming>:
g(...) = ...
VS
produce g:
for x.xo:
produce f:
for x:
f(...) = ...
consume f:
for x.xi in [0, 255]<SMEStreaming>:
g(...) = ...
And yes, your guess is exactly the idea behind this.
| return result; | ||
| } | ||
|
|
||
| int Target::natural_vector_size(const Halide::Type &t, bool is_sme_streaming) const { |
There was a problem hiding this comment.
This is awkwardly different between sme and other offload targets now. An arm pipeline could conceivably have something schedule on host, something schedule on a hexagon dsp, and something schedule for sme. I think unless we figure out a general solution we should just natural_vector_size as it is on main (i.e. it always returns the host vector size) and code will have to call sme_streaming_vector_bits if it's vectorizing a stage scheduled as sme_streaming
There was a problem hiding this comment.
I see. I will apply the change accordingly.
There was a problem hiding this comment.
Done. And updated the PR description as well.
Co-authored-by: Codex <codex@openai.com>
- Removed "Host_" prefix - Updated a few switch-case which was missing Co-authored-by: Codex <codex@openai.com>
|
With halide-llvm |
4 participants
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 usersthe option to control which loop is computed in streaming-mode.
The change introduces a new
Target::SME2feature with supplemental featuresTarget::SME_SVLDDD, where DDD represents streaming vector length in bits (e.g. 128, 256, 512, ...). IfTarget::SME2is enabled, exactly one ofTarget::SME_SVLDDDfeature must be enabled as well.Please note
natural_vector_size()always returns the host vector size and users have to callTarget::sme_streaming_vector_bits()to work out device native vector size if it's vectorizing a stage scheduled as sme_streaming.In Halide lowering, a new
LowerSMEStreamingTaskspass 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_bodyto emit smstart/smstop transitionNoInlineto prevent streaming closure from inlined to non-streaming functionIn CodeGen,
target_vscale()depends on whether streaming-mode or notand 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:
SME2andSME_SVLDDDtarget features on host CPUChecklist