Skip to content

[AArch64][clang][llvm] Add structured sparsity outer product (TMOP) intrinsics #135145

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 15 commits into from
Apr 16, 2025
Merged

[AArch64][clang][llvm] Add structured sparsity outer product (TMOP) intrinsics #135145

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
c0bff6f
[AArch64][clang][llvm] Add structured sparsity outer product (TMOP) i…
jthackray Mar 3, 2025
2624dc6
fixup! Fixup code review issues for TMOP
jthackray Apr 11, 2025
bf4bb10
fixup! Add extra parameter to SME2_ZA_TMOP_Pat for TileOp16 and Vecto…
jthackray Apr 11, 2025
51462fa
fixup! Fix Marian's code review comments
jthackray Apr 11, 2025
66c5ca0
fixup! Remove some range tests, and improve SME2_ZA_TMOP_Pat
jthackray Apr 11, 2025
155ce5d
fixup! Fix remaining CR comments
jthackray Apr 13, 2025
eddca2a
fixup! Numbers in test weren't correctly updated - fix.
jthackray Apr 13, 2025
95a5132
fixup! Bother, another number requiring bit-twiddling
jthackray Apr 13, 2025
25bcd21
fixup! Also set "mayLoad"
jthackray Apr 14, 2025
66ff772
fixup! Move mayLoad/mayStore to reduce code duplication
jthackray Apr 14, 2025
b7bfd9b
fixup! Add FPCR for all FP TMOP instructions
jthackray Apr 15, 2025
8ff3086
fixup! Address Carol's CR comments
jthackray Apr 15, 2025
b6b98f2
fixup! Rename intrinsics
jthackray Apr 15, 2025
ea5d22b
fixup! Fix minor nits
jthackray Apr 16, 2025
8f3e036
fixup! One more intrinsic rename
jthackray Apr 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 42 additions & 0 deletions clang/include/clang/Basic/arm_sme.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -907,6 +907,48 @@ let SMETargetGuard = "sme-f16f16" in {
}


////////////////////////////////////////////////////////////////////////////////
// SME2 - TMOP, SUTMOP, USTMOP

multiclass USTMOP<string za, string t, list<ImmCheck> checks> {
def _ : SInst<"svtmopa_lane_" # za # "[_{d}_{3}]",
"vi2x[i", t, MergeNone, "aarch64_sme_ustmopa",
[IsStreaming, IsInOutZA],
checks>;
}

multiclass SUTMOP<string za, string t, list<ImmCheck> checks> {
def _ : SInst<"svtmopa_lane_" # za # "[_{d}_{3}]",
"vi2u[i", t, MergeNone, "aarch64_sme_sutmopa",
[IsStreaming, IsInOutZA],
checks>;
}

let SMETargetGuard = "sme2,sme-tmop" in {
def SVTMOPA_ZA32 : Inst<"svtmopa_lane_za32[_{d}_{d}]", "vi2d[i", "hbf", MergeNone, "aarch64_sme_tmopa", [IsStreaming, IsInOutZA], [ImmCheck<0, ImmCheck0_3>, ImmCheck<4, ImmCheck0_3>]>;
def SVSTMOPA_ZA32 : Inst<"svtmopa_lane_za32[_{d}_{d}]", "vi2d[i", "cs", MergeNone, "aarch64_sme_stmopa", [IsStreaming, IsInOutZA], [ImmCheck<0, ImmCheck0_3>, ImmCheck<4, ImmCheck0_3>]>;
def SVUTMOPA_ZA32 : Inst<"svtmopa_lane_za32[_{d}_{d}]", "vi2d[i", "UcUs", MergeNone, "aarch64_sme_utmopa", [IsStreaming, IsInOutZA], [ImmCheck<0, ImmCheck0_3>, ImmCheck<4, ImmCheck0_3>]>;

defm SVSUTMOPA_S : SUTMOP<"za32", "c", [ImmCheck<0, ImmCheck0_3>, ImmCheck<4, ImmCheck0_3>]>;
defm SVUSTMOPA_S : USTMOP<"za32", "Uc", [ImmCheck<0, ImmCheck0_3>, ImmCheck<4, ImmCheck0_3>]>;
}

