[SVE][CodeGen] Extend isConstantSplatValue to support ISD::SPLAT_VECTOR

[kmclaughlin]

Summary

Updated the affected scalable_of_scalable tests in sve-gep.ll, as isConstantSplatValue now returns true in DAGCombiner::visitMUL and folds (mul x, 1) -> x

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

148

Is there a truncOrSelf we can use?

[kmclaughlin]

• Removed the EltSize < SplatVal.getBitWidth() check from isConstantSplatVector and used truncOrSelf instead

[kmclaughlin]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

148

There was no truncOrSelf that I could use here, but I've added this in a separate patch (D91445)

[sdesmalen]

Seems like a good change, LGTM!

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

bool ConstantFPSDNode::isValueValidForType(EVT VT,
   return !losesInfo;
 }
 
 //===----------------------------------------------------------------------===//
 //                              ISD Namespace
 //===----------------------------------------------------------------------===//
 
 bool ISD::isConstantSplatVector(const SDNode *N, APInt &SplatVal) {
+  if (N->getOpcode() == ISD::SPLAT_VECTOR) {
+    unsigned EltSize =
+        N->getValueType(0).getVectorElementType().getSizeInBits();
+    if (auto *Op0 = dyn_cast<ConstantSDNode>(N->getOperand(0))) {
+      SplatVal = Op0->getAPIntValue();
+      if (EltSize < SplatVal.getBitWidth())
+        SplatVal = SplatVal.trunc(EltSize);
+      return true;
+    }
+  }
+
   auto *BV = dyn_cast<BuildVectorSDNode>(N);
   if (!BV)
     return false;
 
   APInt SplatUndef;
   unsigned SplatBitSize;
   bool HasUndefs;
   unsigned EltSize = N->getValueType(0).getVectorElementType().getSizeInBits();

llvm/test/CodeGen/AArch64/sve-gep.ll

; CHECK-NEXT:    ret
   %d = getelementptr i8, i8* %base, <vscale x 2 x i32> %idx
   ret <vscale x 2 x i8*> %d
 }
 
 define <vscale x 2 x <vscale x 2 x i64>*> @scalable_of_scalable_1(<vscale x 2 x i64>* %base) {
 ; CHECK-LABEL: scalable_of_scalable_1:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    rdvl x8, #1
-; CHECK-NEXT:    mov z1.d, #1 // =0x1
-; CHECK-NEXT:    mov z0.d, x0
-; CHECK-NEXT:    mov z2.d, x8
-; CHECK-NEXT:    ptrue p0.d
-; CHECK-NEXT:    mla z0.d, p0/m, z2.d, z1.d
+; CHECK-NEXT:    mov z0.d, x8
+; CHECK-NEXT:    mov z1.d, x0
+; CHECK-NEXT:    add z0.d, z1.d, z0.d
 ; CHECK-NEXT:    ret
   %idx = shufflevector <vscale x 2 x i64> insertelement (<vscale x 2 x i64> undef, i64 1, i32 0), <vscale x 2 x i64> zeroinitializer, <vscale x 2 x i32> zeroinitializer
   %d = getelementptr <vscale x 2 x i64>, <vscale x 2 x i64>* %base, <vscale x 2 x i64> %idx
   ret <vscale x 2 x <vscale x 2 x i64>*> %d
 }
 
 define <vscale x 2 x <vscale x 2 x i64>*> @scalable_of_scalable_2(<vscale x 2 x <vscale x 2 x i64>*> %base) {
 ; CHECK-LABEL: scalable_of_scalable_2:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    rdvl x8, #1
-; CHECK-NEXT:    mov z1.d, #1 // =0x1
-; CHECK-NEXT:    mov z2.d, x8
-; CHECK-NEXT:    ptrue p0.d
-; CHECK-NEXT:    mla z0.d, p0/m, z2.d, z1.d
+; CHECK-NEXT:    mov z1.d, x8
+; CHECK-NEXT:    add z0.d, z0.d, z1.d
 ; CHECK-NEXT:    ret
   %idx = shufflevector <vscale x 2 x i64> insertelement (<vscale x 2 x i64> undef, i64 1, i32 0), <vscale x 2 x i64> zeroinitializer, <vscale x 2 x i32> zeroinitializer
   %d = getelementptr <vscale x 2 x i64>, <vscale x 2 x <vscale x 2 x i64>*> %base, <vscale x 2 x i64> %idx
   ret <vscale x 2 x <vscale x 2 x i64>*> %d
 }
 
 define <vscale x 2 x <vscale x 2 x i64>*> @scalable_of_scalable_3(<vscale x 2 x <vscale x 2 x i64>*> %base, <vscale x 2 x i32> %idx) {
 ; CHECK-LABEL: scalable_of_scalable_3: