[AArch64][SVE] Implement isVScaleKnownToBeAPowerOfTwo

[dmgreen]

Summary

According to https://developer.arm.com/documentation/102105/ia-00/?lang=en

Arm is making a retrospective change to the SVE architecture to remove the capability of selecting a non-power-of-two vector length in non-Streaming SVE as well as in Streaming SVE mode. Specific updates as a result of this change will be communicated in due course.

This patch implements the isVScaleKnownToBeAPowerOfTwo method to teach DAG Combines that VScale will be known to be a power of 2, which helps reduce or simplify some expressions (notably the udiv in vector trip count expressions).

llvm/lib/Target/AArch64/AArch64ISelLowering.h

public:
   /// If a change in streaming mode is required on entry to/return from a
   /// function call it emits and returns the corresponding SMSTART or SMSTOP node.
   /// \p Entry tells whether this is before/after the Call, which is necessary
   /// because PSTATE.SM is only queried once.
   SDValue changeStreamingMode(SelectionDAG &DAG, SDLoc DL, bool Enable,
                               SDValue Chain, SDValue InFlag,
                               SDValue PStateSM, bool Entry) const;
 
+  bool isVScaleKnownToBeAPowerOfTwo() const override;
+
   // Normally SVE is only used for byte size vectors that do not fit within a
   // NEON vector. This changes when OverrideNEON is true, allowing SVE to be
   // used for 64bit and 128bit vectors as well.
   bool useSVEForFixedLengthVectorVT(EVT VT, bool OverrideNEON = false) const;
 
 private:
   /// Keep a pointer to the AArch64Subtarget around so that we can
   /// make the right decision when generating code for different targets.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
   }
   }
 }
 
 bool AArch64TargetLowering::mergeStoresAfterLegalization(EVT VT) const {
   return !Subtarget->useSVEForFixedLengthVectors();
 }
 
+bool AArch64TargetLowering::isVScaleKnownToBeAPowerOfTwo() const {
+  return true;
+}
+
 bool AArch64TargetLowering::useSVEForFixedLengthVectorVT(
     EVT VT, bool OverrideNEON) const {
   if (!VT.isFixedLengthVector() || !VT.isSimple())
     return false;
 
   // Don't use SVE for vectors we cannot scalarize if required.
   switch (VT.getVectorElementType().getSimpleVT().SimpleTy) {
   // Fixed length predicates should be promoted to i8.

llvm/test/CodeGen/AArch64/vscale-power-of-two.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64-none-eabi -mattr=+sve -verify-machineinstrs < %s | FileCheck %s
 
 define i64 @vscale_lshr(i64 %TC) {
 ; CHECK-LABEL: vscale_lshr:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    rdvl x8, #1
 ; CHECK-NEXT:    lsr x8, x8, #4
 ; CHECK-NEXT:    lsr x8, x8, #3
-; CHECK-NEXT:    udiv x9, x0, x8
-; CHECK-NEXT:    msub x0, x9, x8, x0
+; CHECK-NEXT:    sub x8, x8, #1
+; CHECK-NEXT:    and x0, x0, x8
 ; CHECK-NEXT:    ret
   %vscale = call i64 @llvm.vscale.i64()
   %shifted = lshr i64 %vscale, 3
   %urem = urem i64 %TC, %shifted
   ret i64 %urem
 }
 
 define i64 @vscale(i64 %TC) {
 ; CHECK-LABEL: vscale:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    rdvl x8, #1
 ; CHECK-NEXT:    lsr x8, x8, #4
-; CHECK-NEXT:    udiv x9, x0, x8
-; CHECK-NEXT:    msub x0, x9, x8, x0
+; CHECK-NEXT:    sub x8, x8, #1
+; CHECK-NEXT:    and x0, x0, x8
 ; CHECK-NEXT:    ret
   %vscale = call i64 @llvm.vscale.i64()
   %urem = urem i64 %TC, %vscale
   ret i64 %urem
 }
 
 define i64 @vscale_shl(i64 %TC) {
 ; CHECK-LABEL: vscale_shl:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    cnth x8
-; CHECK-NEXT:    udiv x9, x0, x8
-; CHECK-NEXT:    msub x0, x9, x8, x0
+; CHECK-NEXT:    sub x8, x8, #1
+; CHECK-NEXT:    and x0, x0, x8
 ; CHECK-NEXT:    ret
   %vscale = call i64 @llvm.vscale.i64()
   %shifted = shl i64 %vscale, 3
   %urem = urem i64 %TC, %shifted
   ret i64 %urem
 }
 
 declare i64 @llvm.vscale.i64()