[AArch64][SVE] Implement punpk[hi|lo] intrinsics

[kmclaughlin]

Summary

Adds the following two intrinsics:

• int_aarch64_sve_punpkhi
• int_aarch64_sve_punpklo

This patch also contains a fix which allows LLVMHalfElementsVectorType
to forward reference overloadable arguments.

[sdesmalen]

LGTM. Looks like a straightforward fix and support for punpk(hi|lo) intrinsics.

llvm/lib/IR/Function.cpp

1213

Nice find!

llvm/test/CodeGen/AArch64/sve-intrinsics-pred-operations.ll

1

nit: You've added the file sve-intrinsics-pred-operations.ll. Is it worth putting these in a separate directory and separate out the tests per intrinsic, e.g. test/CodeGen/AArch64/SVE/sve-intrinsics-punpkhi.ll

llvm/include/llvm/IR/IntrinsicsAArch64.td

 let TargetPrefix = "aarch64" in {  // All intrinsics start with "llvm.aarch64.".
   class AdvSIMD_Merged1VectorArg_Intrinsic
     : Intrinsic<[llvm_anyvector_ty],
                 [LLVMMatchType<0>,
                  LLVMScalarOrSameVectorWidth<0, llvm_i1_ty>,
                  LLVMMatchType<0>],
                 [IntrNoMem]>;
 
+  class AdvSIMD_SVE_PUNPKHI_Intrinsic
+    : Intrinsic<[LLVMHalfElementsVectorType<0>],
+                [llvm_anyvector_ty],
+                [IntrNoMem]>;
+
   // This class of intrinsics are not intended to be useful within LLVM IR but
   // are instead here to support some of the more regid parts of the ACLE.
   class Builtin_SVCVT<string name, LLVMType OUT, LLVMType IN>
   : GCCBuiltin<"__builtin_sve_" # name>,
     Intrinsic<[OUT], [OUT, llvm_nxv16i1_ty, IN], [IntrNoMem]>;
 }
 
 //===----------------------------------------------------------------------===//
 // SVE
 
 let TargetPrefix = "aarch64" in {  // All intrinsics start with "llvm.aarch64.".
 
 //
 // Integer arithmetic
 //
 
 def int_aarch64_sve_abs : AdvSIMD_Merged1VectorArg_Intrinsic;
 def int_aarch64_sve_neg : AdvSIMD_Merged1VectorArg_Intrinsic;
 
 //
 // Floating-point comparisons
 //
 
 def int_aarch64_sve_fcvtzs_i32f16 : Builtin_SVCVT<"svcvt_s32_f16_m", llvm_nxv4i32_ty, llvm_nxv8f16_ty>;
+
+//
+// Predicate operations
+//
+
+def int_aarch64_sve_punpkhi : AdvSIMD_SVE_PUNPKHI_Intrinsic;
+def int_aarch64_sve_punpklo : AdvSIMD_SVE_PUNPKHI_Intrinsic;
 }

llvm/lib/IR/Function.cpp

case IITDescriptor::TruncArgument: {
         NewTy = IntegerType::get(ITy->getContext(), ITy->getBitWidth() / 2);
       else
         return true;
 
       return Ty != NewTy;
     }
     case IITDescriptor::HalfVecArgument:
       // If this is a forward reference, defer the check for later.
