[llvm] [SelectionDAG] Use unaligned store to move AVX registers onto stack for `extractelement` (PR #78422)
Manish Kausik H via llvm-commits
llvm-commits at lists.llvm.org
Sun Jan 28 10:27:40 PST 2024
https://github.com/Nirhar updated https://github.com/llvm/llvm-project/pull/78422
>From 0065c31c8bece502dd74b0d5d5aa192693bf2b4f Mon Sep 17 00:00:00 2001
From: Nirhar <hmanishkausik at gmail.com>
Date: Wed, 17 Jan 2024 15:46:06 +0530
Subject: [PATCH] [SelectionDAG] Use unaligned store to move AVX registers onto
stack for `extractelement`
Prior to this patch, SelectionDAG generated aligned move onto stacks for AVX registers
when the function was marked as a no-realign-stack function. This lead to misalignment
between the stack and the instruction generated. This patch fixes the issue.
Fixes #77730
---
llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp | 20 ++++-
.../AArch64/sve-extract-fixed-vector.ll | 77 ++++++++++---------
...igned_extract_from_vector_through_stack.ll | 20 +++++
3 files changed, 80 insertions(+), 37 deletions(-)
create mode 100644 llvm/test/CodeGen/X86/unaligned_extract_from_vector_through_stack.ll
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp
index adfeea073bff65..cea07d0975bb5e 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp
@@ -21,6 +21,7 @@
#include "llvm/Analysis/ConstantFolding.h"
#include "llvm/Analysis/TargetLibraryInfo.h"
#include "llvm/CodeGen/ISDOpcodes.h"
+#include "llvm/CodeGen/MachineFrameInfo.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/MachineJumpTableInfo.h"
#include "llvm/CodeGen/MachineMemOperand.h"
@@ -1377,6 +1378,20 @@ void SelectionDAGLegalize::LegalizeOp(SDNode *Node) {
}
}
+// Helper function that generates an MMO that considers the alignment of the
+// stack, and the size of the stack object
+static MachineMemOperand *getStackAlignedMMO(SDValue StackPtr,
+ MachineFunction &MF) {
+ auto &MFI = MF.getFrameInfo();
+ int FI = cast<FrameIndexSDNode>(StackPtr)->getIndex();
+ MachinePointerInfo PtrInfo = MachinePointerInfo::getFixedStack(MF, FI);
+ MachineMemOperand *MMO =
+ MF.getMachineMemOperand(PtrInfo, MachineMemOperand::MOStore,
+ MFI.getObjectSize(FI), MFI.getObjectAlign(FI));
+
+ return MMO;
+}
+
SDValue SelectionDAGLegalize::ExpandExtractFromVectorThroughStack(SDValue Op) {
SDValue Vec = Op.getOperand(0);
SDValue Idx = Op.getOperand(1);
@@ -1426,8 +1441,9 @@ SDValue SelectionDAGLegalize::ExpandExtractFromVectorThroughStack(SDValue Op) {
if (!Ch.getNode()) {
// Store the value to a temporary stack slot, then LOAD the returned part.
StackPtr = DAG.CreateStackTemporary(VecVT);
- Ch = DAG.getStore(DAG.getEntryNode(), dl, Vec, StackPtr,
- MachinePointerInfo());
+ MachineMemOperand *StoreMMO =
+ getStackAlignedMMO(StackPtr, DAG.getMachineFunction());
+ Ch = DAG.getStore(DAG.getEntryNode(), dl, Vec, StackPtr, StoreMMO);
}
SDValue NewLoad;
diff --git a/llvm/test/CodeGen/AArch64/sve-extract-fixed-vector.ll b/llvm/test/CodeGen/AArch64/sve-extract-fixed-vector.ll
index d2cbbe0628f0f1..aa92c6604f061e 100644
--- a/llvm/test/CodeGen/AArch64/sve-extract-fixed-vector.ll
+++ b/llvm/test/CodeGen/AArch64/sve-extract-fixed-vector.ll
@@ -1,4 +1,4 @@
-; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve,+bf16 < %s | FileCheck %s --check-prefixes=CHECK
; Should codegen to a nop, since idx is zero.
@@ -84,14 +84,15 @@ define <4 x i32> @extract_v4i32_nxv2i32_idx4(<vscale x 2 x i32> %vec) nounwind #
; CHECK: // %bb.0:
; CHECK-NEXT: str x29, [sp, #-16]! // 8-byte Folded Spill
; CHECK-NEXT: addvl sp, sp, #-1
-; CHECK-NEXT: ptrue p0.d
+; CHECK-NEXT: ptrue p0.d, vl4
; CHECK-NEXT: mov x8, #4 // =0x4
; CHECK-NEXT: mov x9, sp
-; CHECK-NEXT: ptrue p1.d, vl4
-; CHECK-NEXT: st1d { z0.d }, p0, [sp]
-; CHECK-NEXT: ld1d { z0.d }, p1/z, [x9, x8, lsl #3]
-; CHECK-NEXT: uzp1 z0.s, z0.s, z0.s
-; CHECK-NEXT: // kill: def $q0 killed $q0 killed $z0
+; CHECK-NEXT: mov z2.d, z0.d
+; CHECK-NEXT: ld1d { z1.d }, p0/z, [x9, x8, lsl #3]
+; CHECK-NEXT: ptrue p0.d
+; CHECK-NEXT: uzp1 z1.s, z1.s, z1.s
+; CHECK-NEXT: st1d { z2.d }, p0, [sp]
+; CHECK-NEXT: mov v0.16b, v1.16b
; CHECK-NEXT: addvl sp, sp, #1
; CHECK-NEXT: ldr x29, [sp], #16 // 8-byte Folded Reload
; CHECK-NEXT: ret
@@ -149,14 +150,15 @@ define <8 x i16> @extract_v8i16_nxv4i16_idx8(<vscale x 4 x i16> %vec) nounwind #
; CHECK: // %bb.0:
; CHECK-NEXT: str x29, [sp, #-16]! // 8-byte Folded Spill
; CHECK-NEXT: addvl sp, sp, #-1
-; CHECK-NEXT: ptrue p0.s
+; CHECK-NEXT: ptrue p0.s, vl8
; CHECK-NEXT: mov x8, #8 // =0x8
; CHECK-NEXT: mov x9, sp
-; CHECK-NEXT: ptrue p1.s, vl8
-; CHECK-NEXT: st1w { z0.s }, p0, [sp]
-; CHECK-NEXT: ld1w { z0.s }, p1/z, [x9, x8, lsl #2]
-; CHECK-NEXT: uzp1 z0.h, z0.h, z0.h
-; CHECK-NEXT: // kill: def $q0 killed $q0 killed $z0
+; CHECK-NEXT: mov z2.d, z0.d
+; CHECK-NEXT: ld1w { z1.s }, p0/z, [x9, x8, lsl #2]
+; CHECK-NEXT: ptrue p0.s
+; CHECK-NEXT: uzp1 z1.h, z1.h, z1.h
+; CHECK-NEXT: st1w { z2.s }, p0, [sp]
+; CHECK-NEXT: mov v0.16b, v1.16b
; CHECK-NEXT: addvl sp, sp, #1
; CHECK-NEXT: ldr x29, [sp], #16 // 8-byte Folded Reload
; CHECK-NEXT: ret
@@ -182,15 +184,16 @@ define <8 x i16> @extract_v8i16_nxv2i16_idx8(<vscale x 2 x i16> %vec) nounwind #
; CHECK: // %bb.0:
; CHECK-NEXT: str x29, [sp, #-16]! // 8-byte Folded Spill
; CHECK-NEXT: addvl sp, sp, #-1
-; CHECK-NEXT: ptrue p0.d
+; CHECK-NEXT: ptrue p0.d, vl8
; CHECK-NEXT: mov x8, #8 // =0x8
; CHECK-NEXT: mov x9, sp
-; CHECK-NEXT: ptrue p1.d, vl8
-; CHECK-NEXT: st1d { z0.d }, p0, [sp]
-; CHECK-NEXT: ld1d { z0.d }, p1/z, [x9, x8, lsl #3]
-; CHECK-NEXT: uzp1 z0.s, z0.s, z0.s
-; CHECK-NEXT: uzp1 z0.h, z0.h, z0.h
-; CHECK-NEXT: // kill: def $q0 killed $q0 killed $z0
+; CHECK-NEXT: mov z2.d, z0.d
+; CHECK-NEXT: ld1d { z1.d }, p0/z, [x9, x8, lsl #3]
+; CHECK-NEXT: ptrue p0.d
+; CHECK-NEXT: uzp1 z1.s, z1.s, z1.s
+; CHECK-NEXT: st1d { z2.d }, p0, [sp]
+; CHECK-NEXT: uzp1 z1.h, z1.h, z1.h
+; CHECK-NEXT: mov v0.16b, v1.16b
; CHECK-NEXT: addvl sp, sp, #1
; CHECK-NEXT: ldr x29, [sp], #16 // 8-byte Folded Reload
; CHECK-NEXT: ret
@@ -247,14 +250,15 @@ define <16 x i8> @extract_v16i8_nxv8i8_idx16(<vscale x 8 x i8> %vec) nounwind #1
; CHECK: // %bb.0:
; CHECK-NEXT: str x29, [sp, #-16]! // 8-byte Folded Spill
; CHECK-NEXT: addvl sp, sp, #-1
-; CHECK-NEXT: ptrue p0.h
+; CHECK-NEXT: ptrue p0.h, vl16
; CHECK-NEXT: mov x8, #16 // =0x10
; CHECK-NEXT: mov x9, sp
-; CHECK-NEXT: ptrue p1.h, vl16
-; CHECK-NEXT: st1h { z0.h }, p0, [sp]
-; CHECK-NEXT: ld1h { z0.h }, p1/z, [x9, x8, lsl #1]
-; CHECK-NEXT: uzp1 z0.b, z0.b, z0.b
-; CHECK-NEXT: // kill: def $q0 killed $q0 killed $z0
+; CHECK-NEXT: mov z2.d, z0.d
+; CHECK-NEXT: ld1h { z1.h }, p0/z, [x9, x8, lsl #1]
+; CHECK-NEXT: ptrue p0.h
+; CHECK-NEXT: uzp1 z1.b, z1.b, z1.b
+; CHECK-NEXT: st1h { z2.h }, p0, [sp]
+; CHECK-NEXT: mov v0.16b, v1.16b
; CHECK-NEXT: addvl sp, sp, #1
; CHECK-NEXT: ldr x29, [sp], #16 // 8-byte Folded Reload
; CHECK-NEXT: ret
@@ -280,15 +284,16 @@ define <16 x i8> @extract_v16i8_nxv4i8_idx16(<vscale x 4 x i8> %vec) nounwind #1
; CHECK: // %bb.0:
; CHECK-NEXT: str x29, [sp, #-16]! // 8-byte Folded Spill
; CHECK-NEXT: addvl sp, sp, #-1
-; CHECK-NEXT: ptrue p0.s
+; CHECK-NEXT: ptrue p0.s, vl16
; CHECK-NEXT: mov x8, #16 // =0x10
; CHECK-NEXT: mov x9, sp
-; CHECK-NEXT: ptrue p1.s, vl16
-; CHECK-NEXT: st1w { z0.s }, p0, [sp]
-; CHECK-NEXT: ld1w { z0.s }, p1/z, [x9, x8, lsl #2]
-; CHECK-NEXT: uzp1 z0.h, z0.h, z0.h
-; CHECK-NEXT: uzp1 z0.b, z0.b, z0.b
-; CHECK-NEXT: // kill: def $q0 killed $q0 killed $z0
+; CHECK-NEXT: mov z2.d, z0.d
+; CHECK-NEXT: ld1w { z1.s }, p0/z, [x9, x8, lsl #2]
+; CHECK-NEXT: ptrue p0.s
+; CHECK-NEXT: uzp1 z1.h, z1.h, z1.h
+; CHECK-NEXT: st1w { z2.s }, p0, [sp]
+; CHECK-NEXT: uzp1 z1.b, z1.b, z1.b
+; CHECK-NEXT: mov v0.16b, v1.16b
; CHECK-NEXT: addvl sp, sp, #1
; CHECK-NEXT: ldr x29, [sp], #16 // 8-byte Folded Reload
; CHECK-NEXT: ret
@@ -437,8 +442,10 @@ define <2 x i64> @extract_fixed_v2i64_nxv2i64(<vscale x 2 x i64> %vec) nounwind
; CHECK-NEXT: str x29, [sp, #-16]! // 8-byte Folded Spill
; CHECK-NEXT: addvl sp, sp, #-1
; CHECK-NEXT: ptrue p0.d
-; CHECK-NEXT: st1d { z0.d }, p0, [sp]
-; CHECK-NEXT: ldr q0, [sp, #16]
+; CHECK-NEXT: mov z2.d, z0.d
+; CHECK-NEXT: ldr q1, [sp, #16]
+; CHECK-NEXT: mov v0.16b, v1.16b
+; CHECK-NEXT: st1d { z2.d }, p0, [sp]
; CHECK-NEXT: addvl sp, sp, #1
; CHECK-NEXT: ldr x29, [sp], #16 // 8-byte Folded Reload
; CHECK-NEXT: ret
diff --git a/llvm/test/CodeGen/X86/unaligned_extract_from_vector_through_stack.ll b/llvm/test/CodeGen/X86/unaligned_extract_from_vector_through_stack.ll
new file mode 100644
index 00000000000000..52d0c2b509128b
--- /dev/null
+++ b/llvm/test/CodeGen/X86/unaligned_extract_from_vector_through_stack.ll
@@ -0,0 +1,20 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu | FileCheck %s
+
+define i32 @foo(i32 %arg1) #0 {
+; CHECK-LABEL: foo:
+; CHECK: # %bb.0: # %entry
+; CHECK-NEXT: # kill: def $edi killed $edi def $rdi
+; CHECK-NEXT: vxorps %xmm0, %xmm0, %xmm0
+; CHECK-NEXT: vmovups %ymm0, -{{[0-9]+}}(%rsp)
+; CHECK-NEXT: andl $31, %edi
+; CHECK-NEXT: movzbl -40(%rsp,%rdi), %eax
+; CHECK-NEXT: vzeroupper
+; CHECK-NEXT: retq
+entry:
+ %a = extractelement <32 x i8> zeroinitializer, i32 %arg1
+ %b = zext i8 %a to i32
+ ret i32 %b
+}
+
+attributes #0 = { "no-realign-stack" "target-cpu"="skylake-avx512" }
More information about the llvm-commits
mailing list