[AArch64][SME] Disable remat of VL-dependent ops when function changes streaming mode.

[sdesmalen]

Summary

This is a way to prevent the register allocator from inserting instructions
which behave differently for different runtime vector-lengths, inside a
call-sequence which changes the streaming-SVE mode before/after the call.

I've considered using BUNDLEs in Machine IR, but found that using this is
not possible for a few reasons:

• Most passes don't look inside BUNDLEs, but some passes would need to look inside these call-sequence bundles, for example the PrologEpilog pass (to remove the CALLSEQSTART/END), a PostRA pass to remove COPY instructions, or the AArch64PseudoExpand pass.
• Within the streaming-mode-changing call sequence, one of the instructions is a CALLSEQEND. The corresponding CALLSEQBEGIN (AArch64::ADJCALLSTACKUP) is outside this sequence. This means we'd end up with a BUNDLE that has [SMSTART, COPY, BL, ADJCALLSTACKUP, COPY, SMSTOP]. The MachineVerifier doesn't accept this, and we also can't move the CALLSEQSTART into the call sequence.

Maybe in the future we could model this differently by modelling
the runtime vector-length as a value that's used by certain operations
(similar to e.g. NCZV flags) and clobbered by SMSTART/MMSTOP, such that the
register allocator can consider these as actual dependences and avoid
rematerialization. For now we just want to address the immediate problem.

[aemerson]

I have a slight concern that relying on the MFI's flag being set on a DAG codegen call is brittle, in case in future there are other pathways that change streaming mode. Scanning the function seems a more reliable way to set the flag, but for now this is ok with me.

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8540

physical

[sdesmalen]

Use isVirtual() explicitly rather than relying on getRegClassOrNull().
(this also removed the need for the assert)

[sdesmalen]

In D159193#4628454, @aemerson wrote:

I have a slight concern that relying on the MFI's flag being set on a DAG codegen call is brittle, in case in future there are other pathways that change streaming mode. Scanning the function seems a more reliable way to set the flag, but for now this is ok with me.

Thanks for pointing out, I had a patch for that which I forgot to push. 8f469bfedc907a391da7207ba248afdb7e8c555d ensures that any streaming-mode changes are always done through the function changeStreamingMode.

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8540

I've actually simplified it with an explicit check for isVirtual, rather than relying on getRegClassOrNull and then use getRegClass() instead.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8547

Perhaps worth saying "Avoid rematerializing rematerializable instructions...". This might be obvious but for a second I nearly asked "Should this block also include the inc/dec instructions?"

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

