[SelectionDAG][RISCV] Support @llvm.vscale.i64() on 32-bit targets.

[efriedma]

Summary

Not really useful on its own, but D105673 depends on it.

[craig.topper]

LGTM

llvm/test/CodeGen/RISCV/rvv/rvv-vscale.i64.ll

55

vlenb returns a multiple of 8 so we should be able to replace this srli 3 + slli 6 with just slli 3. But I think we have an ordering issue in visiting that prevents it from being simplified. That shouldn't block this patch though.

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

#endif
   case ISD::ROTR:
     ExpandIntRes_Rotate(N, Lo, Hi);
     break;
 
   case ISD::FSHL:
   case ISD::FSHR:
     ExpandIntRes_FunnelShift(N, Lo, Hi);
     break;
+
+  case ISD::VSCALE:
+    ExpandIntRes_VSCALE(N, Lo, Hi);
+    break;
   }
 
   // If Lo/Hi is null, the sub-method took care of registering results etc.
   if (Lo.getNode())
     SetExpandedInteger(SDValue(N, ResNo), Lo, Hi);
 }
 
 /// Lower an atomic node to the appropriate builtin call.

 void DAGTypeLegalizer::ExpandIntRes_FunnelShift(SDNode *N,
                                                 SDValue &Lo, SDValue &Hi) {
   // Lower the funnel shift to shifts and ORs which can be expanded.
   SDValue Res;
   TLI.expandFunnelShift(N, Res, DAG);
   SplitInteger(Res, Lo, Hi);
 }
 
+void DAGTypeLegalizer::ExpandIntRes_VSCALE(SDNode *N, SDValue &Lo,
+                                           SDValue &Hi) {
+  EVT VT = N->getValueType(0);
+  EVT HalfVT =
+      EVT::getIntegerVT(*DAG.getContext(), N->getValueSizeInBits(0) / 2);
+  SDLoc dl(N);
+
+  // We assume VSCALE(1) fits into a legal integer.
+  APInt One(HalfVT.getSizeInBits(), 1);
+  SDValue VScaleBase = DAG.getVScale(dl, HalfVT, One);
+  VScaleBase = DAG.getNode(ISD::ZERO_EXTEND, dl, VT, VScaleBase);
+  SDValue Res = DAG.getNode(ISD::MUL, dl, VT, VScaleBase, N->getOperand(0));
+  SplitInteger(Res, Lo, Hi);
+}
+
 //===----------------------------------------------------------------------===//
 //  Integer Operand Expansion
 //===----------------------------------------------------------------------===//
 
 /// ExpandIntegerOperand - This method is called when the specified operand of
 /// the specified node is found to need expansion.  At this point, all of the
 /// result types of the node are known to be legal, but other operands of the
 /// node may need promotion or expansion as well as the specified one.

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h

private:
   void ExpandIntRes_DIVFIX            (SDNode *N, SDValue &Lo, SDValue &Hi);
 
   void ExpandIntRes_ATOMIC_LOAD       (SDNode *N, SDValue &Lo, SDValue &Hi);
   void ExpandIntRes_VECREDUCE         (SDNode *N, SDValue &Lo, SDValue &Hi);
 
   void ExpandIntRes_Rotate            (SDNode *N, SDValue &Lo, SDValue &Hi);
   void ExpandIntRes_FunnelShift       (SDNode *N, SDValue &Lo, SDValue &Hi);
 
+  void ExpandIntRes_VSCALE            (SDNode *N, SDValue &Lo, SDValue &Hi);
+
   void ExpandShiftByConstant(SDNode *N, const APInt &Amt,
                              SDValue &Lo, SDValue &Hi);
   bool ExpandShiftWithKnownAmountBit(SDNode *N, SDValue &Lo, SDValue &Hi);
   bool ExpandShiftWithUnknownAmountBit(SDNode *N, SDValue &Lo, SDValue &Hi);
 
   // Integer Operand Expansion.
   bool ExpandIntegerOperand(SDNode *N, unsigned OpNo);
   SDValue ExpandIntOp_BR_CC(SDNode *N);

llvm/test/CodeGen/RISCV/rvv/rvv-vscale.i64.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple riscv64 -mattr=+m,+experimental-v < %s \
-; RUN:    | FileCheck %s
+; RUN:    | FileCheck %s -check-prefix=RV64
+; RUN: llc -mtriple riscv32 -mattr=+m,+experimental-v < %s \
+; RUN:    | FileCheck %s -check-prefix=RV32
+
 
 define i64 @vscale_zero() nounwind {
-; CHECK-LABEL: vscale_zero:
-; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    mv a0, zero
-; CHECK-NEXT:    ret
+; RV64-LABEL: vscale_zero:
+; RV64:       # %bb.0: # %entry
+; RV64-NEXT:    mv a0, zero
+; RV64-NEXT:    ret
+;
+; RV32-LABEL: vscale_zero:
+; RV32:       # %bb.0: # %entry
+; RV32-NEXT:    mv a0, zero
+; RV32-NEXT:    mv a1, zero
+; RV32-NEXT:    ret
 entry:
   %0 = call i64 @llvm.vscale.i64()
   %1 = mul i64 %0, 0
   ret i64 %1
 }
 
 define i64 @vscale_one() nounwind {
-; CHECK-LABEL: vscale_one:
-; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    csrr a0, vlenb
-; CHECK-NEXT:    srli a0, a0, 3
-; CHECK-NEXT:    ret
+; RV64-LABEL: vscale_one:
+; RV64:       # %bb.0: # %entry
+; RV64-NEXT:    csrr a0, vlenb
+; RV64-NEXT:    srli a0, a0, 3
+; RV64-NEXT:    ret
+;
+; RV32-LABEL: vscale_one:
+; RV32:       # %bb.0: # %entry
+; RV32-NEXT:    csrr a0, vlenb
+; RV32-NEXT:    srli a0, a0, 3
+; RV32-NEXT:    mv a1, zero
+; RV32-NEXT:    ret
 entry:
   %0 = call i64 @llvm.vscale.i64()
   %1 = mul i64 %0, 1
   ret i64 %1
 }
 
 define i64 @vscale_uimmpow2xlen() nounwind {
-; CHECK-LABEL: vscale_uimmpow2xlen:
-; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    csrr a0, vlenb
-; CHECK-NEXT:    slli a0, a0, 3
-; CHECK-NEXT:    ret
+; RV64-LABEL: vscale_uimmpow2xlen:
+; RV64:       # %bb.0: # %entry
+; RV64-NEXT:    csrr a0, vlenb
+; RV64-NEXT:    slli a0, a0, 3
+; RV64-NEXT:    ret
+;
+; RV32-LABEL: vscale_uimmpow2xlen:
+; RV32:       # %bb.0: # %entry
+; RV32-NEXT:    csrr a0, vlenb
+; RV32-NEXT:    srli a1, a0, 29
+; RV32-NEXT:    srli a0, a0, 3
+; RV32-NEXT:    slli a0, a0, 6
+; RV32-NEXT:    ret
 entry:
   %0 = call i64 @llvm.vscale.i64()
   %1 = mul i64 %0, 64
   ret i64 %1
 }
 
 define i64 @vscale_non_pow2() nounwind {
-; CHECK-LABEL: vscale_non_pow2:
-; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    csrr a0, vlenb
-; CHECK-NEXT:    srli a0, a0, 3
-; CHECK-NEXT:    addi a1, zero, 24
-; CHECK-NEXT:    mul a0, a0, a1
-; CHECK-NEXT:    ret
+; RV64-LABEL: vscale_non_pow2:
+; RV64:       # %bb.0: # %entry
+; RV64-NEXT:    csrr a0, vlenb
+; RV64-NEXT:    srli a0, a0, 3
+; RV64-NEXT:    addi a1, zero, 24
+; RV64-NEXT:    mul a0, a0, a1
+; RV64-NEXT:    ret
+;
+; RV32-LABEL: vscale_non_pow2:
+; RV32:       # %bb.0: # %entry
+; RV32-NEXT:    csrr a0, vlenb
+; RV32-NEXT:    srli a1, a0, 3
+; RV32-NEXT:    addi a2, zero, 24
+; RV32-NEXT:    mul a0, a1, a2
+; RV32-NEXT:    mulhu a1, a1, a2
+; RV32-NEXT:    ret
 entry:
   %0 = call i64 @llvm.vscale.i64()
   %1 = mul i64 %0, 24
   ret i64 %1
 }
 
 declare i64 @llvm.vscale.i64()