[CodeGen][SVE] Calculate correct type legalization for scalable vectors.

[sdesmalen]

Summary

This patch updates TargetLoweringBase::computeRegisterProperties and
TargetLoweringBase::getTypeConversion to support scalable vectors,
and make the right calls on how to legalise them. These changes are required
to legalise both MVTs and EVTs.

[efriedma]

llvm/include/llvm/CodeGen/TargetLowering.h

189

getTypeAction() and getTypeToTransformTo() are public; the unittest can use those APIs.

llvm/lib/CodeGen/TargetLoweringBase.cpp

942

report_fatal_error, and please move this up to just after the NumElts.isOne() check. (The EltVT.isInteger() code also assumes the element count can be divided by two.)

1356

We probably want a new LegalizeTypeAction kind to represent this. I guess just leaving it TypeLegal isn't too terrible for now, though.

[david-arm]

llvm/include/llvm/Support/TypeSize.h

54

Alternatively here we could have a more powerful isEqual() where we pass an unsigned integer of the same type in ElementCount, then you can simply do:

if (EC.isEqual(1)) { ... }

The function would look like

bool isEqual(const unsigned RHS) const {
  return Min == RHS && !Scalable;
}

If you don't want to do that in this patch that's fine, but it might be more generic?

llvm/lib/CodeGen/TargetLoweringBase.cpp

826

Instead of calling EVT::getVectorVT, could we just call SVT.getHalfNumVectorElementsVT() instead? It does the same thing.

[sdesmalen]

• Use public interfaces in tests (getTypeAction and getTypeToTransformTo)
• Changed llvm_unreachable -> report_fatal_error
• Added operator==(unsigned RHS) to ElementCount in favour of isOne()

[sdesmalen]

llvm/lib/CodeGen/TargetLoweringBase.cpp

1356

Something like TypeCannotBeLegalized?

[efriedma]

llvm/lib/CodeGen/TargetLoweringBase.cpp

1356

I was thinking something more like "TypeScalarizeVScaleVector". It's theoretically possible to legalize... although it's complicated by the fact that we can't easily emit loops in SelectionDAG.

[sdesmalen]

Added TypeScalarizeScalableVector enum value.

[efriedma]

LGTM

[foad]

@sdesmalen this is causing test failures in my DEBUG build:

Failing Tests (4):
  LLVM :: CodeGen/AMDGPU/fmax_legacy.f16.ll
  LLVM :: CodeGen/AMDGPU/fmin_legacy.f16.ll
  LLVM :: CodeGen/AMDGPU/select-vectors.ll
  LLVM :: CodeGen/AMDGPU/select.f16.ll

[sdesmalen]

In D80640#2076508, @foad wrote:

@sdesmalen this is causing test failures in my DEBUG build:

Failing Tests (4):
  LLVM :: CodeGen/AMDGPU/fmax_legacy.f16.ll
  LLVM :: CodeGen/AMDGPU/fmin_legacy.f16.ll
  LLVM :: CodeGen/AMDGPU/select-vectors.ll
  LLVM :: CodeGen/AMDGPU/select.f16.ll

Yes sorry for that. I had already reverted the patch and am looking into it. Cheers for pointing out!

llvm/include/llvm/CodeGen/TargetLowering.h

