llvm/docs/AArch64SME.rst

   1 *****************************************************
   2 Support for AArch64 Scalable Matrix Extension in LLVM
   3 *****************************************************
   4
   5 .. contents::
   6    :local:
   7
   8 1. Introduction
   9 ===============
  10
  11 The :ref:`AArch64 SME ACLE <aarch64_sme_acle>` provides a number of
  12 attributes for users to control PSTATE.SM and PSTATE.ZA.
  13 The :ref:`AArch64 SME ABI<aarch64_sme_abi>` describes the requirements for
  14 calls between functions when at least one of those functions uses PSTATE.SM or
  15 PSTATE.ZA.
  16
  17 This document describes how the SME ACLE attributes map to LLVM IR
  18 attributes and how LLVM lowers these attributes to implement the rules and
  19 requirements of the ABI.
  20
  21 Below we describe the LLVM IR attributes and their relation to the C/C++
  22 level ACLE attributes:
  23
  24 ``aarch64_pstate_sm_enabled``
  25     is used for functions with ``__attribute__((arm_streaming))``
  26
  27 ``aarch64_pstate_sm_compatible``
  28     is used for functions with ``__attribute__((arm_streaming_compatible))``
  29
  30 ``aarch64_pstate_sm_body``
  31   is used for functions with ``__attribute__((arm_locally_streaming))`` and is
  32   only valid on function definitions (not declarations)
  33
  34 ``aarch64_pstate_za_new``
  35   is used for functions with ``__attribute__((arm_new_za))``
  36
  37 ``aarch64_pstate_za_shared``
  38   is used for functions with ``__attribute__((arm_shared_za))``
  39
  40 ``aarch64_pstate_za_preserved``
  41   is used for functions with ``__attribute__((arm_preserves_za))``
  42
  43 ``aarch64_expanded_pstate_za``
  44   is used for functions with ``__attribute__((arm_new_za))``
  45
  46 Clang must ensure that the above attributes are added both to the
  47 function's declaration/definition as well as to their call-sites. This is
  48 important for calls to attributed function pointers, where there is no
  49 definition or declaration available.
  50
  51
  52 2. Handling PSTATE.SM
  53 =====================
  54
  55 When changing PSTATE.SM the execution of FP/vector operations may be transferred
  56 to another processing element. This has three important implications:
  57
  58 * The runtime SVE vector length may change.
  59
  60 * The contents of FP/AdvSIMD/SVE registers are zeroed.
  61
  62 * The set of allowable instructions changes.
  63
  64 This leads to certain restrictions on IR and optimizations. For example, it
  65 is undefined behaviour to share vector-length dependent state between functions
  66 that may operate with different values for PSTATE.SM. Front-ends must honour
  67 these restrictions when generating LLVM IR.
  68
  69 Even though the runtime SVE vector length may change, for the purpose of LLVM IR
  70 and almost all parts of CodeGen we can assume that the runtime value for
  71 ``vscale`` does not. If we let the compiler insert the appropriate ``smstart``
  72 and ``smstop`` instructions around call boundaries, then the effects on SVE
  73 state can be mitigated. By limiting the state changes to a very brief window
  74 around the call we can control how the operations are scheduled and how live
  75 values remain preserved between state transitions.
  76
  77 In order to control PSTATE.SM at this level of granularity, we use function and
  78 callsite attributes rather than intrinsics.
  79
  80
  81 Restrictions on attributes
  82 --------------------------
  83
  84 * It is undefined behaviour to pass or return (pointers to) scalable vector
  85   objects to/from functions which may use a different SVE vector length.
  86   This includes functions with a non-streaming interface, but marked with
  87   ``aarch64_pstate_sm_body``.
  88
  89 * It is not allowed for a function to be decorated with both
  90   ``aarch64_pstate_sm_compatible`` and ``aarch64_pstate_sm_enabled``.
  91
  92 * It is not allowed for a function to be decorated with both
  93   ``aarch64_pstate_za_new`` and ``aarch64_pstate_za_preserved``.
  94
  95 * It is not allowed for a function to be decorated with both
  96   ``aarch64_pstate_za_new`` and ``aarch64_pstate_za_shared``.
  97
  98 These restrictions also apply in the higher level SME ACLE, which means we can
  99 emit diagnostics in Clang to signal users about incorrect behaviour.
 100
 101
 102 Compiler inserted streaming-mode changes
 103 ----------------------------------------
 104
 105 The table below describes the transitions in PSTATE.SM the compiler has to
 106 account for when doing calls between functions with different attributes.
 107 In this table, we use the following abbreviations:
 108
 109 ``N``
 110   functions with a normal interface (PSTATE.SM=0 on entry, PSTATE.SM=0 on
 111   return)
 112
 113 ``S``
 114   functions with a Streaming interface (PSTATE.SM=1 on entry, PSTATE.SM=1
 115   on return)
 116
 117 ``SC``
 118   functions with a Streaming-Compatible interface (PSTATE.SM can be
 119   either 0 or 1 on entry, and is unchanged on return).
 120
 121 Functions with ``__attribute__((arm_locally_streaming))`` are excluded from this
 122 table because for the caller the attribute is synonymous to 'streaming', and
 123 for the callee it is merely an implementation detail that is explicitly not
 124 exposed to the caller.
 125
 126 .. table:: Combinations of calls for functions with different attributes
 127
 128    ==== ==== =============================== ============================== ==============================
 129    From To   Before call                     After call                     After exception
 130    ==== ==== =============================== ============================== ==============================
 131    N    N
 132    N    S    SMSTART                         SMSTOP
 133    N    SC
 134    S    N    SMSTOP                          SMSTART                        SMSTART
 135    S    S                                                                   SMSTART
 136    S    SC                                                                  SMSTART
 137    SC   N    If PSTATE.SM before call is 1,  If PSTATE.SM before call is 1, If PSTATE.SM before call is 1,
 138              then SMSTOP                     then SMSTART                   then SMSTART
 139    SC   S    If PSTATE.SM before call is 0,  If PSTATE.SM before call is 0, If PSTATE.SM before call is 1,
 140              then SMSTART                    then SMSTOP                    then SMSTART
 141    SC   SC                                                                  If PSTATE.SM before call is 1,
 142                                                                             then SMSTART
 143    ==== ==== =============================== ============================== ==============================
 144
 145
 146 Because changing PSTATE.SM zeroes the FP/vector registers, it is best to emit
 147 the ``smstart`` and ``smstop`` instructions before register allocation, so that
 148 the register allocator can spill/reload registers around the mode change.
 149
 150 The compiler should also have sufficient information on which operations are
 151 part of the call/function's arguments/result and which operations are part of
 152 the function's body, so that it can place the mode changes in exactly the right
 153 position. The suitable place to do this seems to be SelectionDAG, where it lowers
 154 the call's arguments/return values to implement the specified calling convention.
 155 SelectionDAG provides Chains and Glue to specify the order of operations and give
 156 preliminary control over the instruction's scheduling.
 157
 158
 159 Example of preserving state
 160 ---------------------------
 161
 162 When passing and returning a ``float`` value to/from a function
 163 that has a streaming interface from a function that has a normal interface, the
 164 call-site will need to ensure that the argument/result registers are preserved
 165 and that no other code is scheduled in between the ``smstart/smstop`` and the call.
 166
 167 .. code-block:: llvm
 168
 169     define float @foo(float %f) nounwind {
 170       %res = call float @bar(float %f) "aarch64_pstate_sm_enabled"
 171       ret float %res
 172     }
 173
 174     declare float @bar(float) "aarch64_pstate_sm_enabled"
 175
 176 The program needs to preserve the value of the floating point argument and
 177 return value in register ``s0``:
 178
 179 .. code-block:: none
 180
 181     foo:                                    // @foo
 182     // %bb.0:
 183             stp     d15, d14, [sp, #-80]!           // 16-byte Folded Spill
 184             stp     d13, d12, [sp, #16]             // 16-byte Folded Spill
 185             stp     d11, d10, [sp, #32]             // 16-byte Folded Spill
 186             stp     d9, d8, [sp, #48]               // 16-byte Folded Spill
 187             str     x30, [sp, #64]                  // 8-byte Folded Spill
 188             str     s0, [sp, #76]                   // 4-byte Folded Spill
 189             smstart sm
 190             ldr     s0, [sp, #76]                   // 4-byte Folded Reload
 191             bl      bar
 192             str     s0, [sp, #76]                   // 4-byte Folded Spill
 193             smstop  sm
 194             ldp     d9, d8, [sp, #48]               // 16-byte Folded Reload
 195             ldp     d11, d10, [sp, #32]             // 16-byte Folded Reload
 196             ldp     d13, d12, [sp, #16]             // 16-byte Folded Reload
 197             ldr     s0, [sp, #76]                   // 4-byte Folded Reload
 198             ldr     x30, [sp, #64]                  // 8-byte Folded Reload
 199             ldp     d15, d14, [sp], #80             // 16-byte Folded Reload
 200             ret
 201
 202 Setting the correct register masks on the ISD nodes and inserting the
 203 ``smstart/smstop`` in the right places should ensure this is done correctly.
 204
 205
 206 Instruction Selection Nodes
 207 ---------------------------
 208
 209 .. code-block:: none
 210
 211   AArch64ISD::SMSTART Chain, [SM|ZA|Both], CurrentState, ExpectedState[, RegMask]
 212   AArch64ISD::SMSTOP  Chain, [SM|ZA|Both], CurrentState, ExpectedState[, RegMask]
 213
 214 The ``SMSTART/SMSTOP`` nodes take ``CurrentState`` and ``ExpectedState`` operand for
 215 the case of a conditional SMSTART/SMSTOP. The instruction will only be executed
 216 if CurrentState != ExpectedState.
 217
 218 When ``CurrentState`` and ``ExpectedState`` can be evaluated at compile-time
 219 (i.e. they are both constants) then an unconditional ``smstart/smstop``
 220 instruction is emitted. Otherwise the node is matched to a Pseudo instruction
 221 which expands to a compare/branch and a ``smstart/smstop``. This is necessary to
 222 implement transitions from ``SC -> N`` and ``SC -> S``.
 223
 224
 225 Unchained Function calls
 226 ------------------------
 227 When a function with "``aarch64_pstate_sm_enabled``" calls a function that is not
 228 streaming compatible, the compiler has to insert a SMSTOP before the call and
 229 insert a SMSTOP after the call.
 230
 231 If the function that is called is an intrinsic with no side-effects which in
 232 turn is lowered to a function call (e.g. ``@llvm.cos()``), then the call to
 233 ``@llvm.cos()`` is not part of any Chain; it can be scheduled freely.
 234
 235 Lowering of a Callsite creates a small chain of nodes which:
 236
 237 - starts a call sequence
 238
 239 - copies input values from virtual registers to physical registers specified by
 240   the ABI
 241
 242 - executes a branch-and-link
 243
 244 - stops the call sequence
 245
 246 - copies the output values from their physical registers to virtual registers
 247
 248 When the callsite's Chain is not used, only the result value from the chained
 249 sequence is used, but the Chain itself is discarded.
 250
 251 The ``SMSTART`` and ``SMSTOP`` ISD nodes return a Chain, but no real
 252 values, so when the ``SMSTART/SMSTOP`` nodes are part of a Chain that isn't
 253 used, these nodes are not considered for scheduling and are
 254 removed from the DAG.  In order to prevent these nodes
 255 from being removed, we need a way to ensure the results from the
 256 ``CopyFromReg`` can only be **used after** the ``SMSTART/SMSTOP`` has been
 257 executed.
 258
 259 We can use a CopyToReg -> CopyFromReg sequence for this, which moves the
 260 value to/from a virtual register and chains these nodes with the
 261 SMSTART/SMSTOP to make them part of the expression that calculates
 262 the result value. The resulting COPY nodes are removed by the register
 263 allocator.
 264
 265 The example below shows how this is used in a DAG that does not link
 266 together the result by a Chain, but rather by a value:
 267
 268 .. code-block:: none
 269
 270                t0: ch,glue = AArch64ISD::SMSTOP ...
 271              t1: ch,glue = ISD::CALL ....
 272            t2: res,ch,glue = CopyFromReg t1, ...
 273          t3: ch,glue = AArch64ISD::SMSTART t2:1, ....   <- this is now part of the expression that returns the result value.
 274        t4: ch = CopyToReg t3, Register:f64 %vreg, t2
 275      t5: res,ch = CopyFromReg t4, Register:f64 %vreg
 276    t6: res = FADD t5, t9
 277
 278 We also need this for locally streaming functions, where an ``SMSTART`` needs to
 279 be inserted into the DAG at the start of the function.
 280
 281 Functions with __attribute__((arm_locally_streaming))
 282 -----------------------------------------------------
 283
 284 If a function is marked as ``arm_locally_streaming``, then the runtime SVE
 285 vector length in the prologue/epilogue may be different from the vector length
 286 in the function's body. This happens because we invoke smstart after setting up
 287 the stack-frame and similarly invoke smstop before deallocating the stack-frame.
 288
 289 To ensure we use the correct SVE vector length to allocate the locals with, we
 290 can use the streaming vector-length to allocate the stack-slots through the
 291 ``ADDSVL`` instruction, even when the CPU is not yet in streaming mode.
 292
 293 This only works for locals and not callee-save slots, since LLVM doesn't support
 294 mixing two different scalable vector lengths in one stack frame. That means that the
 295 case where a function is marked ``arm_locally_streaming`` and needs to spill SVE
 296 callee-saves in the prologue is currently unsupported.  However, it is unlikely
 297 for this to happen without user intervention, because ``arm_locally_streaming``
 298 functions cannot take or return vector-length-dependent values. This would otherwise
 299 require forcing both the SVE PCS using '``aarch64_sve_pcs``' combined with using
 300 ``arm_locally_streaming`` in order to encounter this problem. This combination
 301 can be prevented in Clang through emitting a diagnostic.
 302
 303
 304 An example of how the prologue/epilogue would look for a function that is
 305 attributed with ``arm_locally_streaming``:
 306
 307 .. code-block:: c++
 308
 309     #define N 64
 310
 311     void __attribute__((arm_streaming_compatible)) some_use(svfloat32_t *);
 312
 313     // Use a float argument type, to check the value isn't clobbered by smstart.
 314     // Use a float return type to check the value isn't clobbered by smstop.
 315     float __attribute__((noinline, arm_locally_streaming)) foo(float arg) {
 316       // Create local for SVE vector to check local is created with correct
 317       // size when not yet in streaming mode (ADDSVL).
 318       float array[N];
 319       svfloat32_t vector;
 320
 321       some_use(&vector);
 322       svst1_f32(svptrue_b32(), &array[0], vector);
 323       return array[N - 1] + arg;
 324     }
 325
 326 should use ADDSVL for allocating the stack space and should avoid clobbering
 327 the return/argument values.
 328
 329 .. code-block:: none
 330
 331     _Z3foof:                                // @_Z3foof
 332     // %bb.0:                               // %entry
 333             stp     d15, d14, [sp, #-96]!           // 16-byte Folded Spill
 334             stp     d13, d12, [sp, #16]             // 16-byte Folded Spill
 335             stp     d11, d10, [sp, #32]             // 16-byte Folded Spill
 336             stp     d9, d8, [sp, #48]               // 16-byte Folded Spill
 337             stp     x29, x30, [sp, #64]             // 16-byte Folded Spill
 338             add     x29, sp, #64
 339             str     x28, [sp, #80]                  // 8-byte Folded Spill
 340             addsvl  sp, sp, #-1
 341             sub     sp, sp, #256
 342             str     s0, [x29, #28]                  // 4-byte Folded Spill
 343             smstart sm
 344             sub     x0, x29, #64
 345             addsvl  x0, x0, #-1
 346             bl      _Z10some_usePu13__SVFloat32_t
 347             sub     x8, x29, #64
 348             ptrue   p0.s
 349             ld1w    { z0.s }, p0/z, [x8, #-1, mul vl]
 350             ldr     s1, [x29, #28]                  // 4-byte Folded Reload
 351             st1w    { z0.s }, p0, [sp]
 352             ldr     s0, [sp, #252]
 353             fadd    s0, s0, s1
 354             str     s0, [x29, #28]                  // 4-byte Folded Spill
 355             smstop  sm
 356             ldr     s0, [x29, #28]                  // 4-byte Folded Reload
 357             addsvl  sp, sp, #1
 358             add     sp, sp, #256
 359             ldp     x29, x30, [sp, #64]             // 16-byte Folded Reload
 360             ldp     d9, d8, [sp, #48]               // 16-byte Folded Reload
 361             ldp     d11, d10, [sp, #32]             // 16-byte Folded Reload
 362             ldp     d13, d12, [sp, #16]             // 16-byte Folded Reload
 363             ldr     x28, [sp, #80]                  // 8-byte Folded Reload
 364             ldp     d15, d14, [sp], #96             // 16-byte Folded Reload
 365             ret
 366
 367
 368 Preventing the use of illegal instructions in Streaming Mode
 369 ------------------------------------------------------------
 370
 371 * When executing a program in streaming-mode (PSTATE.SM=1) a subset of SVE/SVE2
 372   instructions and most AdvSIMD/NEON instructions are invalid.
 373
 374 * When executing a program in normal mode (PSTATE.SM=0), a subset of SME
 375   instructions are invalid.
 376
 377 * Streaming-compatible functions must only use instructions that are valid when
 378   either PSTATE.SM=0 or PSTATE.SM=1.
 379
 380 The value of PSTATE.SM is not controlled by the feature flags, but rather by the
 381 function attributes. This means that we can compile for '``+sme``' and the compiler
 382 will code-generate any instructions, even if they are not legal under the requested
 383 streaming mode. The compiler needs to use the function attributes to ensure the
 384 compiler doesn't do transformations under the assumption that certain operations
 385 are available at runtime.
 386
 387 We made a conscious choice not to model this with feature flags, because we
 388 still want to support inline-asm in either mode (with the user placing
 389 smstart/smstop manually), and this became rather complicated to implement at the
 390 individual instruction level (see `D120261 <https://reviews.llvm.org/D120261>`_
 391 and `D121208 <https://reviews.llvm.org/D121208>`_) because of limitations in
 392 TableGen.
 393
 394 As a first step, this means we'll disable vectorization (LoopVectorize/SLP)
 395 entirely when the a function has either of the ``aarch64_pstate_sm_enabled``,
 396 ``aarch64_pstate_sm_body`` or ``aarch64_pstate_sm_compatible`` attributes,
 397 in order to avoid the use of vector instructions.
 398
 399 Later on we'll aim to relax these restrictions to enable scalable
 400 auto-vectorization with a subset of streaming-compatible instructions, but that
 401 requires changes to the CostModel, Legalization and SelectionDAG lowering.
 402
 403 We will also emit diagnostics in Clang to prevent the use of
 404 non-streaming(-compatible) operations, e.g. through ACLE intrinsics, when a
 405 function is decorated with the streaming mode attributes.
 406
 407
 408 Other things to consider
 409 ------------------------
 410
 411 * Inlining must be disabled when the call-site needs to toggle PSTATE.SM or
 412   when the callee's function body is executed in a different streaming mode than
 413   its caller. This is needed because function calls are the boundaries for
 414   streaming mode changes.
 415
 416 * Tail call optimization must be disabled when the call-site needs to toggle
 417   PSTATE.SM, such that the caller can restore the original value of PSTATE.SM.
 418
 419
 420 3. Handling PSTATE.ZA
 421 =====================
 422
 423 In contrast to PSTATE.SM, enabling PSTATE.ZA does not affect the SVE vector
 424 length and also doesn't clobber FP/AdvSIMD/SVE registers. This means it is safe
 425 to toggle PSTATE.ZA using intrinsics. This also makes it simpler to setup a
 426 lazy-save mechanism for calls to private-ZA functions (i.e. functions that may
 427 either directly or indirectly clobber ZA state).
 428
 429 For the purpose of handling functions marked with ``aarch64_pstate_za_new``,
 430 we have introduced a new LLVM IR pass (SMEABIPass) that is run just before
 431 SelectionDAG. Any such functions dealt with by this pass are marked with
 432 ``aarch64_expanded_pstate_za``.
 433
 434 Setting up a lazy-save
 435 ----------------------
 436
 437 Committing a lazy-save
 438 ----------------------
 439
 440 Exception handling and ZA
 441 -------------------------
 442
 443 4. Types
 444 ========
 445
 446 AArch64 Predicate-as-Counter Type
 447 ---------------------------------
 448
 449 :Overview:
 450
 451 The predicate-as-counter type represents the type of a predicate-as-counter
 452 value held in a AArch64 SVE predicate register. Such a value contains
 453 information about the number of active lanes, the element width and a bit that
 454 tells whether the generated mask should be inverted. ACLE intrinsics should be
 455 used to move the predicate-as-counter value to/from a predicate vector.
 456
 457 There are certain limitations on the type:
 458
 459 * The type can be used for function parameters and return values.
 460
 461 * The supported LLVM operations on this type are limited to ``load``, ``store``,
 462   ``phi``, ``select`` and ``alloca`` instructions.
 463
 464 The predicate-as-counter type is a scalable type.
 465
 466 :Syntax:
 467
 468 ::
 469
 470       target("aarch64.svcount")
 471
 472
 473
 474 5. References
 475 =============
 476
 477     .. _aarch64_sme_acle:
 478
 479 1.  `SME ACLE Pull-request <https://github.com/ARM-software/acle/pull/188>`__
 480
 481     .. _aarch64_sme_abi:
 482
 483 2.  `SME ABI Pull-request <https://github.com/ARM-software/abi-aa/pull/123>`__