let SMETargetGuard = "sme2,sme-tmop,sme-f16f16" in {
def SVTMOPA_F16 : Inst<"svtmopa_lane_za16[_{d}_{d}]", "vi2d[i", "hb", MergeNone, "aarch64_sme_tmopa", [IsStreaming, IsInOutZA], [ImmCheck<0, ImmCheck0_1>, ImmCheck<4, ImmCheck0_3>]>;
}

let SMETargetGuard = "sme2,sme-tmop,sme-b16b16" in {
def SVTMOPA_BF16 : Inst<"svtmopa_lane_za16[_{d}_{d}]", "vi2d[i", "U", MergeNone, "aarch64_sme_tmopa", [IsStreaming, IsInOutZA], [ImmCheck<0, ImmCheck0_1>, ImmCheck<4, ImmCheck0_3>]>;
}

let SMETargetGuard = "sme2,sme-tmop,sme-f8f16" in {
def SVTMOPA_ZA16_FPM : Inst<"svtmopa_lane_za16[_{d}_{d}]", "vi2.~~[i>", "m", MergeNone, "aarch64_sme_tmopa", [IsStreaming, IsInOutZA], [ImmCheck<0, ImmCheck0_1>, ImmCheck<4, ImmCheck0_3>]>;
}

let SMETargetGuard = "sme2,sme-tmop,sme-f8f32" in {
def SVTMOPA_ZA32_FPM : Inst<"svtmopa_lane_za32[_{d}_{d}]", "vi2.~~[i>", "m", MergeNone, "aarch64_sme_tmopa", [IsStreaming, IsInOutZA], [ImmCheck<0, ImmCheck0_3>, ImmCheck<4, ImmCheck0_3>]>;
}

multiclass ZAReadz<string n_suffix, string vg_num, string t, string i_prefix, list<ImmCheck> ch> {
let SMETargetGuard = "sme2p1" in {
def NAME # _H : SInst<"svreadz_hor_" # n_suffix # "_{d}_vg" # vg_num, vg_num # "im", t,
Expand Down
133 changes: 133 additions & 0 deletions clang/test/CodeGen/AArch64/sme2-intrinsics/acle_sme2_tmop.c
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,133 @@
// NOTE: Assertions have been autogenerated by utils/update_cc_test_checks.py

// REQUIRES: aarch64-registered-target
// RUN: %clang_cc1 -triple aarch64 -target-feature +bf16 -target-feature +sme-tmop -target-feature +sme-f16f16 -target-feature +sme-b16b16 -target-feature +sme-f8f16 -target-feature +sme-f8f32 -target-feature +sme -target-feature +sme2 -disable-O0-optnone -Werror -Wall -emit-llvm -o - %s | opt -S -p mem2reg,instcombine,tailcallelim | FileCheck %s
// RUN: %clang_cc1 -DSME_OVERLOADED_FORMS -triple aarch64 -target-feature +bf16 -target-feature +sme-tmop -target-feature +sme-f16f16 -target-feature +sme-b16b16 -target-feature +sme-f8f16 -target-feature +sme-f8f32 -target-feature +sme -target-feature +sme2 -disable-O0-optnone -Werror -Wall -emit-llvm -o - %s | opt -S -p mem2reg,instcombine,tailcallelim | FileCheck %s
// RUN: %clang_cc1 -triple aarch64 -target-feature +bf16 -target-feature +sme-tmop -target-feature +sme-f16f16 -target-feature +sme-b16b16 -target-feature +sme-f8f16 -target-feature +sme-f8f32 -target-feature +sme -target-feature +sme2 -S -disable-O0-optnone -Werror -Wall -o /dev/null %s

#include <arm_sme.h>

#ifdef SME_OVERLOADED_FORMS
#define SME_ACLE_FUNC(A1,A2_UNUSED,A3) A1##A3
#else
#define SME_ACLE_FUNC(A1,A2,A3) A1##A2##A3
#endif

// CHECK-LABEL: @test_svtmopa_lane_za32_s8_s8(
// CHECK-NEXT: entry:
// CHECK-NEXT: tail call void @llvm.aarch64.sme.stmopa.nxv16i8(i32 1, <vscale x 16 x i8> [[ZN_COERCE0:%.*]], <vscale x 16 x i8> [[ZN_COERCE1:%.*]], <vscale x 16 x i8> [[ZM:%.*]], <vscale x 16 x i8> [[ZK:%.*]], i32 3)
// CHECK-NEXT: ret void
//
void test_svtmopa_lane_za32_s8_s8(svint8x2_t zn, svint8_t zm, svuint8_t zk) __arm_streaming __arm_inout("za") {
SME_ACLE_FUNC(svtmopa_lane_za32,_s8_s8,)(1, zn, zm, zk, 3);
}

// CHECK-LABEL: @test_svtmopa_lane_za32_u8_u8(
// CHECK-NEXT: entry:
// CHECK-NEXT: tail call void @llvm.aarch64.sme.utmopa.nxv16i8(i32 1, <vscale x 16 x i8> [[ZN_COERCE0:%.*]], <vscale x 16 x i8> [[ZN_COERCE1:%.*]], <vscale x 16 x i8> [[ZM:%.*]], <vscale x 16 x i8> [[ZK:%.*]], i32 3)
// CHECK-NEXT: ret void
//
void test_svtmopa_lane_za32_u8_u8(svuint8x2_t zn, svuint8_t zm, svuint8_t zk) __arm_streaming __arm_inout("za") {
SME_ACLE_FUNC(svtmopa_lane_za32,_u8_u8,)(1, zn, zm, zk, 3);
}

// CHECK-LABEL: @test_svtmopa_lane_za32_s8_u8(
// CHECK-NEXT: entry:
// CHECK-NEXT: tail call void @llvm.aarch64.sme.sutmopa.nxv16i8(i32 1, <vscale x 16 x i8> [[ZN_COERCE0:%.*]], <vscale x 16 x i8> [[ZN_COERCE1:%.*]], <vscale x 16 x i8> [[ZM:%.*]], <vscale x 16 x i8> [[ZK:%.*]], i32 3)
// CHECK-NEXT: ret void
//
void test_svtmopa_lane_za32_s8_u8(svint8x2_t zn, svuint8_t zm, svuint8_t zk) __arm_streaming __arm_inout("za") {
SME_ACLE_FUNC(svtmopa_lane_za32,_s8_u8,)(1, zn, zm, zk, 3);
}

// CHECK-LABEL: @test_svtmopa_lane_za32_u8_s8(
// CHECK-NEXT: entry:
// CHECK-NEXT: tail call void @llvm.aarch64.sme.ustmopa.nxv16i8(i32 1, <vscale x 16 x i8> [[ZN_COERCE0:%.*]], <vscale x 16 x i8> [[ZN_COERCE1:%.*]], <vscale x 16 x i8> [[ZM:%.*]], <vscale x 16 x i8> [[ZK:%.*]], i32 3)
// CHECK-NEXT: ret void
//
void test_svtmopa_lane_za32_u8_s8(svuint8x2_t zn, svint8_t zm, svuint8_t zk) __arm_streaming __arm_inout("za") {
SME_ACLE_FUNC(svtmopa_lane_za32,_u8_s8,)(1, zn, zm, zk, 3);
}

// CHECK-LABEL: @test_svtmopa_lane_za32_s16_s16(
// CHECK-NEXT: entry:
// CHECK-NEXT: tail call void @llvm.aarch64.sme.stmopa.nxv8i16(i32 1, <vscale x 8 x i16> [[ZN_COERCE0:%.*]], <vscale x 8 x i16> [[ZN_COERCE1:%.*]], <vscale x 8 x i16> [[ZM:%.*]], <vscale x 16 x i8> [[ZK:%.*]], i32 3)
// CHECK-NEXT: ret void
//
void test_svtmopa_lane_za32_s16_s16(svint16x2_t zn, svint16_t zm, svuint8_t zk) __arm_streaming __arm_inout("za") {
SME_ACLE_FUNC(svtmopa_lane_za32,_s16_s16,)(1, zn, zm, zk, 3);
}

// CHECK-LABEL: @test_svtmopa_lane_za32_u16_u16(
// CHECK-NEXT: entry:
// CHECK-NEXT: tail call void @llvm.aarch64.sme.utmopa.nxv8i16(i32 1, <vscale x 8 x i16> [[ZN_COERCE0:%.*]], <vscale x 8 x i16> [[ZN_COERCE1:%.*]], <vscale x 8 x i16> [[ZM:%.*]], <vscale x 16 x i8> [[ZK:%.*]], i32 3)
// CHECK-NEXT: ret void
//
void test_svtmopa_lane_za32_u16_u16(svuint16x2_t zn, svuint16_t zm, svuint8_t zk) __arm_streaming __arm_inout("za") {
SME_ACLE_FUNC(svtmopa_lane_za32,_u16_u16,)(1, zn, zm, zk, 3);
}

// CHECK-LABEL: @test_svtmopa_lane_za32_f16_f16(
// CHECK-NEXT: entry:
// CHECK-NEXT: tail call void @llvm.aarch64.sme.tmopa.nxv8f16(i32 1, <vscale x 8 x half> [[ZN_COERCE0:%.*]], <vscale x 8 x half> [[ZN_COERCE1:%.*]], <vscale x 8 x half> [[ZM:%.*]], <vscale x 16 x i8> [[ZK:%.*]], i32 3)
// CHECK-NEXT: ret void
//
void test_svtmopa_lane_za32_f16_f16(svfloat16x2_t zn, svfloat16_t zm, svuint8_t zk) __arm_streaming __arm_inout("za") {
SME_ACLE_FUNC(svtmopa_lane_za32,_f16_f16,)(1, zn, zm, zk, 3);
}

// CHECK-LABEL: @test_svtmopa_lane_za32_f32_f32(
// CHECK-NEXT: entry:
// CHECK-NEXT: tail call void @llvm.aarch64.sme.tmopa.nxv4f32(i32 1, <vscale x 4 x float> [[ZN_COERCE0:%.*]], <vscale x 4 x float> [[ZN_COERCE1:%.*]], <vscale x 4 x float> [[ZM:%.*]], <vscale x 16 x i8> [[ZK:%.*]], i32 3)
// CHECK-NEXT: ret void
//
void test_svtmopa_lane_za32_f32_f32(svfloat32x2_t zn, svfloat32_t zm, svuint8_t zk) __arm_streaming __arm_inout("za") {
SME_ACLE_FUNC(svtmopa_lane_za32,_f32_f32,)(1, zn, zm, zk, 3);
}

// CHECK-LABEL: @test_svtmopa_lane_za32_bf16_bf16(
// CHECK-NEXT: entry:
// CHECK-NEXT: tail call void @llvm.aarch64.sme.tmopa.nxv8bf16(i32 1, <vscale x 8 x bfloat> [[ZN_COERCE0:%.*]], <vscale x 8 x bfloat> [[ZN_COERCE1:%.*]], <vscale x 8 x bfloat> [[ZM:%.*]], <vscale x 16 x i8> [[ZK:%.*]], i32 3)
// CHECK-NEXT: ret void
//
void test_svtmopa_lane_za32_bf16_bf16(svbfloat16x2_t zn, svbfloat16_t zm, svuint8_t zk) __arm_streaming __arm_inout("za") {
SME_ACLE_FUNC(svtmopa_lane_za32,_bf16_bf16,)(1, zn, zm, zk, 3);
}

// CHECK-LABEL: @test_svtmopa_lane_za16_f16_f16(
// CHECK-NEXT: entry:
// CHECK-NEXT: tail call void @llvm.aarch64.sme.tmopa.nxv8f16(i32 1, <vscale x 8 x half> [[ZN_COERCE0:%.*]], <vscale x 8 x half> [[ZN_COERCE1:%.*]], <vscale x 8 x half> [[ZM:%.*]], <vscale x 16 x i8> [[ZK:%.*]], i32 3)
// CHECK-NEXT: ret void
//
void test_svtmopa_lane_za16_f16_f16(svfloat16x2_t zn, svfloat16_t zm, svuint8_t zk) __arm_streaming __arm_inout("za") {
SME_ACLE_FUNC(svtmopa_lane_za16,_f16_f16,)(1, zn, zm, zk, 3);
}

// CHECK-LABEL: @test_svtmopa_lane_za16_bf16_bf16(
// CHECK-NEXT: entry:
// CHECK-NEXT: tail call void @llvm.aarch64.sme.tmopa.nxv8bf16(i32 1, <vscale x 8 x bfloat> [[ZN_COERCE0:%.*]], <vscale x 8 x bfloat> [[ZN_COERCE1:%.*]], <vscale x 8 x bfloat> [[ZM:%.*]], <vscale x 16 x i8> [[ZK:%.*]], i32 3)
// CHECK-NEXT: ret void
//
void test_svtmopa_lane_za16_bf16_bf16(svbfloat16x2_t zn, svbfloat16_t zm, svuint8_t zk) __arm_streaming __arm_inout("za") {
SME_ACLE_FUNC(svtmopa_lane_za16,_bf16_bf16,)(1, zn, zm, zk, 3);
}

// CHECK-LABEL: @test_svtmopa_lane_za16_mf8_mf8_fpm(
// CHECK-NEXT: entry:
// CHECK-NEXT: tail call void @llvm.aarch64.set.fpmr(i64 [[FPMR:%.*]])
// CHECK-NEXT: tail call void @llvm.aarch64.sme.tmopa.nxv16i8(i32 1, <vscale x 16 x i8> [[ZN_COERCE0:%.*]], <vscale x 16 x i8> [[ZN_COERCE1:%.*]], <vscale x 16 x i8> [[ZM:%.*]], <vscale x 16 x i8> [[ZK:%.*]], i32 3)
// CHECK-NEXT: ret void
//
void test_svtmopa_lane_za16_mf8_mf8_fpm(svmfloat8x2_t zn, svmfloat8_t zm, svuint8_t zk, fpm_t fpmr) __arm_streaming __arm_inout("za") {
SME_ACLE_FUNC(svtmopa_lane_za16,_mf8_mf8,_fpm)(1, zn, zm, zk, 3, fpmr);
}

// CHECK-LABEL: @test_svtmopa_lane_za32_mf8_mf8_fpm(
// CHECK-NEXT: entry:
// CHECK-NEXT: tail call void @llvm.aarch64.set.fpmr(i64 [[FPMR:%.*]])
// CHECK-NEXT: tail call void @llvm.aarch64.sme.tmopa.nxv16i8(i32 1, <vscale x 16 x i8> [[ZN_COERCE0:%.*]], <vscale x 16 x i8> [[ZN_COERCE1:%.*]], <vscale x 16 x i8> [[ZM:%.*]], <vscale x 16 x i8> [[ZK:%.*]], i32 3)
// CHECK-NEXT: ret void
//
void test_svtmopa_lane_za32_mf8_mf8_fpm(svmfloat8x2_t zn, svmfloat8_t zm, svuint8_t zk, fpm_t fpmr) __arm_streaming __arm_inout("za") {
SME_ACLE_FUNC(svtmopa_lane_za32,_mf8_mf8,_fpm)(1, zn, zm, zk, 3, fpmr);
}
Loading