[RISCV][GISEL] Add support for scalable vector types in lowerReturnVal (#71587)

Scalable vector types from LLVM IR are lowered into physical vector
registers in MIR based on calling convention for return instructions.

Author: michaelmaitland

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

@@ -359,7 +359,7 @@ bool IRTranslator::translateCompare(const User &U,
 bool IRTranslator::translateRet(const User &U, MachineIRBuilder &MIRBuilder) {
   const ReturnInst &RI = cast<ReturnInst>(U);
   const Value *Ret = RI.getReturnValue();
-  if (Ret && DL->getTypeStoreSize(Ret->getType()) == 0)
+  if (Ret && DL->getTypeStoreSize(Ret->getType()).isZero())
     Ret = nullptr;
 
   ArrayRef<Register> VRegs;

llvm/lib/CodeGen/MachineVerifier.cpp

@@ -1966,6 +1966,9 @@ void MachineVerifier::visitMachineInstrBefore(const MachineInstr *MI) {
     if (SrcReg.isVirtual() && DstReg.isPhysical() && SrcSize.isScalable() &&
         !DstSize.isScalable())
       break;
+    if (SrcReg.isVirtual() && DstReg.isPhysical() && SrcSize.isScalable() &&
+        !DstSize.isScalable())
+      break;
 
     if (SrcSize.isNonZero() && DstSize.isNonZero() && SrcSize != DstSize) {
       if (!DstOp.getSubReg() && !SrcOp.getSubReg()) {

llvm/lib/Target/RISCV/GISel/RISCVCallLowering.cpp

@@ -347,7 +347,10 @@ static bool isSupportedArgumentType(Type *T, const RISCVSubtarget &Subtarget,
 }
 
 // TODO: Only integer, pointer and aggregate types are supported now.
-static bool isSupportedReturnType(Type *T, const RISCVSubtarget &Subtarget) {
+// TODO: Remove IsLowerRetVal argument by adding support for vectors in
+// lowerCall.
+static bool isSupportedReturnType(Type *T, const RISCVSubtarget &Subtarget,
+                                  bool IsLowerRetVal = false) {
   // TODO: Integers larger than 2*XLen are passed indirectly which is not
   // supported yet.
   if (T->isIntegerTy())
@@ -368,6 +371,11 @@ static bool isSupportedReturnType(Type *T, const RISCVSubtarget &Subtarget) {
     return true;
   }
 
+  if (IsLowerRetVal && T->isVectorTy() && Subtarget.hasVInstructions() &&
+      T->isScalableTy() &&
+      isLegalElementTypeForRVV(T->getScalarType(), Subtarget))
+    return true;
+
   return false;
 }
 
@@ -380,7 +388,7 @@ bool RISCVCallLowering::lowerReturnVal(MachineIRBuilder &MIRBuilder,
 
   const RISCVSubtarget &Subtarget =
       MIRBuilder.getMF().getSubtarget<RISCVSubtarget>();
-  if (!isSupportedReturnType(Val->getType(), Subtarget))
+  if (!isSupportedReturnType(Val->getType(), Subtarget, /*IsLowerRetVal=*/true))
     return false;
 
   MachineFunction &MF = MIRBuilder.getMF();

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -19638,12 +19638,15 @@ unsigned RISCVTargetLowering::getCustomCtpopCost(EVT VT,
 }
 
 bool RISCVTargetLowering::fallBackToDAGISel(const Instruction &Inst) const {
-  // We don't support scalable vectors in GISel.
+  // At the moment, the only scalable instruction GISel knows how to lower is
+  // ret with scalable argument.
+
   if (Inst.getType()->isScalableTy())
     return true;
 
   for (unsigned i = 0; i < Inst.getNumOperands(); ++i)
-    if (Inst.getOperand(i)->getType()->isScalableTy())
+    if (Inst.getOperand(i)->getType()->isScalableTy() &&
+        !isa<ReturnInst>(&Inst))
       return true;
 
   if (const AllocaInst *AI = dyn_cast<AllocaInst>(&Inst)) {

llvm/test/CodeGen/RISCV/GlobalISel/irtranslator/vec-ret-bf16-err.ll

@@ -0,0 +1,14 @@
+; RUN: not --crash llc -mtriple=riscv32 -mattr=+v -global-isel -stop-after=irtranslator \
+; RUN:   -verify-machineinstrs < %s 2>&1 | FileCheck %s
+; RUN: not --crash llc -mtriple=riscv64 -mattr=+v -global-isel -stop-after=irtranslator \
+; RUN:   -verify-machineinstrs < %s 2>&1 | FileCheck %s
+
+; The purpose of this test is to show that the compiler throws an error when
+; there is no support for bf16 vectors. If the compiler did not throw an error,
+; then it will try to scalarize the argument to an s32, which may drop elements.
+define <vscale x 1 x bfloat> @test_ret_nxv1bf16() {
+entry:
+  ret <vscale x 1 x bfloat> undef
+}
+
+; CHECK: LLVM ERROR: unable to translate instruction: ret (in function: test_ret_nxv1bf16)

llvm/test/CodeGen/RISCV/GlobalISel/irtranslator/vec-ret-f16-err.ll

@@ -0,0 +1,14 @@
+; RUN: not --crash llc -mtriple=riscv32 -mattr=+v -global-isel -stop-after=irtranslator \
+; RUN:   -verify-machineinstrs < %s 2>&1 | FileCheck %s
+; RUN: not --crash llc -mtriple=riscv64 -mattr=+v -global-isel -stop-after=irtranslator \
+; RUN:   -verify-machineinstrs < %s 2>&1 | FileCheck %s
+
+; The purpose of this test is to show that the compiler throws an error when
+; there is no support for f16 vectors. If the compiler did not throw an error,
+; then it will try to scalarize the argument to an s32, which may drop elements.
+define <vscale x 1 x half> @test_ret_nxv1f16() {
+entry:
+  ret <vscale x 1 x half> undef
+}
+
+; CHECK: LLVM ERROR: unable to translate instruction: ret (in function: test_ret_nxv1f16)