[AArch64][SVE] Add DAG combines to improve SVE fixed type FP_EXTEND lowering

[bsmith]

Summary

Currently in some cases FP_EXTEND lowering for SVE fixed types can
introduce EXTRACT_SUBVECTOR nodes between purely fixed types. Lowering
these is less than ideal as it may involve going through memory.

Instead, we can push an extra extend into a load feeding the FP_EXTEND,
so as to ensure the load gets split into the same number of components
as the FP_EXTEND. The extra truncate required to do this can then be
combined with an extend that is introduced during FP_EXTEND lowering.

The net result of this is that the extends produced by the FP_EXTEND
lowering are folded into an extending load, and requirement for
EXTRACT_SUBVECTOR nodes is removed.

[bsmith]

• Fix mistake in bail condition when logic was reversed.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

15242

DAGCombiner::visitFP_EXTEND tries to do something similar, but unfortunately it is a bit too conservative. Look for:

// fold (fpext (load x)) -> (fpext (fptrunc (extload x)))

It seems logical to reuse that code, rather than reimplementing the same thing here. Perhaps you can use something like isVectorLoadExtDesirable instead of isLoadExtLegal to determine whether to fold this.

[bsmith]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

15242

There's actually a subtle difference here, that existing DAG combine is ultimately trying to fold the fpext into the load, creating an FP extending load, which is not something we can do. The DAG combine below is not trying to remove the fpext, but rather ensure the load is of the same size as the fpext so that is gets split in the same way.

I'm also not sure it makes sense for this to be in the generic DAG combine. The only reason this is sensible to do, is that our lowering of FP_EXTEND introduces an integer extend, which can then be folded into the integer truncate this DAG combine introduces. In absence of that fact, this DAG combine would likely just make things worse. It's also only really required because we can't sensibly lower extract_subvector.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

15242

Okay, I see what you mean.

I'm also not sure it makes sense for this to be in the generic DAG combine. The only reason this is sensible to do, is that our lowering of FP_EXTEND introduces an integer extend, which can then be folded into the integer truncate this DAG combine introduces. In absence of that fact, this DAG combine would likely just make things worse. It's also only really required because we can't sensibly lower extract_subvector.

If I understand it correctly, the target-specific part about this is how SVE's FCVT instruction works, i.e. by interpreting the input as unpacked, hence the need to bitcast/zero-extend, which becomes difficult to express in a generic way (in e.g. a DAGCombine).

15248

is there a reason this is specific to fixed-width vectors? It seems the same problem exists for:

%op1 = load <vscale x 8 x half>, <vscale x 8 x half>* %a
%res = fpext <vscale x 8 x half> %op1 to <vscale x 8 x float>
store <vscale x 8 x float> %res, <vscale x 8 x float>* %b

15269–15277

nit: is this loop maybe better expressed with llvm::any_of ?

15303

does the second operand actually matter, or can it be anything/undef?

[bsmith]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

15248

For the scalable case the generated extract_subvector doesn't require a trip through memory, instead it just does an unpack lo/hi pair which is much more reasonable.

(In fact this optimization actually makes things worse for scalable, we end up splitting the load but then combining it back together again with a UZP1, meaning the extract_subvectors are still present, but now also with a split load + recombine).

15269–15277

I'm not sure we can actually do that in this case. Using any_of would automatically dereference the use_iterator which calls getUser(), we also need to be able to call getUse() on the iterator.

[bsmith]

• Don't check 2nd operand of UZP1, it isn't used in this case.
• Redo test changes after recent test reformatting.

[bsmith]

Abandoning in favour of D114628 and D114580

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

#undef LCALLNAME5
   setTargetDAGCombine(ISD::CONCAT_VECTORS);
   setTargetDAGCombine(ISD::INSERT_SUBVECTOR);
   setTargetDAGCombine(ISD::STORE);
   if (Subtarget->supportsAddressTopByteIgnored())
     setTargetDAGCombine(ISD::LOAD);
 
   setTargetDAGCombine(ISD::MUL);
 
+  setTargetDAGCombine(ISD::FP_EXTEND);
+
   setTargetDAGCombine(ISD::SELECT);
   setTargetDAGCombine(ISD::VSELECT);
 
   setTargetDAGCombine(ISD::INTRINSIC_VOID);
   setTargetDAGCombine(ISD::INTRINSIC_W_CHAIN);
   setTargetDAGCombine(ISD::INSERT_VECTOR_ELT);
   setTargetDAGCombine(ISD::EXTRACT_VECTOR_ELT);
   setTargetDAGCombine(ISD::VECREDUCE_ADD);
