[IR] Add LLVM IR support for target("aarch64.svcount") type.

[sdesmalen]

Summary

The C and C++ Language Extensions for AArch64 SME2 [1] adds a new type called
svcount_t which describes a predicate. This is not a predicate vector
mask, but rather a description of a predicate vector mask that can be
expanded into a mask using explicit instructions. The type is a scalable
opaque type.

To implement svcount_t type this patch uses the existing Target Extension Type
mechanism, but adds further support so that this type can be a scalable type.

AArch64 CodeGen support will follow in a separate patch.

[1] https://github.com/ARM-software/acle/pull/217

[nikic]

Any chance this can be implemented as an opaque type as introduced in https://reviews.llvm.org/D135202 instead? Providing a general mechanism for target-specific types is the purpose of opaque types, so it would be great if they work for this use case. cc @jcranmer-intel

[aemerson]

In D136861#3889474, @nikic wrote:

Any chance this can be implemented as an opaque type as introduced in https://reviews.llvm.org/D135202 instead? Providing a general mechanism for target-specific types is the purpose of opaque types, so it would be great if they work for this use case. cc @jcranmer-intel

Can this be implemented in the existing way? Otherwise users of -DCLANG_ENABLE_OPAQUE_POINTERS=OFF/-Xclang -no-opaque-pointers won't be able to use SME intrinsics.

[nikic]

In D136861#3889897, @aemerson wrote:

In D136861#3889474, @nikic wrote:

Any chance this can be implemented as an opaque type as introduced in https://reviews.llvm.org/D135202 instead? Providing a general mechanism for target-specific types is the purpose of opaque types, so it would be great if they work for this use case. cc @jcranmer-intel

Can this be implemented in the existing way? Otherwise users of -DCLANG_ENABLE_OPAQUE_POINTERS=OFF/-Xclang -no-opaque-pointers won't be able to use SME intrinsics.

Despite both having "opaque" in the name, opaque pointers and opaque types aren't really related. Opaque types would work even with -no-opaque-pointers.

[aemerson]

In D136861#3889901, @nikic wrote:

In D136861#3889897, @aemerson wrote:

In D136861#3889474, @nikic wrote:

Any chance this can be implemented as an opaque type as introduced in https://reviews.llvm.org/D135202 instead? Providing a general mechanism for target-specific types is the purpose of opaque types, so it would be great if they work for this use case. cc @jcranmer-intel

Can this be implemented in the existing way? Otherwise users of -DCLANG_ENABLE_OPAQUE_POINTERS=OFF/-Xclang -no-opaque-pointers won't be able to use SME intrinsics.

Despite both having "opaque" in the name, opaque pointers and opaque types aren't really related. Opaque types would work even with -no-opaque-pointers.

Ah ok, that's fine then. Although from looking at D135202 it doesn't look like the design has been settled yet? I don't think we should be writing patches upon speculative changes, so unless D135202 is going to land very soon I'd prefer we continued with the traditional approach.

[nikic]

In D136861#3890025, @aemerson wrote:

Ah ok, that's fine then. Although from looking at D135202 it doesn't look like the design has been settled yet? I don't think we should be writing patches upon speculative changes, so unless D135202 is going to land very soon I'd prefer we continued with the traditional approach.

Yeah, possibly. I'd suggest to still consider the possibility when you put up the RFC for this change, as it's likely going to be an easier sell if it doesn't require the addition of a new target-specific type.

[jcranmer-intel]

In D136861#3890025, @aemerson wrote:

Ah ok, that's fine then. Although from looking at D135202 it doesn't look like the design has been settled yet? I don't think we should be writing patches upon speculative changes, so unless D135202 is going to land very soon I'd prefer we continued with the traditional approach.

At a quick glance of this patch, this does seem like something that could be built on top of opaque types were they already landed, so it's at least worth considering how this type might be implemented as an opaque type. The thing that's the most unsettled in the design is effectively the communication of DataLayout-like type properties of target types, and while I hoped to have a solution to that coded up today, I didn't get enough time to do so.

[efriedma]

From the spec, svreinterpret_c(svreinterpret_b(x)) is equivalent to x, right? So you could just lower svcount_t to <vscale x 16 x i1>, but you want to avoid confusion with non-SME predicate vectors in LLVM IR?

[sdesmalen]

Refactored the patch to use the existing target("<customtype>") mechanism.

[sdesmalen]

