Skip to content

Commit c70c0a8

Browse files
rj-jesusrj-jesus
authored
[AArch64][InstCombine] Combine AES instructions with zero operands. (#142781)
We currently combine (AES (EOR (A, B)), 0) into (AES A, B) for Neon intrinsics when the zero operand appears in the RHS of the AES instruction. This patch extends the combine to support AES SVE intrinsics and the case where the zero operand appears in the LHS of the AES instructions.
1 parent acd264d commit c70c0a8

File tree

2 files changed

+77
-1
lines changed

2 files changed

+77
-1
lines changed

‎llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3076,10 +3076,16 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
30763076
case Intrinsic::arm_neon_aesd:
30773077
case Intrinsic::arm_neon_aese:
30783078
case Intrinsic::aarch64_crypto_aesd:
3079-
case Intrinsic::aarch64_crypto_aese: {
3079+
case Intrinsic::aarch64_crypto_aese:
3080+
case Intrinsic::aarch64_sve_aesd:
3081+
case Intrinsic::aarch64_sve_aese: {
30803082
Value *DataArg = II->getArgOperand(0);
30813083
Value *KeyArg = II->getArgOperand(1);
30823084

3085+
// Accept zero on either operand.
3086+
if (!match(KeyArg, m_ZeroInt()))
3087+
std::swap(KeyArg, DataArg);
3088+
30833089
// Try to use the builtin XOR in AESE and AESD to eliminate a prior XOR
30843090
Value *Data, *Key;
30853091
if (match(KeyArg, m_ZeroInt()) &&

‎llvm/test/Transforms/InstCombine/AArch64/aes-intrinsics.ll‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,17 @@ define <16 x i8> @combineXorAeseZeroARM64(<16 x i8> %data, <16 x i8> %key) {
1313
ret <16 x i8> %data.aes
1414
}
1515

16+
define <16 x i8> @combineXorAeseZeroLhsARM64(<16 x i8> %data, <16 x i8> %key) {
17+
; CHECK-LABEL: define <16 x i8> @combineXorAeseZeroLhsARM64(
18+
; CHECK-SAME: <16 x i8> [[DATA:%.*]], <16 x i8> [[KEY:%.*]]) {
19+
; CHECK-NEXT: [[DATA_AES:%.*]] = tail call <16 x i8> @llvm.aarch64.crypto.aese(<16 x i8> [[DATA]], <16 x i8> [[KEY]])
20+
; CHECK-NEXT: ret <16 x i8> [[DATA_AES]]
21+
;
22+
%data.xor = xor <16 x i8> %data, %key
23+
%data.aes = tail call <16 x i8> @llvm.aarch64.crypto.aese(<16 x i8> zeroinitializer, <16 x i8> %data.xor)
24+
ret <16 x i8> %data.aes
25+
}
26+
1627
define <16 x i8> @combineXorAeseNonZeroARM64(<16 x i8> %data, <16 x i8> %key) {
1728
; CHECK-LABEL: define <16 x i8> @combineXorAeseNonZeroARM64(
1829
; CHECK-SAME: <16 x i8> [[DATA:%.*]], <16 x i8> [[KEY:%.*]]) {
@@ -36,6 +47,17 @@ define <16 x i8> @combineXorAesdZeroARM64(<16 x i8> %data, <16 x i8> %key) {
3647
ret <16 x i8> %data.aes
3748
}
3849

50+
define <16 x i8> @combineXorAesdZeroLhsARM64(<16 x i8> %data, <16 x i8> %key) {
51+
; CHECK-LABEL: define <16 x i8> @combineXorAesdZeroLhsARM64(
52+
; CHECK-SAME: <16 x i8> [[DATA:%.*]], <16 x i8> [[KEY:%.*]]) {
53+
; CHECK-NEXT: [[DATA_AES:%.*]] = tail call <16 x i8> @llvm.aarch64.crypto.aesd(<16 x i8> [[DATA]], <16 x i8> [[KEY]])
54+
; CHECK-NEXT: ret <16 x i8> [[DATA_AES]]
55+
;
56+
%data.xor = xor <16 x i8> %data, %key
57+
%data.aes = tail call <16 x i8> @llvm.aarch64.crypto.aesd(<16 x i8> zeroinitializer, <16 x i8> %data.xor)
58+
ret <16 x i8> %data.aes
59+
}
60+
3961
define <16 x i8> @combineXorAesdNonZeroARM64(<16 x i8> %data, <16 x i8> %key) {
4062
; CHECK-LABEL: define <16 x i8> @combineXorAesdNonZeroARM64(
4163
; CHECK-SAME: <16 x i8> [[DATA:%.*]], <16 x i8> [[KEY:%.*]]) {
@@ -51,3 +73,51 @@ define <16 x i8> @combineXorAesdNonZeroARM64(<16 x i8> %data, <16 x i8> %key) {
5173
declare <16 x i8> @llvm.aarch64.crypto.aese(<16 x i8>, <16 x i8>) #0
5274
declare <16 x i8> @llvm.aarch64.crypto.aesd(<16 x i8>, <16 x i8>) #0
5375

76+
; SVE
77+
78+
define <vscale x 16 x i8> @combineXorAeseZeroLhsSVE(<vscale x 16 x i8> %data, <vscale x 16 x i8> %key) {
79+
; CHECK-LABEL: define <vscale x 16 x i8> @combineXorAeseZeroLhsSVE(
80+
; CHECK-SAME: <vscale x 16 x i8> [[DATA:%.*]], <vscale x 16 x i8> [[KEY:%.*]]) {
81+
; CHECK-NEXT: [[DATA_AES:%.*]] = tail call <vscale x 16 x i8> @llvm.aarch64.sve.aese(<vscale x 16 x i8> [[DATA]], <vscale x 16 x i8> [[KEY]])
82+
; CHECK-NEXT: ret <vscale x 16 x i8> [[DATA_AES]]
83+
;
84+
%data.xor = xor <vscale x 16 x i8> %data, %key
85+
%data.aes = tail call <vscale x 16 x i8> @llvm.aarch64.sve.aese(<vscale x 16 x i8> zeroinitializer, <vscale x 16 x i8> %data.xor)
86+
ret <vscale x 16 x i8> %data.aes
87+
}
88+
89+
define <vscale x 16 x i8> @combineXorAeseZeroRhsSVE(<vscale x 16 x i8> %data, <vscale x 16 x i8> %key) {
90+
; CHECK-LABEL: define <vscale x 16 x i8> @combineXorAeseZeroRhsSVE(
91+
; CHECK-SAME: <vscale x 16 x i8> [[DATA:%.*]], <vscale x 16 x i8> [[KEY:%.*]]) {
92+
; CHECK-NEXT: [[DATA_AES:%.*]] = tail call <vscale x 16 x i8> @llvm.aarch64.sve.aese(<vscale x 16 x i8> [[DATA]], <vscale x 16 x i8> [[KEY]])
93+
; CHECK-NEXT: ret <vscale x 16 x i8> [[DATA_AES]]
94+
;
95+
%data.xor = xor <vscale x 16 x i8> %data, %key
96+
%data.aes = tail call <vscale x 16 x i8> @llvm.aarch64.sve.aese(<vscale x 16 x i8> %data.xor, <vscale x 16 x i8> zeroinitializer)
97+
ret <vscale x 16 x i8> %data.aes
98+
}
99+
100+
define <vscale x 16 x i8> @combineXorAesdZeroLhsSVE(<vscale x 16 x i8> %data, <vscale x 16 x i8> %key) {
101+
; CHECK-LABEL: define <vscale x 16 x i8> @combineXorAesdZeroLhsSVE(
102+
; CHECK-SAME: <vscale x 16 x i8> [[DATA:%.*]], <vscale x 16 x i8> [[KEY:%.*]]) {
103+
; CHECK-NEXT: [[DATA_AES:%.*]] = tail call <vscale x 16 x i8> @llvm.aarch64.sve.aesd(<vscale x 16 x i8> [[DATA]], <vscale x 16 x i8> [[KEY]])
104+
; CHECK-NEXT: ret <vscale x 16 x i8> [[DATA_AES]]
105+
;
106+
%data.xor = xor <vscale x 16 x i8> %data, %key
107+
%data.aes = tail call <vscale x 16 x i8> @llvm.aarch64.sve.aesd(<vscale x 16 x i8> zeroinitializer, <vscale x 16 x i8> %data.xor)
108+
ret <vscale x 16 x i8> %data.aes
109+
}
110+
111+
define <vscale x 16 x i8> @combineXorAesdZeroRhsSVE(<vscale x 16 x i8> %data, <vscale x 16 x i8> %key) {
112+
; CHECK-LABEL: define <vscale x 16 x i8> @combineXorAesdZeroRhsSVE(
113+
; CHECK-SAME: <vscale x 16 x i8> [[DATA:%.*]], <vscale x 16 x i8> [[KEY:%.*]]) {
114+
; CHECK-NEXT: [[DATA_AES:%.*]] = tail call <vscale x 16 x i8> @llvm.aarch64.sve.aesd(<vscale x 16 x i8> [[DATA]], <vscale x 16 x i8> [[KEY]])
115+
; CHECK-NEXT: ret <vscale x 16 x i8> [[DATA_AES]]
116+
;
117+
%data.xor = xor <vscale x 16 x i8> %data, %key
118+
%data.aes = tail call <vscale x 16 x i8> @llvm.aarch64.sve.aesd(<vscale x 16 x i8> %data.xor, <vscale x 16 x i8> zeroinitializer)
119+
ret <vscale x 16 x i8> %data.aes
120+
}
121+
122+
declare <vscale x 16 x i8> @llvm.aarch64.sve.aese(<vscale x 16 x i8>, <vscale x 16 x i8>) #0
123+
declare <vscale x 16 x i8> @llvm.aarch64.sve.aesd(<vscale x 16 x i8>, <vscale x 16 x i8>) #0

0 commit comments

Comments
 (0)