if (Op1.getOperand(0).getOpcode() == AArch64ISD::UZP1) {
       SDValue Z = Op1.getOperand(0).getOperand(1);
       return DAG.getNode(AArch64ISD::UZP1, DL, ResVT, Op0, Z);
     }
   }
 
   return SDValue();
 }
 
+static SDValue performFpExtendCombine(SDNode *N, SelectionDAG &DAG,
+                                      const AArch64Subtarget *Subtarget) {
+  SDLoc DL(N);
+  SDValue Op = N->getOperand(0);
+  EVT VT = N->getValueType(0);
+
+  if (!VT.isFixedLengthVector())
+    return SDValue();
+
+  if (DAG.getTargetLoweringInfo().isTypeLegal(VT) &&
+      Subtarget->useSVEForFixedLengthVectors())
+    return SDValue();
+
+  // In cases where the result of the fp_extend is not legal, it will be
+  // expanded into multiple extract_subvectors which cannot be lowered without
+  // going through memory.
+  //
+  // If we push an extend into the load feeding the fp_extend, we can force the
+  // load to be be expanded into the same number of parts as the fp_extend,
+  // avoiding the need for extract_subvectors completely.
+  //
+  // As part of the lowering of FP_EXTEND for fixed length types uunpklo nodes
+  // will be introduced which will then combine with the truncate introduced
+  // after the load.
+  if (ISD::isNormalLoad(Op.getNode())) {
+    LoadSDNode *LD = cast<LoadSDNode>(Op.getNode());
+
+    // Check if there are other uses. If so, do not combine as it will introduce
+    // an extra load.
+    for (SDNode::use_iterator UI = LD->use_begin(), UE = LD->use_end();
+         UI != UE; ++UI) {
+      if (UI.getUse().getResNo() == 1) // Ignore uses of the chain result.
+        continue;
+      if (*UI != N)
+        return SDValue();
+    }
+
+    SDValue NewLoad = DAG.getExtLoad(
+        ISD::ZEXTLOAD, DL, VT.changeTypeToInteger(), LD->getChain(),
+        LD->getBasePtr(), LD->getMemoryVT().changeTypeToInteger(),
+        LD->getMemOperand());
+
+    DAG.ReplaceAllUsesOfValueWith(SDValue(LD, 1), NewLoad.getValue(1));
+
+    SDValue Trunc = DAG.getNode(
+        ISD::TRUNCATE, DL, Op->getValueType(0).changeTypeToInteger(), NewLoad);
+    SDValue Bitcast = DAG.getNode(ISD::BITCAST, DL, Op->getValueType(0), Trunc);
+
+    return DAG.getNode(ISD::FP_EXTEND, DL, VT, Bitcast);
+  }
+
+  return SDValue();
+}
+
+static SDValue performUunpkloCombine(SDNode *N, SelectionDAG &DAG) {
+  SDLoc DL(N);
+  SDValue Op = N->getOperand(0);
+  EVT VT = N->getValueType(0);
+
+  // uunpklo(uzp1(x, x)) where x = bitcast(zextload) -> x
+  if (Op->getOpcode() == AArch64ISD::UZP1 &&
+      Op->getOperand(0) == Op->getOperand(1)) {
+    EVT HalfVT = Op.getValueType();
+
+    // Ensure the unzip input is the same size as the unpack output
+    if (Op->getOperand(0)->getOpcode() != ISD::BITCAST ||
+        Op->getValueType(0) == VT)
+      return SDValue();
+
+    SDValue Bitcast = Op->getOperand(0);
+
+    // Look through bitcasts and unzips
+    SDValue Input = Bitcast->getOperand(0);
+    while (Input->getOpcode() == ISD::BITCAST ||
+           (Input->getOpcode() == AArch64ISD::UZP1 &&
+            Input->getOperand(0) == Input->getOperand(1)))
+      Input = Input->getOperand(0);
+
+    // Input should come from an extending load
+    if (!isa<MaskedLoadSDNode>(Input) ||
+        cast<MaskedLoadSDNode>(Input)->getExtensionType() != ISD::ZEXTLOAD)
+      return SDValue();
+
+    // Ensure that we don't care about the top half of the input
+    EVT MemVT = cast<MaskedLoadSDNode>(Input)->getMemoryVT();
+    if (isPackedVectorType(MemVT, DAG) &&
+        MemVT.getVectorElementType().getScalarSizeInBits() <=
+            HalfVT.getScalarSizeInBits())
+      return Bitcast->getOperand(0);
+  }
+
+  return SDValue();
+}
+
 static SDValue performGLD1Combine(SDNode *N, SelectionDAG &DAG) {
   unsigned Opc = N->getOpcode();
 
   assert(((Opc >= AArch64ISD::GLD1_MERGE_ZERO && // unsigned gather loads
            Opc <= AArch64ISD::GLD1_IMM_MERGE_ZERO) ||
           (Opc >= AArch64ISD::GLD1S_MERGE_ZERO && // signed gather loads
            Opc <= AArch64ISD::GLD1S_IMM_MERGE_ZERO)) &&
          "Invalid opcode.");
SDValue AArch64TargetLowering::PerformDAGCombine(SDNode *N,
   case AArch64ISD::NVCAST:
     return performNVCASTCombine(N);
   case AArch64ISD::SPLICE:
     return performSpliceCombine(N, DAG);
   case AArch64ISD::UZP1:
     return performUzpCombine(N, DAG);
   case AArch64ISD::SETCC_MERGE_ZERO:
     return performSetccMergeZeroCombine(N, DAG);
+  case ISD::FP_EXTEND:
+    return performFpExtendCombine(N, DAG, Subtarget);
   case AArch64ISD::GLD1_MERGE_ZERO:
   case AArch64ISD::GLD1_SCALED_MERGE_ZERO:
   case AArch64ISD::GLD1_UXTW_MERGE_ZERO:
   case AArch64ISD::GLD1_SXTW_MERGE_ZERO:
   case AArch64ISD::GLD1_UXTW_SCALED_MERGE_ZERO:
   case AArch64ISD::GLD1_SXTW_SCALED_MERGE_ZERO:
   case AArch64ISD::GLD1_IMM_MERGE_ZERO:
   case AArch64ISD::GLD1S_MERGE_ZERO:
   case AArch64ISD::GLD1S_SCALED_MERGE_ZERO:
   case AArch64ISD::GLD1S_UXTW_MERGE_ZERO:
   case AArch64ISD::GLD1S_SXTW_MERGE_ZERO:
   case AArch64ISD::GLD1S_UXTW_SCALED_MERGE_ZERO:
   case AArch64ISD::GLD1S_SXTW_SCALED_MERGE_ZERO:
   case AArch64ISD::GLD1S_IMM_MERGE_ZERO:
     return performGLD1Combine(N, DAG);
   case AArch64ISD::VASHR:
   case AArch64ISD::VLSHR:
     return performVectorShiftCombine(N, *this, DCI);
+  case AArch64ISD::UUNPKLO:
+    return performUunpkloCombine(N, DAG);
   case ISD::INSERT_VECTOR_ELT:
     return performInsertVectorEltCombine(N, DCI);
   case ISD::EXTRACT_VECTOR_ELT:
     return performExtractVectorEltCombine(N, DAG);
   case ISD::VECREDUCE_ADD:
     return performVecReduceAddCombine(N, DCI.DAG, Subtarget);
   case ISD::INTRINSIC_VOID:
   case ISD::INTRINSIC_W_CHAIN:

llvm/test/CodeGen/AArch64/sve-fixed-length-fp-extend-trunc.ll

 ; VBITS_GE_512: ptrue [[PG1:p[0-9]+]].h, vl16
 ; VBITS_GE_512-NEXT: ld1h { [[OP:z[0-9]+]].h }, [[PG1]]/z, [x0]
 ; VBITS_GE_512-NEXT: ptrue [[PG2:p[0-9]+]].s, vl16
 ; VBITS_GE_512-NEXT: uunpklo [[UPK:z[0-9]+]].s, [[OP]].h
 ; VBITS_GE_512-NEXT: fcvt [[RES:z[0-9]+]].s, [[PG2]]/m, [[UPK]].h
 ; VBITS_GE_512-NEXT: st1w { [[RES]].s }, [[PG1]], [x1]
 ; VBITS_GE_512-NEXT: ret
 
-; Ensure sensible type legalisation - fixed type extract_subvector codegen is poor currently.
-; VBITS_EQ_256-DAG: ptrue [[PG1:p[0-9]+]].h, vl16
-; VBITS_EQ_256-DAG: ld1h { [[VEC:z[0-9]+]].h }, [[PG1]]/z, [x0]
-; VBITS_EQ_256-DAG: st1h { [[VEC:z[0-9]+]].h }, [[PG1]], [x8]
-; VBITS_EQ_256-DAG: ldp q[[LO:[0-9]+]], q[[HI:[0-9]+]], [sp]
-; VBITS_EQ_256-DAG: ptrue [[PG2:p[0-9]+]].s, vl8
+; Ensure sensible type legalisation
 ; VBITS_EQ_256-DAG: mov x[[NUMELTS:[0-9]+]], #8
-; VBITS_EQ_256-DAG: uunpklo [[UPK_LO:z[0-9]+]].s, z[[LO]].h
-; VBITS_EQ_256-DAG: uunpklo [[UPK_HI:z[0-9]+]].s, z[[HI]].h
-; VBITS_EQ_256-DAG: fcvt [[RES_LO:z[0-9]+]].s, [[PG2]]/m, [[UPK_LO]].h
-; VBITS_EQ_256-DAG: fcvt [[RES_HI:z[0-9]+]].s, [[PG2]]/m, [[UPK_HI]].h
-; VBITS_EQ_256-DAG: st1w { [[RES_LO]].s }, [[PG2]], [x1]
-; VBITS_EQ_256-DAG: st1w { [[RES_HI]].s }, [[PG2]], [x1, x[[NUMELTS]], lsl #2]
+; VBITS_EQ_256-DAG: ptrue [[PG:p[0-9]+]].s, vl8
+; VBITS_EQ_256-DAG: ld1h { [[VEC_LO:z[0-9]+]].s }, [[PG]]/z, [x0]
+; VBITS_EQ_256-DAG: ld1h { [[VEC_HI:z[0-9]+]].s }, [[PG]]/z, [x0, x[[NUMELTS]], lsl #1]
+; VBITS_EQ_256-DAG: fcvt [[RES_LO:z[0-9]+]].s, [[PG]]/m, [[VEC_LO]].h
+; VBITS_EQ_256-DAG: fcvt [[RES_HI:z[0-9]+]].s, [[PG]]/m, [[VEC_HI]].h
+; VBITS_EQ_256-DAG: st1w { [[RES_LO]].s }, [[PG]], [x1]
+; VBITS_EQ_256-DAG: st1w { [[RES_HI]].s }, [[PG]], [x1, x[[NUMELTS]], lsl #2]
   %op1 = load <16 x half>, <16 x half>* %a
   %res = fpext <16 x half> %op1 to <16 x float>
   store <16 x float> %res, <16 x float>* %b
   ret void
 }
 
 define void @fcvt_v32f16_v32f32(<32 x half>* %a, <32 x float>* %b) #0 {
 ; CHECK-LABEL: fcvt_v32f16_v32f32:

 ; VBITS_GE_512-NEXT: ptrue [[PG:p[0-9]+]].d, vl8
 ; VBITS_GE_512-NEXT: uunpklo [[UPK1:z[0-9]+]].s, z[[OP]].h
 ; VBITS_GE_512-NEXT: uunpklo [[UPK2:z[0-9]+]].d, [[UPK1]].s
 ; VBITS_GE_512-NEXT: fcvt [[RES:z[0-9]+]].d, [[PG]]/m, [[UPK2]].h
 ; VBITS_GE_512-NEXT: st1d { [[RES]].d }, [[PG]], [x1]
 ; VBITS_GE_512-NEXT: ret
 
 ; Ensure sensible type legalisation.
-; VBITS_EQ_256-DAG: ldr q[[OP:[0-9]+]], [x0]
-; VBITS_EQ_256-DAG: ptrue [[PG:p[0-9]+]].d, vl4
 ; VBITS_EQ_256-DAG: mov x[[NUMELTS:[0-9]+]], #4
-; VBITS_EQ_256-DAG: ext v[[HI:[0-9]+]].16b, v[[OP]].16b, v[[OP]].16b, #8
-; VBITS_EQ_256-DAG: uunpklo [[UPK1_LO:z[0-9]+]].s, z[[OP]].h
-; VBITS_EQ_256-DAG: uunpklo [[UPK1_HI:z[0-9]+]].s, z[[HI]].h
-; VBITS_EQ_256-DAG: uunpklo [[UPK2_LO:z[0-9]+]].d, [[UPK1_LO]].s
-; VBITS_EQ_256-DAG: uunpklo [[UPK2_HI:z[0-9]+]].d, [[UPK2_HI]].s
-; VBITS_EQ_256-DAG: fcvt [[RES_LO:z[0-9]+]].d, [[PG]]/m, [[UPK2_LO]].h
-; VBITS_EQ_256-DAG: fcvt [[RES_HI:z[0-9]+]].d, [[PG]]/m, [[UPK2_HI]].h
+; VBITS_EQ_256-DAG: ptrue [[PG:p[0-9]+]].d, vl4
+; VBITS_EQ_254-DAG: ld1h { [[VEC_HI:z[0-9]+]].d }, [[PG]]/z, [x0, x[[NUMELTS]], lsl #1]
+; VBITS_EQ_254-DAG: ld1h { [[VEC_LO:z[0-9]+]].d }, [[PG]]/z, [x0]
+; VBITS_EQ_256-DAG: fcvt [[RES_LO:z[0-9]+]].d, [[PG]]/m, [[VEC_LO]].h
+; VBITS_EQ_256-DAG: fcvt [[RES_HI:z[0-9]+]].d, [[PG]]/m, [[VEC_HI]].h
 ; VBITS_EQ_256-DAG: st1d { [[RES_LO]].d }, [[PG]], [x1]
 ; VBITS_EQ_256-DAG: st1d { [[RES_HI]].d }, [[PG]], [x1, x[[NUMELTS]], lsl #3]
   %op1 = load <8 x half>, <8 x half>* %a
   %res = fpext <8 x half> %op1 to <8 x double>
   store <8 x double> %res, <8 x double>* %b
   ret void
 }
 

 ; VBITS_GE_512: ptrue [[PG1:p[0-9]+]].s, vl8
 ; VBITS_GE_512-NEXT: ld1w { [[OP:z[0-9]+]].s }, [[PG1]]/z, [x0]
 ; VBITS_GE_512-NEXT: ptrue [[PG:p[0-9]+]].d, vl8
 ; VBITS_GE_512-NEXT: uunpklo [[UPK:z[0-9]+]].d, [[OP]].s
 ; VBITS_GE_512-NEXT: fcvt [[RES:z[0-9]+]].d, [[PG1]]/m, [[UPK]].s
 ; VBITS_GE_512-NEXT: st1d { [[RES]].d }, [[PG1]], [x1]
 ; VBITS_GE_512-NEXT: ret
 
-; Ensure sensible type legalisation - fixed type extract_subvector codegen is poor currently.
-; VBITS_EQ_256-DAG: ptrue [[PG1:p[0-9]+]].s, vl8
-; VBITS_EQ_256-DAG: ld1w { [[VEC:z[0-9]+]].s }, [[PG1]]/z, [x0]
-; VBITS_EQ_256-DAG: st1w { [[VEC:z[0-9]+]].s }, [[PG1]], [x8]
-; VBITS_EQ_256-DAG: ldp q[[LO:[0-9]+]], q[[HI:[0-9]+]], [sp]
-; VBITS_EQ_256-DAG: ptrue [[PG2:p[0-9]+]].d, vl4
+; Ensure sensible type legalisation
 ; VBITS_EQ_256-DAG: mov x[[NUMELTS:[0-9]+]], #4
-; VBITS_EQ_256-DAG: uunpklo [[UPK_LO:z[0-9]+]].d, z[[LO]].s
-; VBITS_EQ_256-DAG: uunpklo [[UPK_HI:z[0-9]+]].d, z[[HI]].s
-; VBITS_EQ_256-DAG: fcvt [[RES_LO:z[0-9]+]].d, [[PG2]]/m, [[UPK_LO]].s
-; VBITS_EQ_256-DAG: fcvt [[RES_HI:z[0-9]+]].d, [[PG2]]/m, [[UPK_HI]].s
-; VBITS_EQ_256-DAG: st1d { [[RES_LO]].d }, [[PG2]], [x1]
-; VBITS_EQ_256-DAG: st1d { [[RES_HI]].d }, [[PG2]], [x1, x[[NUMELTS]], lsl #3]
+; VBITS_EQ_256-DAG: ptrue [[PG:p[0-9]+]].d, vl4
+; VBITS_EQ_256-DAG: ld1w { [[VEC_HI:z[0-9]+]].d }, [[PG]]/z, [x0, x[[NUMELTS]], lsl #2]
+; VBITS_EQ_256-DAG: ld1w { [[VEC_LO:z[0-9]+]].d }, [[PG]]/z, [x0]
+; VBITS_EQ_256-DAG: fcvt [[RES_LO:z[0-9]+]].d, [[PG]]/m, [[VEC_LO]].s
+; VBITS_EQ_256-DAG: fcvt [[RES_HI:z[0-9]+]].d, [[PG]]/m, [[VEC_HI]].s
+; VBITS_EQ_256-DAG: st1d { [[RES_LO]].d }, [[PG]], [x1]
+; VBITS_EQ_256-DAG: st1d { [[RES_HI]].d }, [[PG]], [x1, x[[NUMELTS]], lsl #3]
   %op1 = load <8 x float>, <8 x float>* %a
   %res = fpext <8 x float> %op1 to <8 x double>
   store <8 x double> %res, <8 x double>* %b
   ret void
 }
 
 define void @fcvt_v16f32_v16f64(<16 x float>* %a, <16 x double>* %b) #0 {
 ; CHECK-LABEL: fcvt_v16f32_v16f64: