[SVE][CodeGen] Fix bug when falling back to DAG ISel

[david-arm]

Summary

In an earlier commit 584d0d5c1749c13625a5d322178ccb4121eea610 I
added functionality to allow AArch64 CodeGen support for falling
back to DAG ISel when Global ISel encounters scalable vector
types. However, it seems that we were not falling back early
enough as llvm::getLLTForType was still being invoked for scalable
vector types.

I've added a new fallback function to the call lowering class in
order to catch this problem early enough, rather than wait for
lowerFormalArguments to reject scalable vector types.

[sdesmalen]

llvm/lib/Target/AArch64/GISel/AArch64CallLowering.cpp

431

Should this function also check the return value?

nit: use llvm::any_of instead of a loop?

[david-arm]

llvm/lib/Target/AArch64/GISel/AArch64CallLowering.cpp

431

Strictly speaking we didn't have to catch the return value here as it was being caught by the ret instruction, however it is a nice optimisation to catch it early on here so I've added the check.

[sdesmalen]

LGTM!

llvm/include/llvm/CodeGen/GlobalISel/CallLowering.h

public:
 
   /// This hook behaves as the extended lowerReturn function, but for targets
   /// that do not support swifterror value promotion.
   virtual bool lowerReturn(MachineIRBuilder &MIRBuilder, const Value *Val,
                            ArrayRef<Register> VRegs) const {
     return false;
   }
 
+  virtual bool fallBackToDAGISel(const Function &F) const { return false; }
+
   /// This hook must be implemented to lower the incoming (formal)
   /// arguments, described by \p VRegs, for GlobalISel. Each argument
   /// must end up in the related virtual registers described by \p VRegs.
   /// In other words, the first argument should end up in \c VRegs[0],
   /// the second in \c VRegs[1], and so on. For each argument, there will be one
   /// register for each non-aggregate type, as returned by \c computeValueLLTs.
   /// \p MIRBuilder is set to the proper insertion for the argument
   /// lowering.

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

if (!HasMustTailInVarArgFn)
       HasMustTailInVarArgFn = checkForMustTailInVarArgFn(IsVarArg, BB);
   }
 
   MF->getFrameInfo().setHasMustTailInVarArgFunc(HasMustTailInVarArgFn);
 
   // Make our arguments/constants entry block fallthrough to the IR entry block.
   EntryBB->addSuccessor(&getMBB(F.front()));
 