public:
   bool isAligned(Align AlignCheck) const {
     return isSrcAligned(AlignCheck) && isDstAligned(AlignCheck);
   }
 };
 
 /// This base class for TargetLowering contains the SelectionDAG-independent
 /// parts that can be used from the rest of CodeGen.
 class TargetLoweringBase {
 public:
   /// This enum indicates whether operations are valid for a target, and if not,
   /// what action should be used to make them valid.
   enum LegalizeAction : uint8_t {
     Legal,      // The target natively supports this operation.
     Promote,    // This operation should be executed in a larger type.
     Expand,     // Try to expand this to other ops, otherwise use a libcall.
     LibCall,    // Don't try to expand this to other ops, always use a libcall.
     Custom      // Use the LowerOperation hook to implement custom lowering.
public:
 
   /// Return true if the target has BitExtract instructions.
   bool hasExtractBitsInsn() const { return HasExtractBitsInsn; }
 
   /// Return the preferred vector type legalization action.
   virtual TargetLoweringBase::LegalizeTypeAction
   getPreferredVectorAction(MVT VT) const {
     // The default action for one element vectors is to scalarize
-    if (VT.getVectorNumElements() == 1)
+    if (VT.getVectorElementCount() == 1)
       return TypeScalarizeVector;
     // The default action for an odd-width vector is to widen.
     if (!VT.isPow2VectorType())
       return TypeWidenVector;
     // The default action for other vectors is to promote
     return TypePromoteInteger;
   }
 

llvm/include/llvm/Support/TypeSize.h

 // which may be scalable vectors. It provides convenience operators so that
 // it can be used in much the same way as a single scalar value.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_SUPPORT_TYPESIZE_H
 #define LLVM_SUPPORT_TYPESIZE_H
 
+#include "llvm/Support/MathExtras.h"
 #include "llvm/Support/WithColor.h"
 
 #include <cstdint>
 #include <cassert>
 
 namespace llvm {
 
 template <typename T> struct DenseMapInfo;
public:
   }
 
   bool operator==(const ElementCount& RHS) const {
     return Min == RHS.Min && Scalable == RHS.Scalable;
   }
   bool operator!=(const ElementCount& RHS) const {
     return !(*this == RHS);
   }
+  bool operator==(unsigned RHS) const {
+    return Min == RHS && !Scalable;
+  }
+  bool operator!=(unsigned RHS) const {
+    return !(*this == RHS);
+  }
+
+  ElementCount NextPowerOf2() const {
+    return ElementCount(llvm::NextPowerOf2(Min), Scalable);
+  }
 };
 
 // This class is used to represent the size of types. If the type is of fixed
 // size, it will represent the exact size. If the type is a scalable vector,
 // it will represent the known minimum size.
 class TypeSize {
   uint64_t MinSize;   // The known minimum size.
   bool IsScalable;    // If true, then the runtime size is an integer multiple

llvm/lib/CodeGen/TargetLoweringBase.cpp

if (VT.isSimple()) {
 
     assert((LA == TypeLegal || LA == TypeSoftenFloat ||
             LA == TypeSoftPromoteHalf ||
             (NVT.isVector() ||
              ValueTypeActions.getTypeAction(NVT) != TypePromoteInteger)) &&
            "Promote may not follow Expand or Promote");
 
     if (LA == TypeSplitVector)
-      return LegalizeKind(LA,
-                          EVT::getVectorVT(Context, SVT.getVectorElementType(),
-                                           SVT.getVectorNumElements() / 2));
+      return LegalizeKind(LA, SVT.getHalfNumVectorElementsVT());
     if (LA == TypeScalarizeVector)
       return LegalizeKind(LA, SVT.getVectorElementType());
     return LegalizeKind(LA, NVT);
   }
 
   // Handle Extended Scalar Types.
   if (!VT.isVector()) {
     assert(VT.isInteger() && "Float types must be simple");
if (BitSize < 8 || !isPowerOf2_32(BitSize)) {
       return LegalizeKind(TypePromoteInteger, NVT);
     }
 
     return LegalizeKind(TypeExpandInteger,
                         EVT::getIntegerVT(Context, VT.getSizeInBits() / 2));
   }
 
   // Handle vector types.
-  unsigned NumElts = VT.getVectorNumElements();
+  ElementCount NumElts = VT.getVectorElementCount();
   EVT EltVT = VT.getVectorElementType();
 
   // Vectors with only one element are always scalarized.
   if (NumElts == 1)
     return LegalizeKind(TypeScalarizeVector, EltVT);
 
+  if (VT.getVectorElementCount() == ElementCount(1, true))
+    report_fatal_error("Cannot legalize this vector");
+
   // Try to widen vector elements until the element type is a power of two and
   // promote it to a legal type later on, for example:
   // <3 x i8> -> <4 x i8> -> <4 x i32>
   if (EltVT.isInteger()) {
     // Vectors with a number of elements that is not a power of two are always
     // widened, for example <3 x i8> -> <4 x i8>.
     if (!VT.isPow2VectorType()) {
-      NumElts = (unsigned)NextPowerOf2(NumElts);
+      NumElts = NumElts.NextPowerOf2();
       EVT NVT = EVT::getVectorVT(Context, EltVT, NumElts);
       return LegalizeKind(TypeWidenVector, NVT);
     }
 
     // Examine the element type.
     LegalizeKind LK = getTypeConversion(Context, EltVT);
 
     // If type is to be expanded, split the vector.
if (EltVT.isInteger()) {
     // type with a promoted vector element type.
     EltVT = OldEltVT;
   }
 
   // Try to widen the vector until a legal type is found.
   // If there is no wider legal type, split the vector.
   while (true) {
     // Round up to the next power of 2.
-    NumElts = (unsigned)NextPowerOf2(NumElts);
+    NumElts = NumElts.NextPowerOf2();
 
     // If there is no simple vector type with this many elements then there
     // cannot be a larger legal vector type.  Note that this assumes that
     // there are no skipped intermediate vector types in the simple types.
     if (!EltVT.isSimple())
       break;
     MVT LargerVector = MVT::getVectorVT(EltVT.getSimpleVT(), NumElts);
     if (LargerVector == MVT())
       break;
 
     // If this type is legal then widen the vector.
     if (ValueTypeActions.getTypeAction(LargerVector) == TypeLegal)
       return LegalizeKind(TypeWidenVector, LargerVector);
   }
 
   // Widen odd vectors to next power of two.
   if (!VT.isPow2VectorType()) {
     EVT NVT = VT.getPow2VectorType(Context);
     return LegalizeKind(TypeWidenVector, NVT);
   }
 
   // Vectors with illegal element types are expanded.
-  EVT NVT = EVT::getVectorVT(Context, EltVT, VT.getVectorNumElements() / 2);
+  EVT NVT = EVT::getVectorVT(Context, EltVT, VT.getVectorElementCount() / 2);
   return LegalizeKind(TypeSplitVector, NVT);
 }
 
 static unsigned getVectorTypeBreakdownMVT(MVT VT, MVT &IntermediateVT,
                                           unsigned &NumIntermediates,
                                           MVT &RegisterVT,
                                           TargetLoweringBase *TLI) {
   // Figure out the right, legal destination reg to copy into.
void TargetLoweringBase::computeRegisterProperties(
   // Loop over all of the vector value types to see which need transformations.
   for (unsigned i = MVT::FIRST_VECTOR_VALUETYPE;
        i <= (unsigned)MVT::LAST_VECTOR_VALUETYPE; ++i) {
     MVT VT = (MVT::SimpleValueType) i;
     if (isTypeLegal(VT))
       continue;
 
     MVT EltVT = VT.getVectorElementType();
-    unsigned NElts = VT.getVectorNumElements();
+    ElementCount EC = VT.getVectorElementCount();
     bool IsLegalWiderType = false;
     bool IsScalable = VT.isScalableVector();
     LegalizeTypeAction PreferredAction = getPreferredVectorAction(VT);
     switch (PreferredAction) {
     case TypePromoteInteger: {
       MVT::SimpleValueType EndVT = IsScalable ?
                                    MVT::LAST_INTEGER_SCALABLE_VECTOR_VALUETYPE :
                                    MVT::LAST_INTEGER_FIXEDLEN_VECTOR_VALUETYPE;
       // Try to promote the elements of integer vectors. If no legal
       // promotion was found, fall through to the widen-vector method.
       for (unsigned nVT = i + 1;
            (MVT::SimpleValueType)nVT <= EndVT; ++nVT) {
         MVT SVT = (MVT::SimpleValueType) nVT;
         // Promote vectors of integers to vectors with the same number
         // of elements, with a wider element type.
         if (SVT.getScalarSizeInBits() > EltVT.getSizeInBits() &&
-            SVT.getVectorNumElements() == NElts &&
-            SVT.isScalableVector() == IsScalable && isTypeLegal(SVT)) {
+            SVT.getVectorElementCount() == EC && isTypeLegal(SVT)) {
           TransformToType[i] = SVT;
           RegisterTypeForVT[i] = SVT;
           NumRegistersForVT[i] = 1;
           ValueTypeActions.setTypeAction(VT, TypePromoteInteger);
           IsLegalWiderType = true;
           break;
         }
       }
       if (IsLegalWiderType)
         break;
       LLVM_FALLTHROUGH;
     }
 
     case TypeWidenVector:
-      if (isPowerOf2_32(NElts)) {
+      if (isPowerOf2_32(EC.Min)) {
         // Try to widen the vector.
         for (unsigned nVT = i + 1; nVT <= MVT::LAST_VECTOR_VALUETYPE; ++nVT) {
           MVT SVT = (MVT::SimpleValueType) nVT;
-          if (SVT.getVectorElementType() == EltVT
-              && SVT.getVectorNumElements() > NElts
-              && SVT.isScalableVector() == IsScalable && isTypeLegal(SVT)) {
+          if (SVT.getVectorElementType() == EltVT &&
+              SVT.isScalableVector() == IsScalable &&
+              SVT.getVectorElementCount().Min > EC.Min && isTypeLegal(SVT)) {
             TransformToType[i] = SVT;
             RegisterTypeForVT[i] = SVT;
             NumRegistersForVT[i] = 1;
             ValueTypeActions.setTypeAction(VT, TypeWidenVector);
             IsLegalWiderType = true;
             break;
           }
         }
case TypeScalarizeVector: {
       MVT NVT = VT.getPow2VectorType();
       if (NVT == VT) {
         // Type is already a power of 2.  The default action is to split.
         TransformToType[i] = MVT::Other;
         if (PreferredAction == TypeScalarizeVector)
           ValueTypeActions.setTypeAction(VT, TypeScalarizeVector);
         else if (PreferredAction == TypeSplitVector)
           ValueTypeActions.setTypeAction(VT, TypeSplitVector);
+        else if (EC.Min > 1)
+          ValueTypeActions.setTypeAction(VT, TypeSplitVector);
+        else if (EC == 1)
+          ValueTypeActions.setTypeAction(VT, TypeScalarizeVector);
         else
-          // Set type action according to the number of elements.
-          ValueTypeActions.setTypeAction(VT, NElts == 1 ? TypeScalarizeVector
-                                                        : TypeSplitVector);
+          ; // Can't split or scalarize `vscale x 1 x <eltty>``, but cannot use
+            // llvm_unreachable here because targets that support scalable
+            // vectors can address this with widening/promotion, whereas targets
+            // that don't support scalable vectors can't lower them at all and
+            // end up here.
       } else {
         TransformToType[i] = NVT;
         ValueTypeActions.setTypeAction(VT, TypeWidenVector);
       }
       break;
     }
     default:
       llvm_unreachable("Unknown vector legalization action!");

llvm/unittests/CodeGen/AArch64SelectionDAGTest.cpp

 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Target/TargetMachine.h"
 #include "gtest/gtest.h"
 
-using namespace llvm;
-
-namespace {
+namespace llvm {
 
 class AArch64SelectionDAGTest : public testing::Test {
 protected:
   static void SetUpTestCase() {
     InitializeAllTargets();
     InitializeAllTargetMCs();
   }
 
   void SetUp() override {
     StringRef Assembly = "define void @f() { ret void }";
 
     Triple TargetTriple("aarch64--");
     std::string Error;
     const Target *T = TargetRegistry::lookupTarget("", TargetTriple, Error);
     // FIXME: These tests do not depend on AArch64 specifically, but we have to
     // initialize a target. A skeleton Target for unittests would allow us to
     // always run these tests.
     if (!T)
       return;
 
     TargetOptions Options;
     TM = std::unique_ptr<LLVMTargetMachine>(static_cast<LLVMTargetMachine*>(
-        T->createTargetMachine("AArch64", "", "", Options, None, None,
+        T->createTargetMachine("AArch64", "", "+sve", Options, None, None,
                                CodeGenOpt::Aggressive)));
     if (!TM)
       return;
 
     SMDiagnostic SMError;
     M = parseAssemblyString(Assembly, SMError, Context);
     if (!M)
       report_fatal_error(SMError.getMessage());
void SetUp() override {
 
     DAG = std::make_unique<SelectionDAG>(*TM, CodeGenOpt::None);
     if (!DAG)
       report_fatal_error("DAG?");
     OptimizationRemarkEmitter ORE(F);
     DAG->init(*MF, ORE, nullptr, nullptr, nullptr, nullptr, nullptr);
   }
 
+  TargetLoweringBase::LegalizeTypeAction getTypeAction(EVT VT) {
+    return DAG->getTargetLoweringInfo().getTypeAction(Context, VT);
+  }
+
+  EVT getTypeToTransformTo(EVT VT) {
+    return DAG->getTargetLoweringInfo().getTypeToTransformTo(Context, VT);
+  }
+
   LLVMContext Context;
   std::unique_ptr<LLVMTargetMachine> TM;
   std::unique_ptr<Module> M;
   Function *F;
   std::unique_ptr<MachineFunction> MF;
   std::unique_ptr<SelectionDAG> DAG;
 };
 
TEST_F(AArch64SelectionDAGTest, getSplatSourceVector_Scalable_ADD_of_SPLAT_VECTOR) {
   EXPECT_EQ(Val1->getOpcode(), ISD::SPLAT_VECTOR);
   SDValue Op = DAG->getNode(ISD::ADD, Loc, VecVT, Val1, Val2);
 
   int SplatIdx = -1;
   EXPECT_EQ(DAG->getSplatSourceVector(Op, SplatIdx), Op);
   EXPECT_EQ(SplatIdx, 0);
 }
 
-} // end anonymous namespace
+TEST_F(AArch64SelectionDAGTest, getTypeConversion_SplitScalableMVT) {
+  if (!TM)
+    return;
+
+  MVT VT = MVT::nxv4i64;
+  EXPECT_EQ(getTypeAction(VT), TargetLoweringBase::TypeSplitVector);
+  ASSERT_TRUE(getTypeToTransformTo(VT).isScalableVector());
+}
+
+TEST_F(AArch64SelectionDAGTest, getTypeConversion_PromoteScalableMVT) {
+  if (!TM)
+    return;
+
+  MVT VT = MVT::nxv2i32;
+  EXPECT_EQ(getTypeAction(VT), TargetLoweringBase::TypePromoteInteger);
+  ASSERT_TRUE(getTypeToTransformTo(VT).isScalableVector());
+}
+
+TEST_F(AArch64SelectionDAGTest, getTypeConversion_NoScalarizeMVT_nxv1f32) {
+  if (!TM)
+    return;
+
+  MVT VT = MVT::nxv1f32;
+  EXPECT_NE(getTypeAction(VT), TargetLoweringBase::TypeScalarizeVector);
+  ASSERT_TRUE(getTypeToTransformTo(VT).isScalableVector());
+}
+
+TEST_F(AArch64SelectionDAGTest, getTypeConversion_SplitScalableEVT) {
+  if (!TM)
+    return;
+
+  EVT VT = EVT::getVectorVT(Context, MVT::i64, 256, true);
+  EXPECT_EQ(getTypeAction(VT), TargetLoweringBase::TypeSplitVector);
+  EXPECT_EQ(getTypeToTransformTo(VT), VT.getHalfNumVectorElementsVT(Context));
+}
+
+TEST_F(AArch64SelectionDAGTest, getTypeConversion_WidenScalableEVT) {
+  if (!TM)
+    return;
+
+  EVT FromVT = EVT::getVectorVT(Context, MVT::i64, 6, true);
+  EVT ToVT = EVT::getVectorVT(Context, MVT::i64, 8, true);
+
+  EXPECT_EQ(getTypeAction(FromVT), TargetLoweringBase::TypeWidenVector);
+  EXPECT_EQ(getTypeToTransformTo(FromVT), ToVT);
+}
+
+TEST_F(AArch64SelectionDAGTest, getTypeConversion_NoScalarizeEVT_nxv1f128) {
+  if (!TM)
+    return;
+
+  EVT FromVT = EVT::getVectorVT(Context, MVT::f128, 1, true);
+  EXPECT_DEATH(getTypeAction(FromVT), "Cannot legalize this vector");
+}
+
+
+} // end namespace llvm