[CodeGen] Use SPLAT_VECTOR for zeroinitialiser with scalable types

[david-arm]

Summary

When generating code for the LLVM IR zeroinitialiser operation, if
the vector type is scalable we should be using SPLAT_VECTOR instead
of BUILD_VECTOR.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

1562

D77587 introduced FixedVectorType and ScalableVectorType, which means you can instead write:

cast<FixedVectorType>(VecTy)->getNumElements()

That also removes the need for an additional assert.

1570

I'd prefer to see the Scalable and FixedWidth cases more separated here, something like

if (isa<ScalableVectorType>(...))
  // Op = splat vector
else {
  if (EltVT.isFloatingPoint())
    ...
  ...
}

1576

You can then write:

if (isa<ScalableVectorType>(VecTy))
  Op = DAG.getSplatVector(VT, getCurSDLoc(), Op);

[david-arm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

1562

OK sure. It's just that I already got the ElementCount above from VecTy so I wasn't sure it was worth casting VecTy to a FixedVectorType here - it seemd a bit redundant. What's the preferred style in general here? Avoiding using ElementCount wherever possible and cast to Fixed and Scalable types? It would be good to have a consistent style because there are two ways of reaching the same answer.

1570

I'm not sure what you mean here. Are you suggesting duplicating the code into two separate cases? We still need to check whether it is floating point in either case, Wouldn't this be better?

assert(isa<ConstantAggregateZero>(C) && "Unknown vector constant!");
EVT EltVT =
    TLI.getValueType(DAG.getDataLayout(), VecTy->getElementType());

SDValue Op;
if (EltVT.isFloatingPoint())
  Op = DAG.getConstantFP(0, getCurSDLoc(), EltVT);
else
  Op = DAG.getConstant(0, getCurSDLoc(), EltVT);

if (isa<ScalableVectorType>(VecTy))
  return NodeMap[V] = DAG.getSplatVector(VT, getCurSDLoc(), Op);

...  Continue as normal for FixedVectorType

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

1562

@ctetreau touched on that a bit on the mailing list yesterday, see http://lists.llvm.org/pipermail/llvm-dev/2020-April/141162.html.

AIUI:

• If the code only cares about fixed width vectors, use FixedVectorType and getNumElements.
• If the code only cares about scalable vectors, use ScalableVectorType and getMinNumElements or getElementCount.
• If the code cares about both, use VectorType and getElementCount.

1570

You're right, I misread this code.

In that case, does it make sense to move the return from this function into the branches, so that you get:

 if (constantVector) {
  ...
  return DAG.getBuildVector(..)
} else if (ConstantAggregateZero) {
  ...
  if (isa<ScalableVector>(..))
    return DAG.getSplatVector(..)
  else
    return DAG.getBuildVector(..);
}
llvm_unreachable("Unknown vector constant")

[sdesmalen]

LGTM

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

1587

nit: unnecessary newline.

1591

nit: the else is not necessary here.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

if (C->getType()->isStructTy() || C->getType()->isArrayTy()) {
 
       return DAG.getMergeValues(Constants, getCurSDLoc());
     }
 
     if (const BlockAddress *BA = dyn_cast<BlockAddress>(C))
       return DAG.getBlockAddress(BA, VT);
 
     VectorType *VecTy = cast<VectorType>(V->getType());
-    unsigned NumElements = VecTy->getNumElements();
+    ElementCount EltCnt = VecTy->getElementCount();
 
     // Now that we know the number and type of the elements, get that number of
     // elements into the Ops array based on what kind of constant it is.
     SmallVector<SDValue, 16> Ops;
     if (const ConstantVector *CV = dyn_cast<ConstantVector>(C)) {
-      for (unsigned i = 0; i != NumElements; ++i)
+      assert(!EltCnt.Scalable && "ConstantVector used for scalable types");
+      for (unsigned i = 0; i != EltCnt.Min; ++i)
         Ops.push_back(getValue(CV->getOperand(i)));
     } else {
       assert(isa<ConstantAggregateZero>(C) && "Unknown vector constant!");
       EVT EltVT =
           TLI.getValueType(DAG.getDataLayout(), VecTy->getElementType());
 
       SDValue Op;
       if (EltVT.isFloatingPoint())
         Op = DAG.getConstantFP(0, getCurSDLoc(), EltVT);
       else
         Op = DAG.getConstant(0, getCurSDLoc(), EltVT);
-      Ops.assign(NumElements, Op);
+
+      if (EltCnt.Scalable)
+        return NodeMap[V] = DAG.getSplatVector(VT, getCurSDLoc(), Op);
+
+      Ops.assign(EltCnt.Min, Op);
     }
 
     // Create a BUILD_VECTOR node.
     return NodeMap[V] = DAG.getBuildVector(VT, getCurSDLoc(), Ops);
   }
 
   // If this is a static alloca, generate it as the frameindex instead of
   // computation.
   if (const AllocaInst *AI = dyn_cast<AllocaInst>(V)) {
     DenseMap<const AllocaInst*, int>::iterator SI =
       FuncInfo.StaticAllocaMap.find(AI);
     if (SI != FuncInfo.StaticAllocaMap.end())
       return DAG.getFrameIndex(SI->second,
                                TLI.getFrameIndexTy(DAG.getDataLayout()));
   }
 
   // If this is an instruction which fast-isel has deferred, select it now.
   if (const Instruction *Inst = dyn_cast<Instruction>(V)) {
     unsigned InReg = FuncInfo.InitializeRegForValue(Inst);
 

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

+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve < %s | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-none-linux-gnu"
+
+define <vscale x 2 x i64> @test_zeroinit_2xi64() {
+; CHECK-LABEL: test_zeroinit_2xi64
+; CHECK:       mov z0.d, #0
+; CHECK-NEXT:  ret
+  ret <vscale x 2 x i64> zeroinitializer
+}
+
+define <vscale x 4 x i32> @test_zeroinit_4xi32() {
+; CHECK-LABEL: test_zeroinit_4xi32
+; CHECK:       mov z0.s, #0
+; CHECK-NEXT:  ret
+  ret <vscale x 4 x i32> zeroinitializer
+}
+
+define <vscale x 8 x i16> @test_zeroinit_8xi16() {
+; CHECK-LABEL: test_zeroinit_8xi16
+; CHECK:       mov z0.h, #0
+; CHECK-NEXT:  ret
+  ret <vscale x 8 x i16> zeroinitializer
+}
+
+define <vscale x 16 x i8> @test_zeroinit_16xi8() {
+; CHECK-LABEL: test_zeroinit_16xi8
+; CHECK:       mov z0.b, #0
+; CHECK-NEXT:  ret
+  ret <vscale x 16 x i8> zeroinitializer
+}
+
+define <vscale x 2 x double> @test_zeroinit_2xf64() {
+; CHECK-LABEL: test_zeroinit_2xf64
+; CHECK:       mov z0.d, #0
+; CHECK-NEXT:  ret
+  ret <vscale x 2 x double> zeroinitializer
+}
+
+define <vscale x 4 x float> @test_zeroinit_4xf32() {
+; CHECK-LABEL: test_zeroinit_4xf32
+; CHECK:       mov z0.s, #0
+; CHECK-NEXT:  ret
+  ret <vscale x 4 x float> zeroinitializer
+}
+
+define <vscale x 8 x half> @test_zeroinit_8xf16() {
+; CHECK-LABEL: test_zeroinit_8xf16
+; CHECK:       mov z0.h, #0
+; CHECK-NEXT:  ret
+  ret <vscale x 8 x half> zeroinitializer
+}
+
+define <vscale x 2 x i1> @test_zeroinit_2xi1() {
+; CHECK-LABEL: test_zeroinit_2xi1
+; CHECK:       whilelo p0.d, xzr, xzr
+; CHECK-NEXT:  ret
+  ret <vscale x 2 x i1> zeroinitializer
+}
+
+define <vscale x 4 x i1> @test_zeroinit_4xi1() {
+; CHECK-LABEL: test_zeroinit_4xi1
+; CHECK:       whilelo p0.s, xzr, xzr
+; CHECK-NEXT:  ret
+  ret <vscale x 4 x i1> zeroinitializer
+}
+
+define <vscale x 8 x i1> @test_zeroinit_8xi1() {
+; CHECK-LABEL: test_zeroinit_8xi1
+; CHECK:       whilelo p0.h, xzr, xzr
+; CHECK-NEXT:  ret
+  ret <vscale x 8 x i1> zeroinitializer
+}
+
+define <vscale x 16 x i1> @test_zeroinit_16xi1() {
+; CHECK-LABEL: test_zeroinit_16xi1
+; CHECK:       whilelo p0.b, xzr, xzr
+; CHECK-NEXT:  ret
+  ret <vscale x 16 x i1> zeroinitializer
+}