Thanks for all the previous comments! I've rewritten the patch to use the new target-extension types, since that functionality has now landed. This made the patch a lot simpler. I've only had to add a new interface to check if the TargetExtType is scalable.

I will try to be a bit more responsive to comments this time around :)

In D136861#3900512, @efriedma wrote:

From the spec, svreinterpret_c(svreinterpret_b(x)) is equivalent to x, right? So you could just lower svcount_t to <vscale x 16 x i1>, but you want to avoid confusion with non-SME predicate vectors in LLVM IR?

We could let Clang emit the reinterpret cast, but we want the intrinsics to be correct for other front-ends as well think of e.g. MLIR. By giving it a proper type it becomes unambiguous what the format of the predicate value must be when passing it to the intrinsic.

[nikic]

llvm/lib/Transforms/Scalar/SROA.cpp

895

I think it would be more elegant to get the TypeSize return value here and check whether it is isScalable(). That would avoid the need to explicitly check for ScalableVectorType and scalable target extension types. (Feel free to pre-commit as an NFC change and reduce the diff here.)

TypeSize Size = DL.getTypeStoreSize(LI.getType());
if (Size.isScalable())
  return PI.setAborted(&LI);
return handleLoadOrStore(LI.getType(), LI, Offset, Size.getFixedValue(), LI.isVolatile());

[sdesmalen]

llvm/lib/Transforms/Scalar/SROA.cpp

895

That makes sense, thanks for the suggestion. I can do the same (look at allocated type size, rather than the type itself) for all other cases in this file as well.

[sdesmalen]

Split out SROA changes into rG462227f1150f as suggested, and rebased this patch.

[jcranmer-intel]

I'd ideally like to see discussion of this target extension type somewhere in documentation, but alas, AArch64 is one of the backends that lacks a dedicated target documentation page.

[jcranmer-intel]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

1684–1686

A better fix here, I think, is to add a test for the block below so that it only applies to integer or integer vector types.

[sdesmalen]

• Added description of the svcount type to AArch64SME.rst
• Use isIntOrIntVectorTy() in foldSelectInstWithICmp
• Added Type::isScalableTy() convenience function.

[sdesmalen]

@jcranmer-intel thanks for your review, I think I've addressed both your comments now.

[nikic]

llvm/include/llvm/IR/Type.h

216

It doesn't look like these new methods are actually used?

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

1690

This could be committed separately.

[jcranmer-intel]

Thank you for adding documentation for target("aarch64.svcount"). Otherwise, the only other thing I would comment is the same thing that Nikita said with regards to the unused new methods.

[sdesmalen]

llvm/include/llvm/IR/Type.h

216

They are used in the patches that follow it (D143642 and D136862). I thought it made sense to add them to this patch, given that this patch adds the concept of scalable target extension types (with svcount being an instantiation of that). I'm happy to move the interfaces to the other patches if that has the preference.

[nikic]

llvm/include/llvm/IR/Type.h

216

I think it would be better to introduce them later, mainly because I'm not sure we really need them. I think checking isa<ScalabelVectorType> is an anti-pattern in most (but maybe not all) cases.

E.g. for the SCEV change I put up D144624 as an alternative which works on TypeSize only. I think something like this can usually be done (places that have to guard against scalable vectors are generally places that query TypeSize...)

[sdesmalen]

• Removed unused interfaces.
• Rebased after moving out change in foldSelectInstWithICmp.

[nikic]

LGTM from IR perspective. I assume this already has internal consensus on the ARM side.

llvm/docs/AArch64SME.rst

462 ↗

(On Diff #499876)

Probably select as well,

llvm/lib/IR/Type.cpp

877

Could drop the outer check while there is only one type...

[sdesmalen]

Thanks for the reviews! I've addressed the nits before committing the patch.

llvm/include/llvm/IR/Type.h

public:
   bool isX86_MMXTy() const { return getTypeID() == X86_MMXTyID; }
 
   /// Return true if this is X86 AMX.
   bool isX86_AMXTy() const { return getTypeID() == X86_AMXTyID; }
 
   /// Return true if this is a target extension type.
   bool isTargetExtTy() const { return getTypeID() == TargetExtTyID; }
 
+  /// Return true if this is a target extension type with a scalable layout.
+  bool isScalableTargetExtTy() const;
+
   /// Return true if this is a FP type or a vector of FP.
   bool isFPOrFPVectorTy() const { return getScalarType()->isFloatingPointTy(); }
 
   /// Return true if this is 'label'.
   bool isLabelTy() const { return getTypeID() == LabelTyID; }
 
   /// Return true if this is 'metadata'.
   bool isMetadataTy() const { return getTypeID() == MetadataTyID; }
 
   /// Return true if this is 'token'.
   bool isTokenTy() const { return getTypeID() == TokenTyID; }
 
   /// True if this is an instance of IntegerType.

llvm/lib/IR/Type.cpp

const fltSemantics &Type::getFltSemantics() const {
   default: llvm_unreachable("Invalid floating type");
   }
 }
 
 bool Type::isIEEE() const {
   return APFloat::getZero(getFltSemantics()).isIEEE();
 }
 
