Generalize fold

Author: MacDue

llvm/lib/CodeGen/ValueTypes.cpp

@@ -223,6 +223,8 @@ Type *EVT::getTypeForEVT(LLVMContext &Context) const {
   case MVT::x86mmx:  return llvm::FixedVectorType::get(llvm::IntegerType::get(Context, 64), 1);
   case MVT::aarch64svcount:
     return TargetExtType::get(Context, "aarch64.svcount");
+  case MVT::aarch64mfp8:
+    return VectorType::get(IntegerType::get(Context, 8), ElementCount::getFixed(1));
   case MVT::x86amx:  return Type::getX86_AMXTy(Context);
   case MVT::i64x8:   return IntegerType::get(Context, 512);
   case MVT::amdgpuBufferFatPointer:  return IntegerType::get(Context, 160);

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -1394,9 +1394,6 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
       }
     }
 
-    // v1i64 -> v1i8 truncstore represents a bsub FPR8 store.
-    setTruncStoreAction(MVT::v1i64, MVT::v1i8, Legal);
-
     for (auto Op :
          {ISD::FFLOOR, ISD::FNEARBYINT, ISD::FCEIL, ISD::FRINT, ISD::FTRUNC,
           ISD::FROUND, ISD::FROUNDEVEN, ISD::FMAXNUM_IEEE, ISD::FMINNUM_IEEE,
@@ -23981,6 +23978,8 @@ static SDValue combineI8TruncStore(StoreSDNode *ST, SelectionDAG &DAG,
 static unsigned getFPSubregForVT(EVT VT) {
   assert(VT.isSimple() && "Expected simple VT");
   switch (VT.getSimpleVT().SimpleTy) {
+  case MVT::aarch64mfp8:
+    return AArch64::bsub;
   case MVT::f16:
     return AArch64::hsub;
   case MVT::f32:
@@ -23992,22 +23991,6 @@ static unsigned getFPSubregForVT(EVT VT) {
   }
 }
 
-static EVT get64BitVector(EVT ElVT) {
-  assert(ElVT.isSimple() && "Expected simple VT");
-  switch (ElVT.getSimpleVT().SimpleTy) {
-  case MVT::i8:
-    return MVT::v8i8;
-  case MVT::i16:
-    return MVT::v4i16;
-  case MVT::i32:
-    return MVT::v2i32;
-  case MVT::i64:
-    return MVT::v1i64;
-  default:
-    llvm_unreachable("Unexpected VT!");
-  }
-}
-
 static SDValue performSTORECombine(SDNode *N,
                                    TargetLowering::DAGCombinerInfo &DCI,
                                    SelectionDAG &DAG,
@@ -24086,72 +24069,63 @@ static SDValue performSTORECombine(SDNode *N,
     SDValue ExtIdx = Value.getOperand(1);
     EVT VectorVT = Vector.getValueType();
     EVT ElemVT = VectorVT.getVectorElementType();
+
     if (!ValueVT.isInteger())
       return SDValue();
     if (ValueVT != MemVT && !ST->isTruncatingStore())
       return SDValue();
 
-    if (MemVT == MVT::i8) {
-      auto *ExtCst = dyn_cast<ConstantSDNode>(ExtIdx);
-      if (Subtarget->isNeonAvailable() &&
-          (VectorVT == MVT::v8i8 || VectorVT == MVT::v16i8) && ExtCst &&
-          !ExtCst->isZero() && ST->getBasePtr().getOpcode() != ISD::ADD) {
-        // These can lower to st1.b, which is preferable if we're unlikely to
-        // fold the addressing into the store.
-        return SDValue();
-      }
-
-      // Lower as truncstore of v1i64 -> v1i8 (which can lower to a bsub store).
-      SDValue Zero = DAG.getConstant(0, DL, MVT::i64);
-      SDValue ExtVector;
-      EVT VecVT64 = get64BitVector(ElemVT);
-      if (ExtCst && ExtCst->isZero()) {
-        ExtVector =
-            DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, VecVT64, Vector, Zero);
-      } else {
-        SDValue Ext = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL,
-                                  Value.getValueType(), Vector, ExtIdx);
-        ExtVector = DAG.getNode(ISD::INSERT_VECTOR_ELT, DL, VecVT64,
-                                DAG.getUNDEF(VecVT64), Ext, Zero);
-      }
-
-      SDValue Cast = DAG.getNode(AArch64ISD::NVCAST, DL, MVT::v1i64, ExtVector);
-      return DAG.getTruncStore(ST->getChain(), DL, Cast, ST->getBasePtr(),
-                               MVT::v1i8, ST->getMemOperand());
-    }
-
-    // TODO: Handle storing i8s to wider types.
-    if (ElemVT == MVT::i8)
+    // This could generate an additional extract if the index is non-zero and
+    // the extracted value has multiple uses.
+    auto *ExtCst = dyn_cast<ConstantSDNode>(ExtIdx);
+    if ((!ExtCst || !ExtCst->isZero()) && !Value.hasOneUse())
       return SDValue();
 
-    // Heuristic: If there are other users of integer scalars extracted from
-    // this vector that won't fold into the store -- abandon folding. Applying
-    // this fold may extend the vector lifetime and disrupt paired stores.
-    for (const auto &Use : Vector->uses()) {
-      if (Use.getResNo() != Vector.getResNo())
-        continue;
-      const SDNode *User = Use.getUser();
-      if (User->getOpcode() == ISD::EXTRACT_VECTOR_ELT &&
-          (!User->hasOneUse() ||
-           (*User->user_begin())->getOpcode() != ISD::STORE))
-        return SDValue();
+    if (Subtarget->isNeonAvailable() && ElemVT == MemVT &&
+        (VectorVT.is64BitVector() || VectorVT.is128BitVector()) && ExtCst &&
+        !ExtCst->isZero() && ST->getBasePtr().getOpcode() != ISD::ADD) {
+      // These can lower to st1, which is preferable if we're unlikely to fold
+      // the addressing into the store.
+      return SDValue();
     }
 
-    EVT FPElemVT = EVT::getFloatingPointVT(ElemVT.getSizeInBits());
-    EVT FPVectorVT = VectorVT.changeVectorElementType(FPElemVT);
-    SDValue Cast = DAG.getNode(ISD::BITCAST, DL, FPVectorVT, Vector);
-    SDValue Ext =
-        DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, FPElemVT, Cast, ExtIdx);
+    if (MemVT == MVT::i64 || MemVT == MVT::i32) {
+      // Heuristic: If there are other users of w/x integer scalars extracted
+      // from this vector that won't fold into the store -- abandon folding.
+      // Applying this fold may disrupt paired stores.
+      for (const auto &Use : Vector->uses()) {
+        if (Use.getResNo() != Vector.getResNo())
+          continue;
+        const SDNode *User = Use.getUser();
+        if (User->getOpcode() == ISD::EXTRACT_VECTOR_ELT &&
+            (!User->hasOneUse() ||
+             (*User->user_begin())->getOpcode() != ISD::STORE))
+          return SDValue();
+      }
+    }
 
-    EVT FPMemVT = EVT::getFloatingPointVT(MemVT.getSizeInBits());
-    if (ST->isTruncatingStore() && FPMemVT != FPElemVT) {
-      SDValue Trunc = DAG.getTargetExtractSubreg(getFPSubregForVT(FPMemVT), DL,
-                                                 FPMemVT, Ext);
-      return DAG.getStore(ST->getChain(), DL, Trunc, ST->getBasePtr(),
-                          ST->getMemOperand());
+    SDValue ExtVector = Vector;
+    if (!ExtCst || !ExtCst->isZero()) {
+      // Handle extracting from lanes != 0.
+      SDValue Ext = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL,
+                                Value.getValueType(), Vector, ExtIdx);
+      // FIXME: Using a fixed-size vector for the insertion should not be
+      // necessary, but SVE ISEL is missing some folds to avoid fmovs.
+      SDValue Zero = DAG.getConstant(0, DL, MVT::i64);
+      EVT InsertVectorVT = EVT::getVectorVT(
+          *DAG.getContext(), ElemVT,
+          VectorVT.getVectorElementCount().getKnownMinValue(), false);
+      ExtVector = DAG.getNode(ISD::INSERT_VECTOR_ELT, DL, InsertVectorVT,
+                              DAG.getUNDEF(InsertVectorVT), Ext, Zero);
     }
 
-    return DAG.getStore(ST->getChain(), DL, Ext, ST->getBasePtr(),
+    EVT FPMemVT = MemVT == MVT::i8
+                      ? MVT::aarch64mfp8
+                      : EVT::getFloatingPointVT(MemVT.getSizeInBits());
+    SDValue FPSubreg = DAG.getTargetExtractSubreg(getFPSubregForVT(FPMemVT), DL,
+                                                  FPMemVT, ExtVector);
+
+    return DAG.getStore(ST->getChain(), DL, FPSubreg, ST->getBasePtr(),
                         ST->getMemOperand());
   }
 
@@ -28878,10 +28852,6 @@ SDValue AArch64TargetLowering::LowerFixedLengthVectorStoreToSVE(
   auto Pg = getPredicateForFixedLengthVector(DAG, DL, VT);
   auto NewValue = convertToScalableVector(DAG, ContainerVT, Store->getValue());
 
-  // Can be lowered to a bsub store in ISEL.
-  if (VT == MVT::v1i64 && MemVT == MVT::v1i8)
-    return SDValue();
-
   if (VT.isFloatingPoint() && Store->isTruncatingStore()) {
     EVT TruncVT = ContainerVT.changeVectorElementType(
         Store->getMemoryVT().getVectorElementType());

llvm/lib/Target/AArch64/AArch64InstrInfo.td

@@ -4589,11 +4589,11 @@ def : Pat<(truncstorei16 GPR64:$Rt, (am_unscaled16 GPR64sp:$Rn, simm9:$offset)),
 def : Pat<(truncstorei8 GPR64:$Rt, (am_unscaled8 GPR64sp:$Rn, simm9:$offset)),
   (STURBBi (EXTRACT_SUBREG GPR64:$Rt, sub_32), GPR64sp:$Rn, simm9:$offset)>;
 
-// v1i64 -> bsub truncating stores
-def : Pat<(truncstorevi8 v1i64:$VT, (am_unscaled8 GPR64sp:$Rn, simm9:$offset)),
-          (STURBi (aarch64mfp8 (EXTRACT_SUBREG v1i64:$VT, bsub)), GPR64sp:$Rn, simm9:$offset)>;
-def : Pat<(truncstorevi8 v1i64:$VT, (am_indexed8 GPR64sp:$Rn, uimm12s4:$offset)),
-          (STRBui (aarch64mfp8 (EXTRACT_SUBREG v1i64:$VT, bsub)), GPR64sp:$Rn, uimm12s4:$offset)>;
+// aarch64mfp8 (bsub) stores
+def : Pat<(store aarch64mfp8:$Rt, (am_unscaled8 GPR64sp:$Rn, simm9:$offset)),
+          (STURBi FPR8:$Rt, GPR64sp:$Rn, simm9:$offset)>;
+def : Pat<(store aarch64mfp8:$Rt, (am_indexed8 GPR64sp:$Rn, uimm12s4:$offset)),
+          (STRBui FPR8:$Rt, GPR64sp:$Rn, uimm12s4:$offset)>;
 
 // Match stores from lane 0 to the appropriate subreg's store.
 multiclass VecStoreULane0Pat<SDPatternOperator StoreOp,

llvm/test/CodeGen/AArch64/aarch64-neon-vector-insert-uaddlv.ll

@@ -70,10 +70,10 @@ define void @insert_vec_v23i32_uaddlv_from_v8i16(ptr %0) {
 ; CHECK-NEXT:    movi.2d v0, #0000000000000000
 ; CHECK-NEXT:    movi.2d v2, #0000000000000000
 ; CHECK-NEXT:    str wzr, [x0, #88]
-; CHECK-NEXT:    str xzr, [x0, #80]
 ; CHECK-NEXT:    uaddlv.8h s1, v0
 ; CHECK-NEXT:    stp q0, q0, [x0, #16]
 ; CHECK-NEXT:    stp q0, q0, [x0, #48]
+; CHECK-NEXT:    str d0, [x0, #80]
 ; CHECK-NEXT:    mov.s v2[0], v1[0]
 ; CHECK-NEXT:    ucvtf.4s v1, v2
 ; CHECK-NEXT:    str q1, [x0]
@@ -146,12 +146,13 @@ define void @insert_vec_v6i64_uaddlv_from_v4i32(ptr %0) {
 ; CHECK-LABEL: insert_vec_v6i64_uaddlv_from_v4i32:
 ; CHECK:       ; %bb.0: ; %entry
 ; CHECK-NEXT:    movi.2d v0, #0000000000000000
-; CHECK-NEXT:    str xzr, [x0, #16]
+; CHECK-NEXT:    movi.2d v2, #0000000000000000
 ; CHECK-NEXT:    uaddlv.4s d1, v0
+; CHECK-NEXT:    str d0, [x0, #16]
 ; CHECK-NEXT:    fmov x8, d1
 ; CHECK-NEXT:    ucvtf s1, x8
-; CHECK-NEXT:    mov.s v0[0], v1[0]
-; CHECK-NEXT:    str q0, [x0]
+; CHECK-NEXT:    mov.s v2[0], v1[0]
+; CHECK-NEXT:    str q2, [x0]
 ; CHECK-NEXT:    ret
 
 entry: