[SelectionDAG][AArch64] Constant fold in SelectionDAG::getVScale if VScaleMin==VScaleMax.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

14846

Please use clang-format

14850–14857

Seems easier to write:

unsigned VScaleMin = Attr.getVScaleRangeMin();
if (std::optional<unsigned> VScaleMax = Attr.getVScaleRangeMax())
  if (*VScaleMax == VScaleMin)
    return DAG.getConstant(VScaleMin * N->getConstantOperandAPInt(0), SDLoc(N),
                           N->getValueType(0));
return SDValue();

?

[sdesmalen]

Did you also try to see if there is additional value doing this transform at the LLVM IR layer? (I don't know if that could unlock other optimisations since ConstantFolding kicks in earlier)

[craig.topper]

In D145113#4163863, @sdesmalen wrote:

Did you also try to see if there is additional value doing this transform at the LLVM IR layer? (I don't know if that could unlock other optimisations since ConstantFolding kicks in earlier)

I think it’s already in InstSimplify.

[sdesmalen]

In D145113#4163887, @craig.topper wrote:

In D145113#4163863, @sdesmalen wrote:

Did you also try to see if there is additional value doing this transform at the LLVM IR layer? (I don't know if that could unlock other optimisations since ConstantFolding kicks in earlier)

I think it’s already in InstSimplify.

Good, thanks for confirming!

[paulwalker-arm]

Not a strongly held view but is there value in restricting the combine to pre-operation legalisation? I ask because I wonder if they'll be a point during legalisation where somebody specifically wants the ISD::VSCALE node. Perhaps a different/better option is to move the logic into SelectionDAG::getVScale(), which can take a default false bool to disable the optimisation?

[craig.topper]

Move to getVScale

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

 //===- DAGCombiner.cpp - Implement a DAG node combiner --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
private:
     SDValue visitANY_EXTEND(SDNode *N);
     SDValue visitAssertExt(SDNode *N);
     SDValue visitAssertAlign(SDNode *N);
     SDValue visitSIGN_EXTEND_INREG(SDNode *N);
     SDValue visitEXTEND_VECTOR_INREG(SDNode *N);
     SDValue visitTRUNCATE(SDNode *N);
     SDValue visitBITCAST(SDNode *N);
     SDValue visitFREEZE(SDNode *N);
+    SDValue visitVSCALE(SDNode *N);
     SDValue visitBUILD_PAIR(SDNode *N);
     SDValue visitFADD(SDNode *N);
     SDValue visitVP_FADD(SDNode *N);
     SDValue visitSTRICT_FADD(SDNode *N);
     SDValue visitFSUB(SDNode *N);
     SDValue visitFMUL(SDNode *N);
     SDValue visitFMA(SDNode *N);
     SDValue visitFDIV(SDNode *N);
SDValue DAGCombiner::visit(SDNode *N) {
   case ISD::MLOAD:              return visitMLOAD(N);
   case ISD::MSCATTER:           return visitMSCATTER(N);
   case ISD::MSTORE:             return visitMSTORE(N);
   case ISD::LIFETIME_END:       return visitLIFETIME_END(N);
   case ISD::FP_TO_FP16:         return visitFP_TO_FP16(N);
   case ISD::FP16_TO_FP:         return visitFP16_TO_FP(N);
   case ISD::FP_TO_BF16:         return visitFP_TO_BF16(N);
   case ISD::FREEZE:             return visitFREEZE(N);
+  case ISD::VSCALE:             return visitVSCALE(N);
   case ISD::VECREDUCE_FADD:
   case ISD::VECREDUCE_FMUL:
   case ISD::VECREDUCE_ADD:
   case ISD::VECREDUCE_MUL:
   case ISD::VECREDUCE_AND:
   case ISD::VECREDUCE_OR:
   case ISD::VECREDUCE_XOR:
   case ISD::VECREDUCE_SMAX:
SDValue DAGCombiner::visitBITCAST(SDNode *N) {
   return SDValue();
 }
 
 SDValue DAGCombiner::visitBUILD_PAIR(SDNode *N) {
   EVT VT = N->getValueType(0);
   return CombineConsecutiveLoads(N, VT);
 }
 
+SDValue DAGCombiner::visitVSCALE(SDNode *N) {
+  // Constant fold vscale if the min and max value are known.
+  auto Attr = DAG.getMachineFunction().getFunction().getFnAttribute(Attribute::VScaleRange);
+  if (!Attr.isValid())
+    return SDValue();
+
+  unsigned VScaleMin = Attr.getVScaleRangeMin();
+  std::optional<unsigned> VScaleMax = Attr.getVScaleRangeMax();
+  if (!VScaleMax || VScaleMin != VScaleMax)
+    return SDValue();
+
+  APInt C = N->getConstantOperandAPInt(0);
+  C *= VScaleMin;
+  return DAG.getConstant(C, SDLoc(N), N->getValueType(0));
+}
+
 SDValue DAGCombiner::visitFREEZE(SDNode *N) {
   SDValue N0 = N->getOperand(0);
 
   if (DAG.isGuaranteedNotToBeUndefOrPoison(N0, /*PoisonOnly*/ false))
     return N0;
 
   // Fold freeze(op(x, ...)) -> op(freeze(x), ...).
   // Try to push freeze through instructions that propagate but don't produce

llvm/test/CodeGen/AArch64/insert-subvector-res-legalization.ll

 }
 
 define <vscale x 2 x i32> @vec_scalable_subvec_fixed_idx_nonzero_large_i32(<vscale x 2 x i32>* %a, <8 x i32>* %b) #1 {
 ; CHECK-LABEL: vec_scalable_subvec_fixed_idx_nonzero_large_i32:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    addvl sp, sp, #-1
 ; CHECK-NEXT:    ptrue p0.d
-; CHECK-NEXT:    cntd x8
+; CHECK-NEXT:    mov w8, #8
 ; CHECK-NEXT:    ld1w { z0.d }, p0/z, [x0]
-; CHECK-NEXT:    subs x8, x8, #8
+; CHECK-NEXT:    subs x9, x8, #8
 ; CHECK-NEXT:    ptrue p1.d, vl8
-; CHECK-NEXT:    csel x8, xzr, x8, lo
-; CHECK-NEXT:    mov w9, #8
-; CHECK-NEXT:    cmp x8, #8
-; CHECK-NEXT:    csel x8, x8, x9, lo
+; CHECK-NEXT:    csel x9, xzr, x9, lo
+; CHECK-NEXT:    cmp x9, #8
+; CHECK-NEXT:    csel x8, x9, x8, lo
 ; CHECK-NEXT:    mov x9, sp
 ; CHECK-NEXT:    st1d { z0.d }, p0, [sp]
 ; CHECK-NEXT:    ld1w { z0.d }, p1/z, [x1]
 ; CHECK-NEXT:    st1d { z0.d }, p0, [x9, x8, lsl #3]
 ; CHECK-NEXT:    ld1d { z0.d }, p0/z, [sp]
 ; CHECK-NEXT:    addvl sp, sp, #1
 ; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEXT:    ret

llvm/test/CodeGen/AArch64/sve-extract-fixed-vector.ll

 }
 
 ; Goes through memory currently; idx != 0.
 define <4 x i32> @extract_v4i32_nxv2i32_idx4(<vscale x 2 x i32> %vec) nounwind #1 {
 ; CHECK-LABEL: extract_v4i32_nxv2i32_idx4:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    addvl sp, sp, #-1
-; CHECK-NEXT:    cntd x8
+; CHECK-NEXT:    mov w8, #16
 ; CHECK-NEXT:    mov w9, #4
 ; CHECK-NEXT:    subs x8, x8, #4
 ; CHECK-NEXT:    ptrue p0.d
 ; CHECK-NEXT:    csel x8, xzr, x8, lo
 ; CHECK-NEXT:    st1d { z0.d }, p0, [sp]
 ; CHECK-NEXT:    cmp x8, #4
 ; CHECK-NEXT:    ptrue p0.d, vl4
 ; CHECK-NEXT:    csel x8, x8, x9, lo

 }
 
 ; Goes through memory currently; idx != 0.
 define <8 x i16> @extract_v8i16_nxv4i16_idx8(<vscale x 4 x i16> %vec) nounwind #1 {
 ; CHECK-LABEL: extract_v8i16_nxv4i16_idx8:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    addvl sp, sp, #-1
-; CHECK-NEXT:    cntw x8
+; CHECK-NEXT:    mov w8, #32
 ; CHECK-NEXT:    mov w9, #8
 ; CHECK-NEXT:    subs x8, x8, #8
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    csel x8, xzr, x8, lo
 ; CHECK-NEXT:    st1w { z0.s }, p0, [sp]
 ; CHECK-NEXT:    cmp x8, #8
 ; CHECK-NEXT:    ptrue p0.s, vl8
 ; CHECK-NEXT:    csel x8, x8, x9, lo

 }
 
 ; Goes through memory currently; idx != 0.
 define <8 x i16> @extract_v8i16_nxv2i16_idx8(<vscale x 2 x i16> %vec) nounwind #1 {
 ; CHECK-LABEL: extract_v8i16_nxv2i16_idx8:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    addvl sp, sp, #-1
-; CHECK-NEXT:    cntd x8
+; CHECK-NEXT:    mov w8, #16
 ; CHECK-NEXT:    mov w9, #8
 ; CHECK-NEXT:    subs x8, x8, #8
 ; CHECK-NEXT:    ptrue p0.d
 ; CHECK-NEXT:    csel x8, xzr, x8, lo
 ; CHECK-NEXT:    st1d { z0.d }, p0, [sp]
 ; CHECK-NEXT:    cmp x8, #8
 ; CHECK-NEXT:    ptrue p0.d, vl8
 ; CHECK-NEXT:    csel x8, x8, x9, lo

 }
 
 ; Goes through memory currently; idx != 0.
 define <16 x i8> @extract_v16i8_nxv8i8_idx16(<vscale x 8 x i8> %vec) nounwind #1 {
 ; CHECK-LABEL: extract_v16i8_nxv8i8_idx16:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    addvl sp, sp, #-1
-; CHECK-NEXT:    cnth x8
+; CHECK-NEXT:    mov w8, #64
 ; CHECK-NEXT:    mov w9, #16
 ; CHECK-NEXT:    subs x8, x8, #16
 ; CHECK-NEXT:    ptrue p0.h
 ; CHECK-NEXT:    csel x8, xzr, x8, lo
 ; CHECK-NEXT:    st1h { z0.h }, p0, [sp]
 ; CHECK-NEXT:    cmp x8, #16
 ; CHECK-NEXT:    ptrue p0.h, vl16
 ; CHECK-NEXT:    csel x8, x8, x9, lo

 }
 
 ; Goes through memory currently; idx != 0.
 define <16 x i8> @extract_v16i8_nxv4i8_idx16(<vscale x 4 x i8> %vec) nounwind #1 {
 ; CHECK-LABEL: extract_v16i8_nxv4i8_idx16:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    addvl sp, sp, #-1
-; CHECK-NEXT:    cntw x8
+; CHECK-NEXT:    mov w8, #32
 ; CHECK-NEXT:    mov w9, #16
 ; CHECK-NEXT:    subs x8, x8, #16
 ; CHECK-NEXT:    ptrue p0.s
 ; CHECK-NEXT:    csel x8, xzr, x8, lo
 ; CHECK-NEXT:    st1w { z0.s }, p0, [sp]
 ; CHECK-NEXT:    cmp x8, #16
 ; CHECK-NEXT:    ptrue p0.s, vl16
 ; CHECK-NEXT:    csel x8, x8, x9, lo

 }
 
 ; Goes through memory currently; idx != 0.
 define <16 x i8> @extract_v16i8_nxv2i8_idx16(<vscale x 2 x i8> %vec) nounwind #1 {
 ; CHECK-LABEL: extract_v16i8_nxv2i8_idx16:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    addvl sp, sp, #-1
-; CHECK-NEXT:    cntd x8
-; CHECK-NEXT:    mov w9, #16
-; CHECK-NEXT:    subs x8, x8, #16
+; CHECK-NEXT:    mov w8, #16
 ; CHECK-NEXT:    ptrue p0.d
-; CHECK-NEXT:    csel x8, xzr, x8, lo
+; CHECK-NEXT:    subs x9, x8, #16
 ; CHECK-NEXT:    st1d { z0.d }, p0, [sp]
-; CHECK-NEXT:    cmp x8, #16
-; CHECK-NEXT:    csel x8, x8, x9, lo
+; CHECK-NEXT:    csel x9, xzr, x9, lo
+; CHECK-NEXT:    cmp x9, #16
+; CHECK-NEXT:    csel x8, x9, x8, lo
 ; CHECK-NEXT:    mov x9, sp
 ; CHECK-NEXT:    ld1d { z0.d }, p0/z, [x9, x8, lsl #3]
 ; CHECK-NEXT:    uzp1 z0.s, z0.s, z0.s
 ; CHECK-NEXT:    uzp1 z0.h, z0.h, z0.h
 ; CHECK-NEXT:    uzp1 z0.b, z0.b, z0.b
 ; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
 ; CHECK-NEXT:    addvl sp, sp, #1
 ; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload

 
 ; Fixed length clamping
 
 define <2 x i64> @extract_fixed_v2i64_nxv2i64(<vscale x 2 x i64> %vec) nounwind #0 {
 ; CHECK-LABEL: extract_fixed_v2i64_nxv2i64:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    addvl sp, sp, #-1
-; CHECK-NEXT:    cntd x8
-; CHECK-NEXT:    mov w9, #2
-; CHECK-NEXT:    sub x8, x8, #2
 ; CHECK-NEXT:    ptrue p0.d
-; CHECK-NEXT:    cmp x8, #2
 ; CHECK-NEXT:    st1d { z0.d }, p0, [sp]
-; CHECK-NEXT:    csel x8, x8, x9, lo
-; CHECK-NEXT:    mov x9, sp
-; CHECK-NEXT:    lsl x8, x8, #3
-; CHECK-NEXT:    ldr q0, [x9, x8]
+; CHECK-NEXT:    ldr q0, [sp, #16]
 ; CHECK-NEXT:    addvl sp, sp, #1
 ; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEXT:    ret
   %retval = call <2 x i64> @llvm.vector.extract.v2i64.nxv2i64(<vscale x 2 x i64> %vec, i64 2)
   ret <2 x i64> %retval
 }
 
 define <4 x i64> @extract_fixed_v4i64_nxv2i64(<vscale x 2 x i64> %vec) nounwind #0 {
 ; CHECK-LABEL: extract_fixed_v4i64_nxv2i64:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    addvl sp, sp, #-1
-; CHECK-NEXT:    cntd x9
-; CHECK-NEXT:    mov w10, #4
-; CHECK-NEXT:    subs x9, x9, #4
+; CHECK-NEXT:    mov w9, #4
 ; CHECK-NEXT:    ptrue p0.d
-; CHECK-NEXT:    csel x9, xzr, x9, lo
+; CHECK-NEXT:    subs x10, x9, #4
 ; CHECK-NEXT:    st1d { z0.d }, p0, [sp]
-; CHECK-NEXT:    cmp x9, #4
-; CHECK-NEXT:    csel x9, x9, x10, lo
+; CHECK-NEXT:    csel x10, xzr, x10, lo
+; CHECK-NEXT:    cmp x10, #4
+; CHECK-NEXT:    csel x9, x10, x9, lo
 ; CHECK-NEXT:    mov x10, sp
 ; CHECK-NEXT:    ld1d { z0.d }, p0/z, [x10, x9, lsl #3]
 ; CHECK-NEXT:    st1d { z0.d }, p0, [x8]
 ; CHECK-NEXT:    addvl sp, sp, #1
 ; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEXT:    ret
   %retval = call <4 x i64> @llvm.vector.extract.v4i64.nxv2i64(<vscale x 2 x i64> %vec, i64 4)
   ret <4 x i64> %retval

llvm/test/CodeGen/AArch64/sve-insert-vector.ll

 
 ; Fixed length clamping
 
 define <vscale x 2 x i64> @insert_fixed_v2i64_nxv2i64(<vscale x 2 x i64> %vec, <2 x i64> %subvec) nounwind #0 {
 ; CHECK-LABEL: insert_fixed_v2i64_nxv2i64:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    addvl sp, sp, #-1
-; CHECK-NEXT:    cntd x8
-; CHECK-NEXT:    mov w9, #2
-; CHECK-NEXT:    sub x8, x8, #2
 ; CHECK-NEXT:    ptrue p0.d
-; CHECK-NEXT:    cmp x8, #2
 ; CHECK-NEXT:    st1d { z0.d }, p0, [sp]
-; CHECK-NEXT:    csel x8, x8, x9, lo
-; CHECK-NEXT:    mov x9, sp
-; CHECK-NEXT:    lsl x8, x8, #3
-; CHECK-NEXT:    str q1, [x9, x8]
+; CHECK-NEXT:    str q1, [sp, #16]
 ; CHECK-NEXT:    ld1d { z0.d }, p0/z, [sp]
 ; CHECK-NEXT:    addvl sp, sp, #1
 ; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEXT:    ret
   %retval = call <vscale x 2 x i64> @llvm.vector.insert.nxv2i64.v2i64(<vscale x 2 x i64> %vec, <2 x i64> %subvec, i64 2)
   ret <vscale x 2 x i64> %retval
 }
 
 define <vscale x 2 x i64> @insert_fixed_v4i64_nxv2i64(<vscale x 2 x i64> %vec, <4 x i64>* %ptr) nounwind #0 {
 ; CHECK-LABEL: insert_fixed_v4i64_nxv2i64:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    addvl sp, sp, #-1
-; CHECK-NEXT:    cntd x8
+; CHECK-NEXT:    mov w8, #4
 ; CHECK-NEXT:    ptrue p0.d
-; CHECK-NEXT:    subs x8, x8, #4
+; CHECK-NEXT:    subs x9, x8, #4
 ; CHECK-NEXT:    ld1d { z1.d }, p0/z, [x0]
-; CHECK-NEXT:    csel x8, xzr, x8, lo
-; CHECK-NEXT:    mov w9, #4
-; CHECK-NEXT:    cmp x8, #4
+; CHECK-NEXT:    csel x9, xzr, x9, lo
 ; CHECK-NEXT:    st1d { z0.d }, p0, [sp]
-; CHECK-NEXT:    csel x8, x8, x9, lo
+; CHECK-NEXT:    cmp x9, #4
+; CHECK-NEXT:    csel x8, x9, x8, lo
 ; CHECK-NEXT:    mov x9, sp
 ; CHECK-NEXT:    st1d { z1.d }, p0, [x9, x8, lsl #3]
 ; CHECK-NEXT:    ld1d { z0.d }, p0/z, [sp]
 ; CHECK-NEXT:    addvl sp, sp, #1
 ; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEXT:    ret
   %subvec = load <4 x i64>, <4 x i64>* %ptr
   %retval = call <vscale x 2 x i64> @llvm.vector.insert.nxv2i64.v4i64(<vscale x 2 x i64> %vec, <4 x i64> %subvec, i64 4)