-      return D.getArgumentNumber() >= ArgTys.size() ||
-             !isa<VectorType>(ArgTys[D.getArgumentNumber()]) ||
+      if (D.getArgumentNumber() >= ArgTys.size())
+        return IsDeferredCheck || DeferCheck(Ty);
+      return !isa<VectorType>(ArgTys[D.getArgumentNumber()]) ||
              VectorType::getHalfElementsVectorType(
                      cast<VectorType>(ArgTys[D.getArgumentNumber()])) != Ty;
     case IITDescriptor::SameVecWidthArgument: {
       if (D.getArgumentNumber() >= ArgTys.size()) {
         // Defer check and subsequent check for the vector element type.
         Infos = Infos.slice(1);
         return IsDeferredCheck || DeferCheck(Ty);
       }

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

let Predicates = [HasSVE] in {
   defm REV_PP : sve_int_perm_reverse_p<"rev">;
   defm REV_ZZ : sve_int_perm_reverse_z<"rev">;
 
   defm SUNPKLO_ZZ : sve_int_perm_unpk<0b00, "sunpklo">;
   defm SUNPKHI_ZZ : sve_int_perm_unpk<0b01, "sunpkhi">;
   defm UUNPKLO_ZZ : sve_int_perm_unpk<0b10, "uunpklo">;
   defm UUNPKHI_ZZ : sve_int_perm_unpk<0b11, "uunpkhi">;
 
-  def  PUNPKLO_PP : sve_int_perm_punpk<0b0, "punpklo">;
-  def  PUNPKHI_PP : sve_int_perm_punpk<0b1, "punpkhi">;
+  defm PUNPKLO_PP : sve_int_perm_punpk<0b0, "punpklo", int_aarch64_sve_punpklo>;
+  defm PUNPKHI_PP : sve_int_perm_punpk<0b1, "punpkhi", int_aarch64_sve_punpkhi>;
 
   defm MOVPRFX_ZPzZ : sve_int_movprfx_pred_zero<0b000, "movprfx">;
   defm MOVPRFX_ZPmZ : sve_int_movprfx_pred_merge<0b001, "movprfx">;
   def MOVPRFX_ZZ : sve_int_bin_cons_misc_0_c<0b00000001, "movprfx", ZPRAny>;
   def FEXPA_ZZ_H : sve_int_bin_cons_misc_0_c<0b01000000, "fexpa", ZPR16>;
   def FEXPA_ZZ_S : sve_int_bin_cons_misc_0_c<0b10000000, "fexpa", ZPR32>;
   def FEXPA_ZZ_D : sve_int_bin_cons_misc_0_c<0b11000000, "fexpa", ZPR64>;
 

llvm/lib/Target/AArch64/SVEInstrFormats.td

let Predicates = [HasSVE] in {
   defm PTRUE  : sve_int_ptrue<0b000, "ptrue">;
   defm PTRUES : sve_int_ptrue<0b001, "ptrues">;
 }
 
 //===----------------------------------------------------------------------===//
 // SVE pattern match helpers.
 //===----------------------------------------------------------------------===//
 
+class SVE_1_Op_Pat<ValueType vtd, SDPatternOperator op, ValueType vt1,
+                   Instruction inst>
+: Pat<(vtd (op vt1:$Op1)),
+      (inst $Op1)>;
+
 class SVE_3_Op_Pat<ValueType vtd, SDPatternOperator op, ValueType vt1,
                    ValueType vt2, ValueType vt3, Instruction inst>
 : Pat<(vtd (op vt1:$Op1, vt2:$Op2, vt3:$Op3)),
       (inst $Op1, $Op2, $Op3)>;
 
 //===----------------------------------------------------------------------===//
 // SVE Predicate Misc Group
 //===----------------------------------------------------------------------===//
: I<(outs PPR16:$Pd), (ins PPR8:$Pn),
   let Inst{31-17} = 0b000001010011000;
   let Inst{16}    = opc;
   let Inst{15-9}  = 0b0100000;
   let Inst{8-5}   = Pn;
   let Inst{4}     = 0b0;
   let Inst{3-0}   = Pd;
 }
 
+multiclass sve_int_perm_punpk<bit opc, string asm, SDPatternOperator op> {
+  def NAME : sve_int_perm_punpk<opc, asm>;
+
+  def : SVE_1_Op_Pat<nxv8i1, op, nxv16i1, !cast<Instruction>(NAME)>;
+  def : SVE_1_Op_Pat<nxv4i1, op, nxv8i1,  !cast<Instruction>(NAME)>;
+  def : SVE_1_Op_Pat<nxv2i1, op, nxv4i1,  !cast<Instruction>(NAME)>;
+}
+
 class sve_int_rdffr_pred<bit s, string asm>
 : I<(outs PPR8:$Pd), (ins PPRAny:$Pg),
   asm, "\t$Pd, $Pg/z",
   "",
   []>, Sched<[]> {
   bits<4> Pd;
   bits<4> Pg;
   let Inst{31-23} = 0b001001010;

llvm/test/CodeGen/AArch64/sve-intrinsics-pred-operations.ll

+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve < %s | FileCheck %s
+
+;
+; PUNPKHI
+;
+
+define <vscale x 8 x i1> @punpkhi_b16(<vscale x 16 x i1> %a) {
+; CHECK-LABEL: punpkhi_b16
+; CHECK: punpkhi p0.h, p0.b
+; CHECK-NEXT: ret
+  %res = call <vscale x 8 x i1> @llvm.aarch64.sve.punpkhi.nxv8i1(<vscale x 16 x i1> %a)
+  ret <vscale x 8 x i1> %res
+}
+
+define <vscale x 4 x i1> @punpkhi_b8(<vscale x 8 x i1> %a) {
+; CHECK-LABEL: punpkhi_b8
+; CHECK: punpkhi p0.h, p0.b
+; CHECK-NEXT: ret
+  %res = call <vscale x 4 x i1> @llvm.aarch64.sve.punpkhi.nxv4i1(<vscale x 8 x i1> %a)
+  ret <vscale x 4 x i1> %res
+}
+
+define <vscale x 2 x i1> @punpkhi_b4(<vscale x 4 x i1> %a) {
+; CHECK-LABEL: punpkhi_b4
+; CHECK: punpkhi p0.h, p0.b
+; CHECK-NEXT: ret
+  %res = call <vscale x 2 x i1> @llvm.aarch64.sve.punpkhi.nxv2i1(<vscale x 4 x i1> %a)
+  ret <vscale x 2 x i1> %res
+}
+
+;
+; PUNPKLO
+;
+
+define <vscale x 8 x i1> @punpklo_b16(<vscale x 16 x i1> %a) {
+; CHECK-LABEL: punpklo_b16
+; CHECK: punpklo p0.h, p0.b
+; CHECK-NEXT: ret
+  %res = call <vscale x 8 x i1> @llvm.aarch64.sve.punpklo.nxv8i1(<vscale x 16 x i1> %a)
+  ret <vscale x 8 x i1> %res
+}
+
+define <vscale x 4 x i1> @punpklo_b8(<vscale x 8 x i1> %a) {
+; CHECK-LABEL: punpklo_b8
+; CHECK: punpklo p0.h, p0.b
+; CHECK-NEXT: ret
+  %res = call <vscale x 4 x i1> @llvm.aarch64.sve.punpklo.nxv4i1(<vscale x 8 x i1> %a)
+  ret <vscale x 4 x i1> %res
+}
+
+define <vscale x 2 x i1> @punpklo_b4(<vscale x 4 x i1> %a) {
+; CHECK-LABEL: punpklo_b4
+; CHECK: punpklo p0.h, p0.b
+; CHECK-NEXT: ret
+  %res = call <vscale x 2 x i1> @llvm.aarch64.sve.punpklo.nxv2i1(<vscale x 4 x i1> %a)
+  ret <vscale x 2 x i1> %res
+}
+
+declare <vscale x 8 x i1> @llvm.aarch64.sve.punpkhi.nxv8i1(<vscale x 16 x i1>)
+declare <vscale x 4 x i1> @llvm.aarch64.sve.punpkhi.nxv4i1(<vscale x 8 x i1>)
+declare <vscale x 2 x i1> @llvm.aarch64.sve.punpkhi.nxv2i1(<vscale x 4 x i1>)
+
+declare <vscale x 8 x i1> @llvm.aarch64.sve.punpklo.nxv8i1(<vscale x 16 x i1>)
+declare <vscale x 4 x i1> @llvm.aarch64.sve.punpklo.nxv4i1(<vscale x 8 x i1>)
+declare <vscale x 2 x i1> @llvm.aarch64.sve.punpklo.nxv2i1(<vscale x 4 x i1>)