[AArch64][SVE] Add patterns for logical immediate operations.

[dancgr]

Summary

Add pattern matching for the following SVE logical vector and immediate instructions:

• and/bic, orr/orn, eor/eon.

[efriedma]

llvm/include/llvm/IR/IntrinsicsAArch64.td

1064

I don't really like anyint here... can we restrict the type somehow?

Actually, I'm not sure how this even works with your testcases; how is LLVM computing the type of llvm.aarch64.sve.orr.imm.nxv16i8? You aren't specifying the type of the integer.

[dancgr]

llvm/include/llvm/IR/IntrinsicsAArch64.td

1064

We can split this into two intrinsics, one i32 and one i64, both will be almost exactly the same with only the immediate type difference. I don't see a problem in that, I was just trying to reuse some code to avoid unnecessary complexity. But I'm not opposed to having two separate ones either.

In the following test case I'm explicitly setting the input as i64 on the declare statement. If I were to put i32 there, It would fail to find a matching pattern.

define <vscale x 16 x i8> @orr_i8(<vscale x 16 x i8> %a) {
; CHECK-LABEL: orr_i8:
; CHECK: orr z0.b, z0.b, #0xf
; CHECK-NEXT: ret
  %res = call <vscale x 16 x i8> @llvm.aarch64.sve.orr.imm.nxv16i8(<vscale x 16 x i8> %a, 
                                                                   i64 15)
  ret <vscale x 16 x i8> %res
}

declare <vscale x 16 x i8> @llvm.aarch64.sve.orr.imm.nxv16i8(<vscale x 16 x i8>, i64)

[efriedma]

llvm/include/llvm/IR/IntrinsicsAArch64.td

1064

Why do we need both an i32 and an i64 variant? The i64 variant seems to cover all the relevant cases (and you haven't added any tests for the i32 variant).

I figured out the source of my confusion about the names of the intrinsics; apparently IR autoupgrade is "fixing" the name of llvm.aarch64.sve.orr.imm.nxv16i8 to refer to the actual name, llvm.aarch64.sve.orr.imm.nxv16i8.i64.

[dancgr]

llvm/include/llvm/IR/IntrinsicsAArch64.td

1064

Oh, I get it now. The AdvSIMD_1VectorArg_Imm_Intrinsic was introduced previously for the add/sub/sqadd imm. instrinsics (which use i32). And then I changed it to anyint to be able to use the same class for and/orr/eor imm., because those use i64.

[dancgr]

llvm/include/llvm/IR/IntrinsicsAArch64.td

1064

But I don't mind changing that. I am just trying to understand the reasoning so I can apply it to my future patches.

[efriedma]

llvm/include/llvm/IR/IntrinsicsAArch64.td

1064

A new "class" doesn't have any associated cost; it's just a pattern for defining an intrinsic. if you need a new one, just add it.

[dancgr]

Add specific intrinsic for i64.

[efriedma]

LGTM

[sdesmalen]

Thanks @dancgr, LGTM!

There is no harm in adding these intrinsics, but it is worth pointing out that our downstream compiler does not have explicit intrinsics for the immediate forms because:

• There are no specific C/C++ intrinsics that require the operand to be an immediate. That means using the immediate form is purely an optimisation of the general ((unpredicated) vector, vector) case.
• We can use a pattern to match e.g. the AArch64dup (SVELogicalImm32 i64:$imm) and emit the immediate-form of the instruction.

[dancgr]

I will be merging this patch then, and I will submit a short patch for the AArch64dup (SVELogicalImm32 i64:$imm) patterns and the equivalent ones for the add/sub instructions.

llvm/include/llvm/IR/IntrinsicsAArch64.td

: Intrinsic<[llvm_anyvector_ty],
                   llvm_anyvector_ty,
                   llvm_i64_ty
                 ],
                 [IntrReadMem, IntrArgMemOnly]>;
 
 class AdvSIMD_1VectorArg_Imm_Intrinsic
     : Intrinsic<[llvm_anyvector_ty],
                 [LLVMMatchType<0>,
-                 llvm_i32_ty],
+                 llvm_anyint_ty],
                 [IntrNoMem, ImmArg<1>]>;
 
 //
 // Loads
 //
 
 def int_aarch64_sve_ldnt1 : AdvSIMD_1Vec_PredLoad_Intrinsic;
 

 def int_aarch64_sve_orn         : AdvSIMD_Pred2VectorArg_Intrinsic;
 def int_aarch64_sve_nor         : AdvSIMD_Pred2VectorArg_Intrinsic;
 def int_aarch64_sve_nand        : AdvSIMD_Pred2VectorArg_Intrinsic;
 def int_aarch64_sve_orrs        : AdvSIMD_Pred2VectorArg_Intrinsic;
 def int_aarch64_sve_orns        : AdvSIMD_Pred2VectorArg_Intrinsic;
 def int_aarch64_sve_nors        : AdvSIMD_Pred2VectorArg_Intrinsic;
 def int_aarch64_sve_nands       : AdvSIMD_Pred2VectorArg_Intrinsic;
 
+def int_aarch64_sve_orr_imm  : AdvSIMD_1VectorArg_Imm_Intrinsic;
+def int_aarch64_sve_eor_imm  : AdvSIMD_1VectorArg_Imm_Intrinsic;
+def int_aarch64_sve_and_imm  : AdvSIMD_1VectorArg_Imm_Intrinsic;
+
 //
 // Conversion
 //
 
 def int_aarch64_sve_sxtb : AdvSIMD_Merged1VectorArg_Intrinsic;
 def int_aarch64_sve_sxth : AdvSIMD_Merged1VectorArg_Intrinsic;
 def int_aarch64_sve_sxtw : AdvSIMD_Merged1VectorArg_Intrinsic;
 def int_aarch64_sve_uxtb : AdvSIMD_Merged1VectorArg_Intrinsic;

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

bool SelectDupZeroOrUndef(SDValue N) {
     return false;
   }
 
   template<MVT::SimpleValueType VT>
   bool SelectSVEAddSubImm(SDValue N, SDValue &Imm, SDValue &Shift) {
     return SelectSVEAddSubImm(N, VT, Imm, Shift);
   }
 
+  template<MVT::SimpleValueType VT>
+  bool SelectSVELogicalImm(SDValue N, SDValue &Imm) {
+    return SelectSVELogicalImm(N, VT, Imm);
+  }
+
   /// Form sequences of consecutive 64/128-bit registers for use in NEON
   /// instructions making use of a vector-list (e.g. ldN, tbl). Vecs must have
   /// between 1 and 4 elements. If it contains a single element that is returned
   /// unchanged; otherwise a REG_SEQUENCE value is returned.
   SDValue createDTuple(ArrayRef<SDValue> Vecs);
   SDValue createQTuple(ArrayRef<SDValue> Vecs);
 
   /// Generic helper for the createDTuple/createQTuple
bool SelectCVTFixedPosOperand(SDValue N, SDValue &FixedPos) {
     return SelectCVTFixedPosOperand(N, FixedPos, RegWidth);
   }
 
   bool SelectCVTFixedPosOperand(SDValue N, SDValue &FixedPos, unsigned Width);
 
   bool SelectCMP_SWAP(SDNode *N);
 
   bool SelectSVEAddSubImm(SDValue N, MVT VT, SDValue &Imm, SDValue &Shift);
+
+  bool SelectSVELogicalImm(SDValue N, MVT VT, SDValue &Imm);
 };
 } // end anonymous namespace
 
 /// isIntImmediate - This method tests to see if the node is a constant
 /// operand. If so Imm will receive the 32-bit value.
 static bool isIntImmediate(const SDNode *N, uint64_t &Imm) {
   if (const ConstantSDNode *C = dyn_cast<const ConstantSDNode>(N)) {
     Imm = C->getZExtValue();
if (auto CNode = dyn_cast<ConstantSDNode>(N)) {
     default:
       break;
     }
   }
 
   return false;
 }
 
+bool AArch64DAGToDAGISel::SelectSVELogicalImm(SDValue N, MVT VT, SDValue &Imm) {
+  if (auto CNode = dyn_cast<ConstantSDNode>(N)) {
+    uint64_t ImmVal = CNode->getZExtValue();
+    SDLoc DL(N);
+
+    // Shift mask depending on type size.
+    switch (VT.SimpleTy) {
+      case MVT::i8:
+        ImmVal &= 0xFF;
+        ImmVal |= (ImmVal << 8);
+      case MVT::i16:
+        ImmVal &= 0xFFFF;
+        ImmVal |= (ImmVal << 16);
+      case MVT::i32:
+        ImmVal &= 0xFFFFFFFF;
+        ImmVal |= (ImmVal << 32);
+        break;
+      default:
+        break;
+    }
+
+    uint64_t encoding;
+    if (AArch64_AM::processLogicalImmediate(ImmVal, 64, encoding)) {
+      Imm = CurDAG->getTargetConstant(encoding, DL, MVT::i64);
+      return true;
+    }
+  }
+  return false;
+}
 
 bool AArch64DAGToDAGISel::trySelectStackSlotTagP(SDNode *N) {
   // tagp(FrameIndex, IRGstack, tag_offset):
   // since the offset between FrameIndex and IRGstack is a compile-time
   // constant, this can be lowered to a single ADDG instruction.
   if (!(isa<FrameIndexSDNode>(N->getOperand(1)))) {
     return false;
   }

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

let Predicates = [HasSVE] in {
   defm SMAXV_VPZ : sve_int_reduce_1<0b000, "smaxv", AArch64smaxv_pred>;
   defm UMAXV_VPZ : sve_int_reduce_1<0b001, "umaxv", AArch64umaxv_pred>;
   defm SMINV_VPZ : sve_int_reduce_1<0b010, "sminv", AArch64sminv_pred>;
   defm UMINV_VPZ : sve_int_reduce_1<0b011, "uminv", AArch64uminv_pred>;
   defm ORV_VPZ   : sve_int_reduce_2<0b000, "orv", AArch64orv_pred>;
   defm EORV_VPZ  : sve_int_reduce_2<0b001, "eorv", AArch64eorv_pred>;
   defm ANDV_VPZ  : sve_int_reduce_2<0b010, "andv", AArch64andv_pred>;
 
-  defm ORR_ZI : sve_int_log_imm<0b00, "orr", "orn">;
-  defm EOR_ZI : sve_int_log_imm<0b01, "eor", "eon">;
-  defm AND_ZI : sve_int_log_imm<0b10, "and", "bic">;
+  defm ORR_ZI : sve_int_log_imm<0b00, "orr", "orn", int_aarch64_sve_orr_imm>;
+  defm EOR_ZI : sve_int_log_imm<0b01, "eor", "eon", int_aarch64_sve_eor_imm>;
+  defm AND_ZI : sve_int_log_imm<0b10, "and", "bic", int_aarch64_sve_and_imm>;
 
   defm SMAX_ZI   : sve_int_arith_imm1<0b00, "smax", simm8>;
   defm SMIN_ZI   : sve_int_arith_imm1<0b10, "smin", simm8>;
   defm UMAX_ZI   : sve_int_arith_imm1<0b01, "umax", imm0_255>;
   defm UMIN_ZI   : sve_int_arith_imm1<0b11, "umin", imm0_255>;
 
   defm MUL_ZI    : sve_int_arith_imm2<"mul">;
   defm MUL_ZPmZ   : sve_int_bin_pred_arit_2<0b000, "mul", int_aarch64_sve_mul>;

llvm/lib/Target/AArch64/SVEInstrFormats.td

def addsub_imm8_opt_lsl_i64 : imm8_opt_lsl<64, "uint64_t", SVEAddSubImmOperand64, [{
   return AArch64_AM::isSVEAddSubImm<int64_t>(Imm);
 }]>;
 
 def SVEAddSubImm8Pat  : ComplexPattern<i32, 2, "SelectSVEAddSubImm<MVT::i8>", []>;
 def SVEAddSubImm16Pat : ComplexPattern<i32, 2, "SelectSVEAddSubImm<MVT::i16>", []>;
 def SVEAddSubImm32Pat : ComplexPattern<i32, 2, "SelectSVEAddSubImm<MVT::i32>", []>;
 def SVEAddSubImm64Pat : ComplexPattern<i32, 2, "SelectSVEAddSubImm<MVT::i64>", []>;
 
+def SVELogicalImm8Pat  : ComplexPattern<i64, 1, "SelectSVELogicalImm<MVT::i8>", []>;
+def SVELogicalImm16Pat : ComplexPattern<i64, 1, "SelectSVELogicalImm<MVT::i16>", []>;
+def SVELogicalImm32Pat : ComplexPattern<i64, 1, "SelectSVELogicalImm<MVT::i32>", []>;
+def SVELogicalImm64Pat : ComplexPattern<i64, 1, "SelectSVELogicalImm<MVT::i64>", []>;
+
+
 class SVEExactFPImm<string Suffix, string ValA, string ValB> : AsmOperandClass {
   let Name = "SVEExactFPImmOperand" # Suffix;
   let DiagnosticType = "Invalid" # Name;
   let ParserMethod = "tryParseFPImm<false>";
   let PredicateMethod = "isExactFPImm<" # ValA # ", " # ValB # ">";
   let RenderMethod = "addExactFPImmOperands<" # ValA # ", " # ValB # ">";
 }
 

 : Pat<(vtd (op vt1:$Op1)),
       (inst $Op1)>;
 
 class SVE_1_Op_Imm_OptLsl_Pat<ValueType vt, SDPatternOperator op, ZPRRegOp zprty,
                               ComplexPattern cpx, Instruction inst>
   : Pat<(vt (op (vt zprty:$Op1), (i32 (cpx i32:$imm, i32:$shift)))),
         (inst $Op1, i32:$imm, i32:$shift)>;
 
+class SVE_1_Op_Imm_Log_Pat<ValueType vt, SDPatternOperator op, ZPRRegOp zprty,
+                           ComplexPattern cpx, Instruction inst>
+  : Pat<(vt (op (vt zprty:$Op1), (i64 (cpx i64:$imm)))),
+        (inst $Op1, i64:$imm)>;
+
 class SVE_2_Op_Pat<ValueType vtd, SDPatternOperator op, ValueType vt1,
                    ValueType vt2, Instruction inst>
 : Pat<(vtd (op vt1:$Op1, vt2:$Op2)),
       (inst $Op1, $Op2)>;
 
 class SVE_2_Op_Pat_Reduce_To_Neon<ValueType vtd, SDPatternOperator op, ValueType vt1,
                    ValueType vt2, Instruction inst, SubRegIndex sub>
 : Pat<(vtd (op vt1:$Op1, vt2:$Op2)),
: I<(outs ZPR64:$Zdn), (ins ZPR64:$_Zdn, logical_imm64:$imms13),
   let Inst{4-0}   = Zdn;
 
   let Constraints = "$Zdn = $_Zdn";
   let DecoderMethod = "DecodeSVELogicalImmInstruction";
   let DestructiveInstType = Destructive;
   let ElementSize = ElementSizeNone;
 }
 
-multiclass sve_int_log_imm<bits<2> opc, string asm, string alias> {
+multiclass sve_int_log_imm<bits<2> opc, string asm, string alias, SDPatternOperator op> {
   def NAME : sve_int_log_imm<opc, asm>;
 
+  def : SVE_1_Op_Imm_Log_Pat<nxv16i8, op, ZPR8,  SVELogicalImm8Pat,  !cast<Instruction>(NAME)>;
+  def : SVE_1_Op_Imm_Log_Pat<nxv8i16, op, ZPR16, SVELogicalImm16Pat, !cast<Instruction>(NAME)>;
+  def : SVE_1_Op_Imm_Log_Pat<nxv4i32, op, ZPR32, SVELogicalImm32Pat, !cast<Instruction>(NAME)>;
+  def : SVE_1_Op_Imm_Log_Pat<nxv2i64, op, ZPR64, SVELogicalImm64Pat, !cast<Instruction>(NAME)>;
+
   def : InstAlias<asm # "\t$Zdn, $Zdn, $imm",
                   (!cast<Instruction>(NAME) ZPR8:$Zdn, sve_logical_imm8:$imm), 4>;
   def : InstAlias<asm # "\t$Zdn, $Zdn, $imm",
                   (!cast<Instruction>(NAME) ZPR16:$Zdn, sve_logical_imm16:$imm), 3>;
   def : InstAlias<asm # "\t$Zdn, $Zdn, $imm",
                   (!cast<Instruction>(NAME) ZPR32:$Zdn, sve_logical_imm32:$imm), 2>;
 
   def : InstAlias<alias # "\t$Zdn, $Zdn, $imm",

llvm/test/CodeGen/AArch64/sve-int-log-imm.ll

+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve < %s | FileCheck %s
+
+define <vscale x 16 x i8> @orr_i8(<vscale x 16 x i8> %a) {
+; CHECK-LABEL: orr_i8:
+; CHECK: orr z0.b, z0.b, #0xf
+; CHECK-NEXT: ret
+  %res = call <vscale x 16 x i8> @llvm.aarch64.sve.orr.imm.nxv16i8(<vscale x 16 x i8> %a, 
+                                                                   i64 15)
+  ret <vscale x 16 x i8> %res
+}
+
+define <vscale x 8 x i16> @orr_i16(<vscale x 8 x i16> %a) {
+; CHECK-LABEL: orr_i16:
+; CHECK: orr z0.h, z0.h, #0xfc07
+; CHECK-NEXT: ret
+  %res = call <vscale x 8 x i16> @llvm.aarch64.sve.orr.imm.nxv8i16(<vscale x 8 x i16> %a, 
+                                                                   i64 64519)
+  ret <vscale x 8 x i16> %res
+}
+
+define <vscale x 4 x i32> @orr_i32(<vscale x 4 x i32> %a) {
+; CHECK-LABEL: orr_i32:
+; CHECK: orr z0.s, z0.s, #0xffff00
+; CHECK-NEXT: ret
+  %res = call <vscale x 4 x i32> @llvm.aarch64.sve.orr.imm.nxv4i32(<vscale x 4 x i32> %a, 
+                                                                   i64 16776960)
+  ret <vscale x 4 x i32> %res
+}
+
+define <vscale x 2 x i64> @orr_i64(<vscale x 2 x i64> %a) {
+; CHECK-LABEL: orr_i64:
+; CHECK: orr z0.d, z0.d, #0xfffc000000000000
+; CHECK-NEXT: ret
+  %res = call <vscale x 2 x i64> @llvm.aarch64.sve.orr.imm.nxv2i64(<vscale x 2 x i64> %a, 
+                                                                   i64 18445618173802708992)
+  ret <vscale x 2 x i64> %res
+}
+
+define <vscale x 16 x i8> @eor_i8(<vscale x 16 x i8> %a) {
+; CHECK-LABEL: eor_i8:
+; CHECK: eor z0.b, z0.b, #0xf
+; CHECK-NEXT: ret
+  %res = call <vscale x 16 x i8> @llvm.aarch64.sve.eor.imm.nxv16i8(<vscale x 16 x i8> %a, 
+                                                                   i64 15)
+  ret <vscale x 16 x i8> %res
+}
+
+define <vscale x 8 x i16> @eor_i16(<vscale x 8 x i16> %a) {
+; CHECK-LABEL: eor_i16:
+; CHECK: eor z0.h, z0.h, #0xfc07
+; CHECK-NEXT: ret
+  %res = call <vscale x 8 x i16> @llvm.aarch64.sve.eor.imm.nxv8i16(<vscale x 8 x i16> %a, 
+                                                                   i64 64519)
+  ret <vscale x 8 x i16> %res
+}
+
+define <vscale x 4 x i32> @eor_i32(<vscale x 4 x i32> %a) {
+; CHECK-LABEL: eor_i32:
+; CHECK: eor z0.s, z0.s, #0xffff00
+; CHECK-NEXT: ret
+  %res = call <vscale x 4 x i32> @llvm.aarch64.sve.eor.imm.nxv4i32(<vscale x 4 x i32> %a, 
+                                                                   i64 16776960)
+  ret <vscale x 4 x i32> %res
+}
+
+define <vscale x 2 x i64> @eor_i64(<vscale x 2 x i64> %a) {
+; CHECK-LABEL: eor_i64:
+; CHECK: eor z0.d, z0.d, #0xfffc000000000000
+; CHECK-NEXT: ret
+  %res = call <vscale x 2 x i64> @llvm.aarch64.sve.eor.imm.nxv2i64(<vscale x 2 x i64> %a, 
+                                                                   i64 18445618173802708992)
+  ret <vscale x 2 x i64> %res
+}
+
+define <vscale x 16 x i8> @and_i8(<vscale x 16 x i8> %a) {
+; CHECK-LABEL: and_i8:
+; CHECK: and z0.b, z0.b, #0xf
+; CHECK-NEXT: ret
+  %res = call <vscale x 16 x i8> @llvm.aarch64.sve.and.imm.nxv16i8(<vscale x 16 x i8> %a, 
+                                                                   i64 15)
+  ret <vscale x 16 x i8> %res
+}
+
+define <vscale x 8 x i16> @and_i16(<vscale x 8 x i16> %a) {
+; CHECK-LABEL: and_i16:
+; CHECK: and z0.h, z0.h, #0xfc07
+; CHECK-NEXT: ret
+  %res = call <vscale x 8 x i16> @llvm.aarch64.sve.and.imm.nxv8i16(<vscale x 8 x i16> %a, 
+                                                                   i64 64519)
+  ret <vscale x 8 x i16> %res
+}
+
+define <vscale x 4 x i32> @and_i32(<vscale x 4 x i32> %a) {
+; CHECK-LABEL: and_i32:
+; CHECK: and z0.s, z0.s, #0xffff00
+; CHECK-NEXT: ret
+  %res = call <vscale x 4 x i32> @llvm.aarch64.sve.and.imm.nxv4i32(<vscale x 4 x i32> %a, 
+                                                                   i64 16776960)
+  ret <vscale x 4 x i32> %res
+}
+
+define <vscale x 2 x i64> @and_i64(<vscale x 2 x i64> %a) {
+; CHECK-LABEL: and_i64:
+; CHECK: and z0.d, z0.d, #0xfffc000000000000
+; CHECK-NEXT: ret
+  %res = call <vscale x 2 x i64> @llvm.aarch64.sve.and.imm.nxv2i64(<vscale x 2 x i64> %a,
+                                                                   i64 18445618173802708992)
+  ret <vscale x 2 x i64> %res
+}
+
+declare <vscale x 16 x i8> @llvm.aarch64.sve.orr.imm.nxv16i8(<vscale x 16 x i8>, i64)
+declare <vscale x 8 x i16> @llvm.aarch64.sve.orr.imm.nxv8i16(<vscale x 8 x i16>, i64)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.orr.imm.nxv4i32(<vscale x 4 x i32>, i64)
+declare <vscale x 2 x i64> @llvm.aarch64.sve.orr.imm.nxv2i64(<vscale x 2 x i64>, i64)
+declare <vscale x 16 x i8> @llvm.aarch64.sve.eor.imm.nxv16i8(<vscale x 16 x i8>, i64)
+declare <vscale x 8 x i16> @llvm.aarch64.sve.eor.imm.nxv8i16(<vscale x 8 x i16>, i64)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.eor.imm.nxv4i32(<vscale x 4 x i32>, i64)
+declare <vscale x 2 x i64> @llvm.aarch64.sve.eor.imm.nxv2i64(<vscale x 2 x i64>, i64)
+declare <vscale x 16 x i8> @llvm.aarch64.sve.and.imm.nxv16i8(<vscale x 16 x i8>, i64)
+declare <vscale x 8 x i16> @llvm.aarch64.sve.and.imm.nxv8i16(<vscale x 8 x i16>, i64)
+declare <vscale x 4 x i32> @llvm.aarch64.sve.and.imm.nxv4i32(<vscale x 4 x i32>, i64)
+declare <vscale x 2 x i64> @llvm.aarch64.sve.and.imm.nxv2i64(<vscale x 2 x i64>, i64)