+  if (CLI->fallBackToDAGISel(F)) {
+    OptimizationRemarkMissed R("gisel-irtranslator", "GISelFailure",
+                               F.getSubprogram(), &F.getEntryBlock());
+    R << "unable to lower function: " << ore::NV("Prototype", F.getType());
+    reportTranslationError(*MF, *TPC, *ORE, R);
+    return false;
+  }
+
   // Lower the actual args into this basic block.
   SmallVector<ArrayRef<Register>, 8> VRegArgs;
   for (const Argument &Arg: F.args()) {
     if (DL->getTypeStoreSize(Arg.getType()).isZero())
       continue; // Don't handle zero sized types.
     ArrayRef<Register> VRegs = getOrCreateVRegs(Arg);
     VRegArgs.push_back(VRegs);
 

llvm/lib/Target/AArch64/GISel/AArch64CallLowering.h

 class AArch64CallLowering: public CallLowering {
 public:
   AArch64CallLowering(const AArch64TargetLowering &TLI);
 
   bool lowerReturn(MachineIRBuilder &MIRBuilder, const Value *Val,
                    ArrayRef<Register> VRegs,
                    Register SwiftErrorVReg) const override;
 
+  bool fallBackToDAGISel(const Function &F) const override;
+
   bool lowerFormalArguments(MachineIRBuilder &MIRBuilder, const Function &F,
                             ArrayRef<ArrayRef<Register>> VRegs) const override;
 
   bool lowerCall(MachineIRBuilder &MIRBuilder,
                  CallLoweringInfo &Info) const override;
 
   /// Returns true if the call can be lowered as a tail call.
   bool

llvm/lib/Target/AArch64/GISel/AArch64CallLowering.cpp

static void handleMustTailForwardedRegisters(MachineIRBuilder &MIRBuilder,
 
   // Add the forwards to the MachineBasicBlock and MachineFunction.
   for (const auto &F : Forwards) {
     MBB.addLiveIn(F.PReg);
     MIRBuilder.buildCopy(Register(F.VReg), Register(F.PReg));
   }
 }
 
+bool AArch64CallLowering::fallBackToDAGISel(const Function &F) const {
+  if (isa<ScalableVectorType>(F.getReturnType()))
+    return true;
+  return llvm::any_of(F.args(), [](const Argument &A) {
+    return isa<ScalableVectorType>(A.getType());
+  });
+}
+
 bool AArch64CallLowering::lowerFormalArguments(
     MachineIRBuilder &MIRBuilder, const Function &F,
     ArrayRef<ArrayRef<Register>> VRegs) const {
   MachineFunction &MF = MIRBuilder.getMF();
   MachineBasicBlock &MBB = MIRBuilder.getMBB();
   MachineRegisterInfo &MRI = MF.getRegInfo();
   auto &DL = F.getParent()->getDataLayout();
 
   SmallVector<ArgInfo, 8> SplitArgs;
   unsigned i = 0;
   for (auto &Arg : F.args()) {
-    if (isa<ScalableVectorType>(Arg.getType()))
-      return false;
-
     if (DL.getTypeStoreSize(Arg.getType()).isZero())
       continue;
 
     ArgInfo OrigArg{VRegs[i], Arg.getType()};
     setArgFlags(OrigArg, i + AttributeList::FirstArgIndex, DL, F);
 
     splitToValueTypes(OrigArg, SplitArgs, DL, MRI, F.getCallingConv());
     ++i;

llvm/test/CodeGen/AArch64/GlobalISel/arm64-fallback.ll

 ; FALLBACK-WITH-REPORT-ERR: warning: Instruction selection used fallback path for direct_mem
 ; FALLBACK-WITH-REPORT-OUT-LABEL: direct_mem
 define void @direct_mem(i32 %x, i32 %y) {
 entry:
   tail call void asm sideeffect "", "imr,imr,~{memory}"(i32 %x, i32 %y)
   ret void
 }
 
-; FALLBACK-WITH-REPORT-ERR: remark: <unknown>:0:0: unable to lower arguments{{.*}}scalable_arg
+; FALLBACK-WITH-REPORT-ERR: remark: <unknown>:0:0: unable to lower function{{.*}}scalable_arg
 ; FALLBACK-WITH-REPORT-OUT-LABEL: scalable_arg
 define <vscale x 16 x i8> @scalable_arg(<vscale x 16 x i1> %pred, i8* %addr) #1 {
   %res = call <vscale x 16 x i8> @llvm.aarch64.sve.ld1.nxv16i8(<vscale x 16 x i1> %pred, i8* %addr)
   ret <vscale x 16 x i8> %res
 }
 
-; FALLBACK-WITH-REPORT-ERR: remark: <unknown>:0:0: unable to translate instruction{{.*}}scalable_call
-; FALLBACK-WITH-REPORT-OUT-LABEL: scalable_call
-define <vscale x 16 x i8> @scalable_call(i8* %addr) #1 {
+; FALLBACK-WITH-REPORT-ERR: remark: <unknown>:0:0: unable to lower function{{.*}}scalable_ret
+; FALLBACK-WITH-REPORT-OUT-LABEL: scalable_ret
+define <vscale x 16 x i8> @scalable_ret(i8* %addr) #1 {
   %pred = call <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 0)
   %res = call <vscale x 16 x i8> @llvm.aarch64.sve.ld1.nxv16i8(<vscale x 16 x i1> %pred, i8* %addr)
   ret <vscale x 16 x i8> %res
 }
 
+; FALLBACK-WITH-REPORT-ERR: remark: <unknown>:0:0: unable to translate instruction{{.*}}scalable_call
+; FALLBACK-WITH-REPORT-OUT-LABEL: scalable_call
+define i8 @scalable_call(i8* %addr) #1 {
+  %pred = call <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 0)
+  %vec = call <vscale x 16 x i8> @llvm.aarch64.sve.ld1.nxv16i8(<vscale x 16 x i1> %pred, i8* %addr)
+  %res = extractelement <vscale x 16 x i8> %vec, i32 0
+  ret i8 %res
+}
+
 attributes #1 = { "target-features"="+sve" }
 
 declare <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 %pattern)
 declare <vscale x 16 x i8> @llvm.aarch64.sve.ld1.nxv16i8(<vscale x 16 x i1>, i8*)