[ARM][MVE] Add intrinsics for immediate shifts. (reland)

This adds the family of `vshlq_n` and `vshrq_n` ACLE intrinsics, which
shift every lane of a vector left or right by a compile-time
immediate. They mostly work by expanding to the IR `shl`, `lshr` and
`ashr` operations, with their second operand being a vector splat of
the immediate.

There's a fiddly special case, though. ACLE specifies that the
immediate in `vshrq_n` can take values up to //and including// the bit
size of the vector lane. But LLVM IR thinks that shifting right by the
full size of the lane is UB, and feels free to replace the `lshr` with
an `undef` half way through the optimization pipeline. Hence, to keep
this legal in source code, I have to detect it at codegen time.
Logical (unsigned) right shifts by the element size are handled by
simply emitting the zero vector; arithmetic ones are converted into a
shift of one bit less, which will always give the same output.

In order to do that check, I also had to enhance the tablegen
MveEmitter so that it can cope with converting a builtin function's
operand into a bare integer to pass to a code-generating subfunction.
Previously the only bare integers it knew how to handle were flags
generated from within `arm_mve.td`.

Reviewers: dmgreen, miyuki, MarkMurrayARM, ostannard

Reviewed By: dmgreen, MarkMurrayARM

Subscribers: echristo, hokein, rdhindsa, kristof.beyls, hiraditya, cfe-commits, llvm-commits

Tags: #clang, #llvm

Differential Revision: https://reviews.llvm.org/D71065

Author: statham-arm

clang/include/clang/Basic/arm_mve.td

@@ -609,6 +609,33 @@ defm vstrhq: scatter_offset_both<!listconcat(T.All16, T.Int32), u16, 1>;
 defm vstrwq: scatter_offset_both<T.All32, u32, 2>;
 defm vstrdq: scatter_offset_both<T.Int64, u64, 3>;
 
+multiclass PredicatedImmediateVectorShift<
+    Immediate immtype, string predIntrName, list<dag> unsignedFlag = []> {
+  foreach predIntr = [IRInt<predIntrName, [Vector, Predicate]>] in {
+    def _m_n: Intrinsic<Vector, (args Vector:$inactive, Vector:$v,
+                                      immtype:$sh, Predicate:$pred),
+       !con((predIntr $v, $sh), !dag(predIntr, unsignedFlag, ?),
+            (predIntr $pred, $inactive))>;
+    def _x_n: Intrinsic<Vector, (args Vector:$v, immtype:$sh,
+                                      Predicate:$pred),
+       !con((predIntr $v, $sh), !dag(predIntr, unsignedFlag, ?),
+            (predIntr $pred, (undef Vector)))>;
+  }
+}
+
+let params = T.Int in {
+  def vshlq_n: Intrinsic<Vector, (args Vector:$v, imm_0toNm1:$sh),
+                         (shl $v, (splat (Scalar $sh)))>;
+  defm vshlq: PredicatedImmediateVectorShift<imm_0toNm1, "shl_imm_predicated">;
+
+  let pnt = PNT_NType in {
+    def vshrq_n: Intrinsic<Vector, (args Vector:$v, imm_1toN:$sh),
+                            (immshr $v, $sh, (unsignedflag Scalar))>;
+    defm vshrq: PredicatedImmediateVectorShift<imm_1toN, "shr_imm_predicated",
+                                     [(unsignedflag Scalar)]>;
+  }
+}
+
 // Base class for the scalar shift intrinsics.
 class ScalarShift<Type argtype, dag shiftCountArg, dag shiftCodeGen>:
   Intrinsic<argtype, !con((args argtype:$value), shiftCountArg), shiftCodeGen> {

clang/include/clang/Basic/arm_mve_defs.td

@@ -66,6 +66,10 @@ def xor: IRBuilder<"CreateXor">;
 def sub: IRBuilder<"CreateSub">;
 def shl: IRBuilder<"CreateShl">;
 def lshr: IRBuilder<"CreateLShr">;
+def immshr: CGHelperFn<"MVEImmediateShr"> {
+  let special_params = [IRBuilderIntParam<1, "unsigned">,
+                        IRBuilderIntParam<2, "bool">];
+}
 def fadd: IRBuilder<"CreateFAdd">;
 def fmul: IRBuilder<"CreateFMul">;
 def fsub: IRBuilder<"CreateFSub">;
@@ -318,8 +322,8 @@ def imm_simd_vmvn : Immediate<u32, IB_UEltValue> {
 //
 // imm_0toNm1 is the same but with the range offset by 1, i.e. 0 to N-1
 // inclusive.
-def imm_1toN : Immediate<u32, IB_EltBit<1>>;
-def imm_0toNm1 : Immediate<u32, IB_EltBit<0>>;
+def imm_1toN : Immediate<sint, IB_EltBit<1>>;
+def imm_0toNm1 : Immediate<sint, IB_EltBit<0>>;
 
 // imm_lane has to be the index of a vector lane in the main vector type, i.e
 // it can range from 0 to (128 / size of scalar)-1 inclusive. (e.g. vgetq_lane)

clang/lib/CodeGen/CGBuiltin.cpp

@@ -6916,13 +6916,43 @@ Value *CodeGenFunction::EmitARMBuiltinExpr(unsigned BuiltinID,
   }
 }
 
+template<typename Integer>
+static Integer GetIntegerConstantValue(const Expr *E, ASTContext &Context) {
+  llvm::APSInt IntVal;
+  bool IsConst = E->isIntegerConstantExpr(IntVal, Context);
+  assert(IsConst && "Sema should have checked this was a constant");
+  (void)IsConst;
+  return IntVal.getExtValue();
+}
+
 static llvm::Value *SignOrZeroExtend(CGBuilderTy &Builder, llvm::Value *V,
                                      llvm::Type *T, bool Unsigned) {
   // Helper function called by Tablegen-constructed ARM MVE builtin codegen,
   // which finds it convenient to specify signed/unsigned as a boolean flag.
   return Unsigned ? Builder.CreateZExt(V, T) : Builder.CreateSExt(V, T);
 }
 
+static llvm::Value *MVEImmediateShr(CGBuilderTy &Builder, llvm::Value *V,
+                                    uint32_t Shift, bool Unsigned) {
+  // MVE helper function for integer shift right. This must handle signed vs
+  // unsigned, and also deal specially with the case where the shift count is
+  // equal to the lane size. In LLVM IR, an LShr with that parameter would be
+  // undefined behavior, but in MVE it's legal, so we must convert it to code
+  // that is not undefined in IR.
+  unsigned LaneBits =
+      V->getType()->getVectorElementType()->getPrimitiveSizeInBits();
+  if (Shift == LaneBits) {
+    // An unsigned shift of the full lane size always generates zero, so we can
+    // simply emit a zero vector. A signed shift of the full lane size does the
+    // same thing as shifting by one bit fewer.
+    if (Unsigned)
+      return llvm::Constant::getNullValue(V->getType());
+    else
+      --Shift;
+  }
+  return Unsigned ? Builder.CreateLShr(V, Shift) : Builder.CreateAShr(V, Shift);
+}
+
 static llvm::Value *ARMMVEVectorSplat(CGBuilderTy &Builder, llvm::Value *V) {
   // MVE-specific helper function for a vector splat, which infers the element
   // count of the output vector by knowing that MVE vectors are all 128 bits