static bool checkZExtBool(SDValue Arg, const SelectionDAG &DAG) {
   KnownBits Bits = DAG.computeKnownBits(Arg, 4);
   bool ZExtBool = (Bits.Zero & RequredZero) == RequredZero;
   return ZExtBool;
 }
 
 SDValue AArch64TargetLowering::changeStreamingMode(
     SelectionDAG &DAG, SDLoc DL, bool Enable,
     SDValue Chain, SDValue InGlue, SDValue PStateSM, bool Entry) const {
+  MachineFunction &MF = DAG.getMachineFunction();
+  AArch64FunctionInfo *FuncInfo = MF.getInfo<AArch64FunctionInfo>();
+  FuncInfo->setHasStreamingModeChanges(true);
+
   const AArch64RegisterInfo *TRI = Subtarget->getRegisterInfo();
   SDValue RegMask = DAG.getRegisterMask(TRI->getSMStartStopCallPreservedMask());
   SDValue MSROp =
       DAG.getTargetConstant((int32_t)AArch64SVCR::SVCRSM, DL, MVT::i32);
 
   SDValue ExpectedSMVal =
       DAG.getTargetConstant(Entry ? Enable : !Enable, DL, MVT::i64);
   SmallVector<SDValue> Ops = {Chain, MSROp, PStateSM, ExpectedSMVal, RegMask};

llvm/lib/Target/AArch64/AArch64InstrInfo.h

public:
 
   static void decomposeStackOffsetForFrameOffsets(const StackOffset &Offset,
                                                   int64_t &NumBytes,
                                                   int64_t &NumPredicateVectors,
                                                   int64_t &NumDataVectors);
   static void decomposeStackOffsetForDwarfOffsets(const StackOffset &Offset,
                                                   int64_t &ByteSized,
                                                   int64_t &VGSized);
+
+  bool isReallyTriviallyReMaterializable(const MachineInstr &MI) const override;
 #define GET_INSTRINFO_HELPER_DECLS
 #include "AArch64GenInstrInfo.inc"
 
 protected:
   /// If the specific machine instruction is an instruction that moves/copies
   /// value from one register to another register return destination and source
   /// registers as machine operands.
   std::optional<DestSourcePair>

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

 
 unsigned llvm::getBLRCallOpcode(const MachineFunction &MF) {
   if (MF.getSubtarget<AArch64Subtarget>().hardenSlsBlr())
     return AArch64::BLRNoIP;
   else
     return AArch64::BLR;
 }
 
+bool AArch64InstrInfo::isReallyTriviallyReMaterializable(
+    const MachineInstr &MI) const {
+  const MachineFunction &MF = *MI.getMF();
+  const AArch64FunctionInfo &AFI = *MF.getInfo<AArch64FunctionInfo>();
+
+  // If the function contains changes to streaming mode, then there
+  // is a danger that rematerialised instructions end up between
+  // instruction sequences (e.g. call sequences, or prolog/epilogue)
+  // where the streaming-SVE mode is temporarily changed.
+  if (AFI.hasStreamingModeChanges()) {
+    // Avoid rematerializing instructions that use/define scalable values,
+    // such as 'pfalse' or 'ptrue', which result in different results
+    // when the runtime vector length is different.
+    const MachineRegisterInfo &MRI = MF.getRegInfo();
+    if (any_of(MI.operands(), [&MRI](const MachineOperand &MO) {
+          if (!MO.isReg())
+            return false;
+          if (const TargetRegisterClass *RC =
+                  MRI.getRegClassOrNull(MO.getReg()))
+            return AArch64::ZPRRegClass.hasSubClassEq(RC) ||
+                   AArch64::PPRRegClass.hasSubClassEq(RC);
+          assert(!MO.getReg().isVirtual() && "Expected a phsyical register");
+          return AArch64::ZPRRegClass.contains(MO.getReg()) ||
+                 AArch64::PPRRegClass.contains(MO.getReg());
+        }))
+      return false;
+
+    // Avoid rematerializing instructions that return a value that is
+    // different depending on vector length, even when it is not returned
+    // in a scalable vector/predicate register.
+    switch (MI.getOpcode()) {
+    default:
+      break;
+    case AArch64::RDVLI_XI:
+    case AArch64::ADDVL_XXI:
+    case AArch64::ADDPL_XXI:
+    case AArch64::CNTB_XPiI:
+    case AArch64::CNTH_XPiI:
+    case AArch64::CNTW_XPiI:
+    case AArch64::CNTD_XPiI:
+      return false;
+    }
+  }
+
+  return TargetInstrInfo::isReallyTriviallyReMaterializable(MI);
+}
+
 #define GET_INSTRINFO_HELPERS
 #define GET_INSTRMAP_INFO
 #include "AArch64GenInstrInfo.inc"

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h

class AArch64FunctionInfo final : public MachineFunctionInfo {
 
   /// The function has Scalable Vector or Scalable Predicate register argument
   /// or return type
   bool IsSVECC = false;
 
   /// The frame-index for the TPIDR2 object used for lazy saves.
   Register LazySaveTPIDR2Obj = 0;
 
+  /// Whether this function changes streaming mode within the function.
+  bool HasStreamingModeChanges = false;
 
   /// True if the function need unwind information.
   mutable std::optional<bool> NeedsDwarfUnwindInfo;
 
   /// True if the function need asynchronous unwind information.
   mutable std::optional<bool> NeedsAsyncDwarfUnwindInfo;
 
 public:
#endif
   void setSwiftAsyncContextFrameIdx(int FI) {
     SwiftAsyncContextFrameIdx = FI;
   }
   int getSwiftAsyncContextFrameIdx() const { return SwiftAsyncContextFrameIdx; }
 
   bool needsDwarfUnwindInfo(const MachineFunction &MF) const;
   bool needsAsyncDwarfUnwindInfo(const MachineFunction &MF) const;
 
+  bool hasStreamingModeChanges() const { return HasStreamingModeChanges; }
+  void setHasStreamingModeChanges(bool HasChanges) {
+    HasStreamingModeChanges = HasChanges;
+  }
+
 private:
   // Hold the lists of LOHs.
   MILOHContainer LOHContainerSet;
   SetOfInstructions LOHRelated;
 
   SmallVector<std::pair<unsigned, MCSymbol *>, 2> JumpTableEntryInfo;
 };
 

llvm/test/CodeGen/AArch64/sme-disable-rematerialize-with-streaming-mode-changes.ll

+; RUN: llc < %s | FileCheck %s
+
+target triple = "aarch64"
+
+
+define void @dont_rematerialize_cntd(i32 %N) #0 {
+; CHECK-LABEL: dont_rematerialize_cntd:
+; CHECK:        cntd
+; CHECK:        smstop sm
+; CHECK-NOT:    cntd
+; CHECK:        bl      foo
+; CHECK:        smstart  sm
+entry:
+  %cmp2 = icmp sgt i32 %N, 0
+  br i1 %cmp2, label %for.body, label %for.cond.cleanup
+
+for.body:                                         ; preds = %entry, %for.body
+  %index.03 = phi i32 [ %inc, %for.body ], [ 0, %entry ]
+  call void asm sideeffect "", "~{x19},~{x20},~{x21},~{x22},~{x23},~{x24},~{x25},~{x26},~{x27}"() nounwind
+  %.tr = call i32 @llvm.vscale.i32()
+  %conv = shl nuw nsw i32 %.tr, 4
+  call void @foo(i32 %conv)
+  %inc = add nuw nsw i32 %index.03, 1
+  %exitcond.not = icmp eq i32 %inc, %N
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.body, %entry
+  ret void
+}
+
+; This test doesn't strictly make sense, because it passes a scalable predicate
+; to a function, which makes little sense if the VL is not the same in/out of
+; streaming-SVE mode. If the VL is known to be the same, then we could just as
+; well rematerialize the `ptrue` inside the call sequence. However, the purpose
+; of this test is more to ensure that the logic works, which may also trigger
+; when the value is not being passed as argument (e.g. when it is hoisted from
+; a loop and placed inside the call sequence).
+;
+; FIXME: This test also exposes another bug, where the 'mul vl' addressing mode
+; is used before/after the smstop. This will be fixed in a future patch.
+define void @dont_rematerialize_ptrue(i32 %N) #0 {
+; CHECK-LABEL: dont_rematerialize_ptrue:
+; CHECK:        ptrue [[PTRUE:p[0-9]+]].b
+; CHECK:        str [[PTRUE]], [[[SPILL_ADDR:.*]]]
+; CHECK:        smstop sm
+; CHECK:        ldr p0, [[[SPILL_ADDR]]]
+; CHECK-NOT:    ptrue
+; CHECK:        bl      bar
+; CHECK:        smstart  sm
+entry:
+  %cmp2 = icmp sgt i32 %N, 0
+  br i1 %cmp2, label %for.body, label %for.cond.cleanup
+
+for.body:                                         ; preds = %entry, %for.body
+  %index.03 = phi i32 [ %inc, %for.body ], [ 0, %entry ]
+  call void asm sideeffect "", "~{x19},~{x20},~{x21},~{x22},~{x23},~{x24},~{x25},~{x26},~{x27}"() nounwind
+  %ptrue.ins = insertelement <vscale x 16 x i1> poison, i1 1, i32 0
+  %ptrue = shufflevector <vscale x 16 x i1> %ptrue.ins, <vscale x 16 x i1> poison, <vscale x 16 x i32> zeroinitializer
+  call void @bar(<vscale x 16 x i1> %ptrue)
+  %inc = add nuw nsw i32 %index.03, 1
+  %exitcond.not = icmp eq i32 %inc, %N
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.body, %entry
+  ret void
+}
+declare void @foo(i32)
+declare void @bar(<vscale x 16 x i1>)
+declare i32 @llvm.vscale.i32()
+
+attributes #0 = { "aarch64_pstate_sm_enabled" "frame-pointer"="non-leaf" "target-features"="+sme,+sve" }