[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 ↗

(On Diff #496098)

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 ↗

(On Diff #496098)

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

1703–1705 ↗

(On Diff #497588)

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 ↗

(On Diff #499202)

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

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

1705 ↗

(On Diff #499202)

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 ↗

(On Diff #499202)

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 ↗

(On Diff #499202)

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

Probably select as well,

llvm/lib/IR/Type.cpp

868

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/docs/AArch64SME.rst

 ----------------------
 
 Committing a lazy-save
 ----------------------
 
 Exception handling and ZA
 -------------------------
 
-4. References
+4. Types
+========
+
+AArch64 Predicate-as-Counter Type
+---------------------------------
+
+:Overview:
+
+The predicate-as-counter type represents the type of a predicate-as-counter
+value held in a AArch64 SVE predicate register. Such a value contains
+information about the number of active lanes, the element width and a bit that
+tells whether the generated mask should be inverted. ACLE intrinsics should be
+used to move the predicate-as-counter value to/from a predicate vector.
+
+There are certain limitations on the type:
+
+* The type can be used for function parameters and return values.
+
+* The supported LLVM operations on this type are limited to ``load``, ``store``,
+  ``phi``, ``select`` and ``alloca`` instructions.
+
+The predicate-as-counter type is a scalable type.
+
+:Syntax:
+
+::
+
+      target("aarch64.svcount")
+
+
+
+5. References
 =============
 
     .. _aarch64_sme_acle:
 
 1.  `SME ACLE Pull-request <https://github.com/ARM-software/acle/pull/188>`__
 
     .. _aarch64_sme_abi:
 
 2.  `SME ABI Pull-request <https://github.com/ARM-software/abi-aa/pull/123>`__

llvm/lib/IR/Type.cpp

 
 static TargetTypeInfo getTargetTypeInfo(const TargetExtType *Ty) {
   LLVMContext &C = Ty->getContext();
   StringRef Name = Ty->getName();
   if (Name.startswith("spirv.") || Name.startswith("opencl.")) {
     return TargetTypeInfo(Type::getInt8PtrTy(C, 0), TargetExtType::HasZeroInit,
                           TargetExtType::CanBeGlobal);
   }
+
+  // Opaque types in the AArch64 name space.
+  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/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
+}