[Clang][AArch64][SME] Add intrinsics for ZA array load/store (LDR/STR)

[bryanpkc]

Summary

This patch adds support for the following SME ACLE intrinsics (as defined
in https://arm-software.github.io/acle/main/acle.html):

• svldr_vnum_za
• svstr_vnum_za

Co-authored-by: Sagar Kulkarni <[e-mail]>

[bryanpkc]

Rebased and cleaned up the patch.

[bryanpkc]

Instead of defining new type flags, EmitAArch64SMEBuiltinExpr is updated to switch on BuiltinID, as suggested by @sdesmalen.

[bryanpkc]

Also added SME attributes to the tests.

[bryanpkc]

Removed the attribute macro from tests, as @sdesmalen suggested.

[sdesmalen]

clang/test/CodeGen/aarch64-sme-intrinsics/acle_sme_str.c

25

This is missing an add of 15 to %slice_base.

[bryanpkc]

clang/test/CodeGen/aarch64-sme-intrinsics/acle_sme_str.c

25

Thanks for catching that. Fixed.

[sdesmalen]

Thanks for the quick fix!

clang/include/clang/Basic/TargetBuiltins.h

#undef LLVM_GET_SME_IMMCHECKTYPES
     MergeType getMergeType() const {
       return (MergeType)((Flags & MergeTypeMask) >> MergeTypeShift);
     }
 
     bool isLoad() const { return Flags & IsLoad; }
     bool isStore() const { return Flags & IsStore; }
     bool isMove() const { return Flags & IsMove; }
     bool isZero() const { return Flags & IsZero; }
+    bool isLoadReg() const { return Flags & IsLoadReg; }
+    bool isStoreReg() const { return Flags & IsStoreReg; }
 
     uint64_t getBits() const { return Flags; }
     bool isFlagSet(uint64_t Flag) const { return Flags & Flag; }
   };
 
   /// Hexagon builtins
   namespace Hexagon {
     enum {

clang/include/clang/Basic/arm_sme.td

 def SVLD1_HOR_VNUM_ZA64 : MInst<"svld1_hor_vnum_za64", "vimiPQl", "l", [IsLoad, IsOverloadNone, IsStreaming, IsSharedZA], MemEltTyDefault, "aarch64_sme_ld1d_horiz">;
 def SVLD1_HOR_VNUM_ZA128 : MInst<"svld1_hor_vnum_za128", "vimiPQl", "q", [IsLoad, IsOverloadNone, IsStreaming, IsSharedZA], MemEltTyDefault, "aarch64_sme_ld1q_horiz">;
 def SVLD1_VER_VNUM_ZA8 : MInst<"svld1_ver_vnum_za8", "vimiPQl", "c", [IsLoad, IsOverloadNone, IsStreaming, IsSharedZA], MemEltTyDefault, "aarch64_sme_ld1b_vert">;
 def SVLD1_VER_VNUM_ZA16 : MInst<"svld1_ver_vnum_za16", "vimiPQl", "s", [IsLoad, IsOverloadNone, IsStreaming, IsSharedZA], MemEltTyDefault, "aarch64_sme_ld1h_vert">;
 def SVLD1_VER_VNUM_ZA32 : MInst<"svld1_ver_vnum_za32", "vimiPQl", "i", [IsLoad, IsOverloadNone, IsStreaming, IsSharedZA], MemEltTyDefault, "aarch64_sme_ld1w_vert">;
 def SVLD1_VER_VNUM_ZA64 : MInst<"svld1_ver_vnum_za64", "vimiPQl", "l", [IsLoad, IsOverloadNone, IsStreaming, IsSharedZA], MemEltTyDefault, "aarch64_sme_ld1d_vert">;
 def SVLD1_VER_VNUM_ZA128 : MInst<"svld1_ver_vnum_za128", "vimiPQl", "q", [IsLoad, IsOverloadNone, IsStreaming, IsSharedZA], MemEltTyDefault, "aarch64_sme_ld1q_vert">;
 
+def SVLDR_VNUM_ZA : MInst<"svldr_vnum_za", "vmiQ", "", [IsOverloadNone, IsLoadReg, IsStreamingCompatible, IsSharedZA], MemEltTyDefault, "aarch64_sme_ldr">;
+
 ////////////////////////////////////////////////////////////////////////////////
 // Stores
 
 def SVST1_HOR_ZA8 : MInst<"svst1_hor_za8", "vimiP%", "c", [IsStore, IsOverloadNone, IsStreaming, IsSharedZA, IsPreservesZA], MemEltTyDefault, "aarch64_sme_st1b_horiz">;
 def SVST1_HOR_ZA16 : MInst<"svst1_hor_za16", "vimiP%", "s", [IsStore, IsOverloadNone, IsStreaming, IsSharedZA, IsPreservesZA], MemEltTyDefault, "aarch64_sme_st1h_horiz">;
 def SVST1_HOR_ZA32 : MInst<"svst1_hor_za32", "vimiP%", "i", [IsStore, IsOverloadNone, IsStreaming, IsSharedZA, IsPreservesZA], MemEltTyDefault, "aarch64_sme_st1w_horiz">;
 def SVST1_HOR_ZA64 : MInst<"svst1_hor_za64", "vimiP%", "l", [IsStore, IsOverloadNone, IsStreaming, IsSharedZA, IsPreservesZA], MemEltTyDefault, "aarch64_sme_st1d_horiz">;
 def SVST1_HOR_ZA128 : MInst<"svst1_hor_za128", "vimiP%", "q", [IsStore, IsOverloadNone, IsStreaming, IsSharedZA, IsPreservesZA], MemEltTyDefault, "aarch64_sme_st1q_horiz">;

 def SVST1_HOR_VNUM_ZA64 : MInst<"svst1_hor_vnum_za64", "vimiP%l", "l", [IsStore, IsOverloadNone, IsStreaming, IsSharedZA, IsPreservesZA], MemEltTyDefault, "aarch64_sme_st1d_horiz">;
 def SVST1_HOR_VNUM_ZA128 : MInst<"svst1_hor_vnum_za128", "vimiP%l", "q", [IsStore, IsOverloadNone, IsStreaming, IsSharedZA, IsPreservesZA], MemEltTyDefault, "aarch64_sme_st1q_horiz">;
 def SVST1_VER_VNUM_ZA8 : MInst<"svst1_ver_vnum_za8", "vimiP%l", "c", [IsStore, IsOverloadNone, IsStreaming, IsSharedZA, IsPreservesZA], MemEltTyDefault, "aarch64_sme_st1b_vert">;
 def SVST1_VER_VNUM_ZA16 : MInst<"svst1_ver_vnum_za16", "vimiP%l", "s", [IsStore, IsOverloadNone, IsStreaming, IsSharedZA, IsPreservesZA], MemEltTyDefault, "aarch64_sme_st1h_vert">;
 def SVST1_VER_VNUM_ZA32 : MInst<"svst1_ver_vnum_za32", "vimiP%l", "i", [IsStore, IsOverloadNone, IsStreaming, IsSharedZA, IsPreservesZA], MemEltTyDefault, "aarch64_sme_st1w_vert">;
 def SVST1_VER_VNUM_ZA64 : MInst<"svst1_ver_vnum_za64", "vimiP%l", "l", [IsStore, IsOverloadNone, IsStreaming, IsSharedZA, IsPreservesZA], MemEltTyDefault, "aarch64_sme_st1d_vert">;
 def SVST1_VER_VNUM_ZA128 : MInst<"svst1_ver_vnum_za128", "vimiP%l", "q", [IsStore, IsOverloadNone, IsStreaming, IsSharedZA, IsPreservesZA], MemEltTyDefault, "aarch64_sme_st1q_vert">;
 
+def SVSTR_VNUM_ZA : MInst<"svstr_vnum_za", "vmi%", "", [IsOverloadNone, IsStoreReg, IsStreamingCompatible, IsSharedZA, IsPreservesZA], MemEltTyDefault, "aarch64_sme_str">;
+
 ////////////////////////////////////////////////////////////////////////////////
 // SME - Read horizontal/vertical ZA slices
 
 def SVREAD_HOR_ZA8 : SInst<"svread_hor_za8[_{d}]", "ddPimi", "cUc", MergeOp1, "aarch64_sme_read_horiz", [IsMove, IsStreaming, IsSharedZA, IsPreservesZA]>;
 def SVREAD_HOR_ZA16 : SInst<"svread_hor_za16[_{d}]", "ddPimi", "sUshb", MergeOp1, "aarch64_sme_read_horiz", [IsMove, IsStreaming, IsSharedZA, IsPreservesZA]>;
 def SVREAD_HOR_ZA32 : SInst<"svread_hor_za32[_{d}]", "ddPimi", "iUif", MergeOp1, "aarch64_sme_read_horiz", [IsMove, IsStreaming, IsSharedZA, IsPreservesZA]>;
 def SVREAD_HOR_ZA64 : SInst<"svread_hor_za64[_{d}]", "ddPimi", "lUld", MergeOp1, "aarch64_sme_read_horiz", [IsMove, IsStreaming, IsSharedZA, IsPreservesZA]>;
 def SVREAD_HOR_ZA128 : SInst<"svread_hor_za128[_{d}]", "ddPimi", "csilUcUsUiUlhbfd", MergeOp1, "aarch64_sme_readq_horiz", [IsMove, IsStreaming, IsSharedZA, IsPreservesZA]>;

clang/include/clang/Basic/arm_sve_sme_incl.td

 def IsTupleGet                : FlagType<0x200000000>;
 def IsTupleSet                : FlagType<0x400000000>;
 def IsStreaming               : FlagType<0x800000000>;
 def IsStreamingCompatible     : FlagType<0x1000000000>;
 def IsSharedZA                : FlagType<0x2000000000>;
 def IsPreservesZA             : FlagType<0x4000000000>;
 def IsMove                    : FlagType<0x8000000000>;
 def IsZero                    : FlagType<0x10000000000>;
+def IsLoadReg                 : FlagType<0x20000000000>;
+def IsStoreReg                : FlagType<0x40000000000>;
 
 // These must be kept in sync with the flags in include/clang/Basic/TargetBuiltins.h
 class ImmCheckType<int val> {
   int Value = val;
 }
 def ImmCheck0_31                : ImmCheckType<0>;  // 0..31 (used for e.g. predicate patterns)
 def ImmCheck1_16                : ImmCheckType<1>;  // 1..16
 def ImmCheckExtract             : ImmCheckType<2>;  // 0..(2048/sizeinbits(elt) - 1)

clang/lib/CodeGen/CGBuiltin.cpp

Value *CodeGenFunction::EmitSMEZero(SMETypeFlags TypeFlags,
                                     unsigned IntID) {
   // svzero_za() intrinsic zeros the entire za tile and has no paramters.
   if (Ops.size() == 0)
     Ops.push_back(llvm::ConstantInt::get(Int32Ty, 255));
   Function *F = CGM.getIntrinsic(IntID, {});
   return Builder.CreateCall(F, Ops);
 }
 
+Value *CodeGenFunction::EmitSMELdrStr(SMETypeFlags TypeFlags,
+                                      SmallVectorImpl<Value *> &Ops,
+                                      unsigned IntID) {
+  Function *Vscale = CGM.getIntrinsic(Intrinsic::vscale, Int64Ty);
+  llvm::Value *VscaleCall = Builder.CreateCall(Vscale, {}, "vscale");
+  llvm::Value *MulVL = Builder.CreateMul(
+      VscaleCall,
+      Builder.getInt64(16 * cast<llvm::ConstantInt>(Ops[1])->getZExtValue()),
+      "mulvl");
+  Ops[2] = Builder.CreateGEP(Int8Ty, Ops[2], MulVL);
+  Ops.erase(&Ops[1]);
+  Function *F = CGM.getIntrinsic(IntID, {});
+  return Builder.CreateCall(F, Ops);
+}
+
 // Limit the usage of scalable llvm IR generated by the ACLE by using the
 // sve dup.x intrinsic instead of IRBuilder::CreateVectorSplat.
 Value *CodeGenFunction::EmitSVEDupX(Value *Scalar, llvm::Type *Ty) {
   auto F = CGM.getIntrinsic(Intrinsic::aarch64_sve_dup_x, Ty);
   return Builder.CreateCall(F, Scalar);
 }
 
 Value *CodeGenFunction::EmitSVEDupX(Value* Scalar) {
auto *Builtin = findARMVectorIntrinsicInMap(AArch64SMEIntrinsicMap, BuiltinID,
                                               AArch64SMEIntrinsicsProvenSorted);
   SMETypeFlags TypeFlags(Builtin->TypeModifier);
   if (TypeFlags.isLoad() || TypeFlags.isStore())
     return EmitSMELd1St1(TypeFlags, Ops, Builtin->LLVMIntrinsic);
   else if (TypeFlags.isMove())
     return EmitSMEReadWrite(TypeFlags, Ops, Builtin->LLVMIntrinsic);
   else if (TypeFlags.isZero())
     return EmitSMEZero(TypeFlags, Ops, Builtin->LLVMIntrinsic);
+  else if (TypeFlags.isLoadReg() || TypeFlags.isStoreReg())
+    return EmitSMELdrStr(TypeFlags, Ops, Builtin->LLVMIntrinsic);
 
   /// Should not happen
   return nullptr;
 }
 
 Value *CodeGenFunction::EmitAArch64BuiltinExpr(unsigned BuiltinID,
                                                const CallExpr *E,
                                                llvm::Triple::ArchType Arch) {

clang/lib/CodeGen/CodeGenFunction.h

llvm::Value *EmitSMELd1St1(SMETypeFlags TypeFlags,
                              llvm::SmallVectorImpl<llvm::Value *> &Ops,
                              unsigned IntID);
   llvm::Value *EmitSMEReadWrite(SMETypeFlags TypeFlags,
                                 llvm::SmallVectorImpl<llvm::Value *> &Ops,
                                 unsigned IntID);
   llvm::Value *EmitSMEZero(SMETypeFlags TypeFlags,
                            llvm::SmallVectorImpl<llvm::Value *> &Ops,
                            unsigned IntID);
+  llvm::Value *EmitSMELdrStr(SMETypeFlags TypeFlags,
+                             llvm::SmallVectorImpl<llvm::Value *> &Ops,
+                             unsigned IntID);
   llvm::Value *EmitAArch64SMEBuiltinExpr(unsigned BuiltinID, const CallExpr *E);
 
   llvm::Value *EmitAArch64BuiltinExpr(unsigned BuiltinID, const CallExpr *E,
                                       llvm::Triple::ArchType Arch);
   llvm::Value *EmitBPFBuiltinExpr(unsigned BuiltinID, const CallExpr *E);
 
   llvm::Value *BuildVector(ArrayRef<llvm::Value*> Ops);
   llvm::Value *EmitX86BuiltinExpr(unsigned BuiltinID, const CallExpr *E);

clang/test/CodeGen/aarch64-sme-intrinsics/acle_sme_ldr.c

+// REQUIRES: aarch64-registered-target
+// RUN: %clang_cc1 -D__ARM_FEATURE_SME -triple aarch64-none-linux-gnu -target-feature +sme -target-feature +sve -S -O1 -Werror -emit-llvm -o - %s | FileCheck %s -check-prefixes=CHECK,CHECK-C
+// RUN: %clang_cc1 -D__ARM_FEATURE_SME -triple aarch64-none-linux-gnu -target-feature +sme -target-feature +sve -S -O1 -Werror -emit-llvm -o - -x c++ %s | FileCheck %s -check-prefixes=CHECK,CHECK-CXX
+// RUN: %clang_cc1 -D__ARM_FEATURE_SME -triple aarch64-none-linux-gnu -target-feature +sme -target-feature +sve -S -O1 -Werror -o /dev/null %s
+
+#include <arm_sme.h>
+
+// CHECK-C-LABEL: @test_svldr_vnum_za(
+// CHECK-CXX-LABEL: @_Z18test_svldr_vnum_zajPKv(
+// CHECK-NEXT:  entry:
+// CHECK-NEXT:    tail call void @llvm.aarch64.sme.ldr(i32 [[SLICE_BASE:%.*]], ptr [[PTR:%.*]])
+// CHECK-NEXT:    ret void
+//
+void test_svldr_vnum_za(uint32_t slice_base, const void *ptr) {
+  svldr_vnum_za(slice_base, 0, ptr);
+}
+
+// CHECK-C-LABEL: @test_svldr_vnum_za_1(
+// CHECK-CXX-LABEL: @_Z20test_svldr_vnum_za_1jPKv(
+// CHECK-NEXT:  entry:
+// CHECK-NEXT:    [[VSCALE:%.*]] = tail call i64 @llvm.vscale.i64()
+// CHECK-NEXT:    [[MULVL:%.*]] = mul nuw nsw i64 [[VSCALE]], 240
+// CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i8, ptr [[PTR:%.*]], i64 [[MULVL]]
+// CHECK-NEXT:    tail call void @llvm.aarch64.sme.ldr(i32 [[SLICE_BASE:%.*]], ptr [[TMP0]])
+// CHECK-NEXT:    ret void
+//
+void test_svldr_vnum_za_1(uint32_t slice_base, const void *ptr) {
+  svldr_vnum_za(slice_base, 15, ptr);
+}
+
+// CHECK-C-LABEL: @test_svldr_vnum_za_2(
+// CHECK-CXX-LABEL: @_Z20test_svldr_vnum_za_2jPKv(
+// CHECK-NEXT:  entry:
+// CHECK-NEXT:    [[VSCALE:%.*]] = tail call i64 @llvm.vscale.i64()
+// CHECK-NEXT:    [[MULVL:%.*]] = shl nuw nsw i64 [[VSCALE]], 8
+// CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i8, ptr [[PTR:%.*]], i64 [[MULVL]]
+// CHECK-NEXT:    tail call void @llvm.aarch64.sme.ldr(i32 [[SLICE_BASE:%.*]], ptr [[TMP0]])
+// CHECK-NEXT:    ret void
+//
+void test_svldr_vnum_za_2(uint32_t slice_base, const void *ptr) {
+  svldr_vnum_za(slice_base, 16, ptr);
+}

clang/test/CodeGen/aarch64-sme-intrinsics/acle_sme_str.c

+// REQUIRES: aarch64-registered-target
+// RUN: %clang_cc1 -D__ARM_FEATURE_SME -triple aarch64-none-linux-gnu -target-feature +sme -target-feature +sve -S -O1 -Werror -emit-llvm -o - %s | FileCheck %s -check-prefixes=CHECK,CHECK-C
+// RUN: %clang_cc1 -D__ARM_FEATURE_SME -triple aarch64-none-linux-gnu -target-feature +sme -target-feature +sve -S -O1 -Werror -emit-llvm -o - -x c++ %s | FileCheck %s -check-prefixes=CHECK,CHECK-CXX
+// RUN: %clang_cc1 -D__ARM_FEATURE_SME -triple aarch64-none-linux-gnu -target-feature +sme -target-feature +sve -S -O1 -Werror -o /dev/null %s
+
+#include <arm_sme.h>
+
+
+// CHECK-C-LABEL: @test_svstr_vnum_za(
+// CHECK-CXX-LABEL: @_Z18test_svstr_vnum_zajPv(
+// CHECK-NEXT:  entry:
+// CHECK-NEXT:    tail call void @llvm.aarch64.sme.str(i32 [[SLICE_BASE:%.*]], ptr [[PTR:%.*]])
+// CHECK-NEXT:    ret void
+//
+void test_svstr_vnum_za(uint32_t slice_base, void *ptr) {
+  svstr_vnum_za(slice_base, 0, ptr);
+}
+
+// CHECK-C-LABEL: @test_svstr_vnum_za_1(
+// CHECK-CXX-LABEL: @_Z20test_svstr_vnum_za_1jPv(
+// CHECK-NEXT:  entry:
+// CHECK-NEXT:    [[VSCALE:%.*]] = tail call i64 @llvm.vscale.i64()
+// CHECK-NEXT:    [[MULVL:%.*]] = mul nuw nsw i64 [[VSCALE]], 240
+// CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i8, ptr [[PTR:%.*]], i64 [[MULVL]]
+// CHECK-NEXT:    tail call void @llvm.aarch64.sme.str(i32 [[SLICE_BASE:%.*]], ptr [[TMP0]])
+// CHECK-NEXT:    ret void
+//
+void test_svstr_vnum_za_1(uint32_t slice_base, void *ptr) {
+  svstr_vnum_za(slice_base, 15, ptr);
+}
+
+// CHECK-C-LABEL: @test_svstr_vnum_za_2(
+// CHECK-CXX-LABEL: @_Z20test_svstr_vnum_za_2jPv(
+// CHECK-NEXT:  entry:
+// CHECK-NEXT:    [[VSCALE:%.*]] = tail call i64 @llvm.vscale.i64()
+// CHECK-NEXT:    [[MULVL:%.*]] = shl nuw nsw i64 [[VSCALE]], 8
+// CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i8, ptr [[PTR:%.*]], i64 [[MULVL]]
+// CHECK-NEXT:    tail call void @llvm.aarch64.sme.str(i32 [[SLICE_BASE:%.*]], ptr [[TMP0]])
+// CHECK-NEXT:    ret void
+//
+void test_svstr_vnum_za_2(uint32_t slice_base, void *ptr) {
+  svstr_vnum_za(slice_base, 16, ptr);
+}