+bool Type::isScalableTargetExtTy() const {
+  if (auto *TT = dyn_cast<TargetExtType>(this))
+    return isa<ScalableVectorType>(TT->getLayoutType());
+  return false;
+}
+
 Type *Type::getFloatingPointTy(LLVMContext &C, const fltSemantics &S) {
   Type *Ty;
   if (&S == &APFloat::IEEEhalf())
     Ty = Type::getHalfTy(C);
   else if (&S == &APFloat::BFloat())
     Ty = Type::getBFloatTy(C);
   else if (&S == &APFloat::IEEEsingle())
     Ty = Type::getFloatTy(C);

 
 static TargetTypeInfo getTargetTypeInfo(const TargetExtType *Ty) {
   LLVMContext &C = Ty->getContext();
   StringRef Name = Ty->getName();
   if (Name.startswith("spirv.")) {
     return TargetTypeInfo(Type::getInt8PtrTy(C, 0), TargetExtType::HasZeroInit,
                           TargetExtType::CanBeGlobal);
   }
+
+  // Opaque types in the AArch64 name space.
+  if (Name.startswith("aarch64.")) {
+    if (Name == "aarch64.svcount")
+      return TargetTypeInfo(ScalableVectorType::get(Type::getInt1Ty(C), 16));
+  }
+
   return TargetTypeInfo(Type::getVoidTy(C));
 }
 
 Type *TargetExtType::getLayoutType() const {
   return getTargetTypeInfo(this).LayoutType;
 }
 
 bool TargetExtType::hasProperty(Property Prop) const {
   uint64_t Properties = getTargetTypeInfo(this).Properties;
   return (Properties & Prop) == Prop;
 }

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

if (Pred == ICmpInst::ICMP_SGT && match(CmpRHS, m_AllOnes()) &&
     Builder.SetInsertPoint(&SI);
     Value *IsNeg = Builder.CreateIsNeg(CmpLHS, ICI->getName());
     replaceOperand(SI, 0, IsNeg);
     SI.swapValues();
     SI.swapProfMetadata();
     return &SI;
   }
 
+  // The code below doesn't work on scalable target types.
+  if (TrueVal->getType()->isScalableTargetExtTy())
+    return nullptr;
+
   // FIXME: This code is nearly duplicated in InstSimplify. Using/refactoring
   // decomposeBitTestICmp() might help.
   {
     unsigned BitWidth =
         DL.getTypeSizeInBits(TrueVal->getType()->getScalarType());
     APInt MinSignedValue = APInt::getSignedMinValue(BitWidth);
     Value *X;
     const APInt *Y, *C;
     bool TrueWhenUnset;
     bool IsBitTest = false;
     if (ICmpInst::isEquality(Pred) &&

llvm/lib/Transforms/Scalar/SROA.cpp

private:
 
   void visitLoadInst(LoadInst &LI) {
     assert((!LI.isSimple() || LI.getType()->isSingleValueType()) &&
            "All simple FCA loads should have been pre-split");
 
     if (!IsOffsetKnown)
       return PI.setAborted(&LI);
 
-    if (isa<ScalableVectorType>(LI.getType()))
+    if (isa<ScalableVectorType>(LI.getType()) ||
+        LI.getType()->isScalableTargetExtTy())
       return PI.setAborted(&LI);
 
     uint64_t Size = DL.getTypeStoreSize(LI.getType()).getFixedValue();
     return handleLoadOrStore(LI.getType(), LI, Offset, Size, LI.isVolatile());
   }
 
   void visitStoreInst(StoreInst &SI) {
     Value *ValOp = SI.getValueOperand();
     if (ValOp == *U)
       return PI.setEscapedAndAborted(&SI);
     if (!IsOffsetKnown)
       return PI.setAborted(&SI);
 
-    if (isa<ScalableVectorType>(ValOp->getType()))
+    if (isa<ScalableVectorType>(ValOp->getType()) ||
+        ValOp->getType()->isScalableTargetExtTy())
       return PI.setAborted(&SI);
 
     uint64_t Size = DL.getTypeStoreSize(ValOp->getType()).getFixedValue();
 
     // If this memory access can be shown to *statically* extend outside the
     // bounds of the allocation, it's behavior is undefined, so simply
     // ignore it. Note that this is more strict than the generic clamping
     // behavior of insertUse. We also try to handle cases which might run the
if (AI.use_empty()) {
     Changed = true;
     return {Changed, CFGChanged};
   }
   const DataLayout &DL = AI.getModule()->getDataLayout();
 
   // Skip alloca forms that this analysis can't handle.
   auto *AT = AI.getAllocatedType();
   if (AI.isArrayAllocation() || !AT->isSized() || isa<ScalableVectorType>(AT) ||
-      DL.getTypeAllocSize(AT).getFixedValue() == 0)
+      AT->isScalableTargetExtTy() || DL.getTypeAllocSize(AT).getFixedValue() == 0)
     return {Changed, CFGChanged};
 
   // First, split any FCA loads and stores touching this alloca to promote
   // better splitting and promotion opportunities.
   IRBuilderTy IRB(&AI);
   AggLoadStoreRewriter AggRewriter(DL, IRB);
   Changed |= AggRewriter.rewrite(AI);
 
PreservedAnalyses SROAPass::runImpl(Function &F, DomTreeUpdater &RunDTU,
   C = &F.getContext();
   DTU = &RunDTU;
   AC = &RunAC;
 
   BasicBlock &EntryBB = F.getEntryBlock();
   for (BasicBlock::iterator I = EntryBB.begin(), E = std::prev(EntryBB.end());
        I != E; ++I) {
     if (AllocaInst *AI = dyn_cast<AllocaInst>(I)) {
-      if (isa<ScalableVectorType>(AI->getAllocatedType())) {
+      if (isa<ScalableVectorType>(AI->getAllocatedType()) ||
+          AI->getAllocatedType()->isScalableTargetExtTy()) {
         if (isAllocaPromotable(AI))
           PromotableAllocas.push_back(AI);
       } else {
         Worklist.insert(AI);
       }
     }
   }
 

llvm/test/Transforms/InstCombine/AArch64/sme-svcount.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes='instcombine' -S < %s | FileCheck %s
+
+define target("aarch64.svcount") @test_alloca_store_reload(target("aarch64.svcount") %val) nounwind {
+; CHECK-LABEL: @test_alloca_store_reload(
+; CHECK-NEXT:    ret target("aarch64.svcount") [[VAL:%.*]]
+;
+  %ptr = alloca target("aarch64.svcount"), align 1
+  store target("aarch64.svcount") %val, ptr %ptr
+  %res = load target("aarch64.svcount"), ptr %ptr
+  ret target("aarch64.svcount") %res
+}
+
+; Test that instcombine doesn't try to query the (scalable) size of target("aarch64.svcount")
+; in foldSelectInstWithICmp.
+define target("aarch64.svcount") @test_combine_on_select(target("aarch64.svcount") %x, target("aarch64.svcount") %y, i32 %k) {
+; CHECK-LABEL: @test_combine_on_select(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i32 [[K:%.*]], 42
+; CHECK-NEXT:    [[X_Y:%.*]] = select i1 [[CMP]], target("aarch64.svcount") [[X:%.*]], target("aarch64.svcount") [[Y:%.*]]
+; CHECK-NEXT:    ret target("aarch64.svcount") [[X_Y]]
+;
+  %cmp = icmp sgt i32 %k, 42
+  %x.y = select i1 %cmp, target("aarch64.svcount") %x, target("aarch64.svcount") %y
+  ret target("aarch64.svcount") %x.y
+}

llvm/test/Transforms/SROA/aarch64-sme-svcount.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes='sroa' -S < %s | FileCheck %s
+
+define target("aarch64.svcount") @test_alloca_store_reload(target("aarch64.svcount") %val) nounwind {
+; CHECK-LABEL: @test_alloca_store_reload(
+; CHECK-NEXT:    ret target("aarch64.svcount") [[VAL:%.*]]
+;
+  %ptr = alloca target("aarch64.svcount"), align 1
+  store target("aarch64.svcount") %val, ptr %ptr
+  %res = load target("aarch64.svcount"), ptr %ptr
+  ret target("aarch64.svcount") %res
+}