[llvm] [X86] Fold vpextrq $1 from a spilled vector into a direct memory load (PR #203339)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-x86

Author: Ye Tian (TianYe717)

<details>
<summary>Changes</summary>

Fixes https://github.com/llvm/llvm-project/issues/199181

Before: `-O3 -mcpu=emeraldrapids` generates `vmovdqa (%rsp), %xmm0` + `vpextrq $1, %xmm0, %rax` when extracting lane 1 from a spilled v2i64 vector.

After: The reload+extract is folded into a single `movq 8(%rsp), %rax` .

Fix: Added a custom fold in `foldMemoryOperandCustom` for `VPEXTRQrri` / `VPEXTRQZrri` / `PEXTRQrri` . When extracting lane 1 from a spilled stack slot (OpNum == 1, imm == 1), emit `MOV64rm` with an 8-byte offset instead.

---
Full diff: https://github.com/llvm/llvm-project/pull/203339.diff


2 Files Affected:

- (modified) llvm/lib/Target/X86/X86InstrInfo.cpp (+15) 
- (added) llvm/test/CodeGen/X86/vpextrq-fold.ll (+68) 


``````````diff
diff --git a/llvm/lib/Target/X86/X86InstrInfo.cpp b/llvm/lib/Target/X86/X86InstrInfo.cpp
index 15d2e10aa0f08..ac458904146db 100644
--- a/llvm/lib/Target/X86/X86InstrInfo.cpp
+++ b/llvm/lib/Target/X86/X86InstrInfo.cpp
@@ -7432,6 +7432,21 @@ MachineInstr *X86InstrInfo::foldMemoryOperandCustom(
                        InsertPt, MI))
       return NewMI;
     break;
+  case X86::VPEXTRQrri:
+  case X86::VPEXTRQZrri:
+  case X86::PEXTRQrri:
+    // Fold: extractelt(v2i64 vector, 1) where vector is a spilled stack slot.
+    // Instead of reloading the full vector and extracting with vpextrq, load
+    // the upper 8 bytes directly from the spill slot at offset 8.
+    if (OpNum == 1 && MI.getOperand(2).getImm() == 1 &&
+        Alignment >= Align(8)) {
+      MachineInstrBuilder MIB = BuildMI(*InsertPt->getParent(), InsertPt,
+                                        MI.getDebugLoc(), get(X86::MOV64rm));
+      MIB.add(MI.getOperand(0)); // dst register
+      addOperands(MIB, MOs, 8);  // offset 8 for upper lane
+      return MIB;
+    }
+    break;
   }
 
   return nullptr;
diff --git a/llvm/test/CodeGen/X86/vpextrq-fold.ll b/llvm/test/CodeGen/X86/vpextrq-fold.ll
new file mode 100644
index 0000000000000..197aebb18cf5a
--- /dev/null
+++ b/llvm/test/CodeGen/X86/vpextrq-fold.ll
@@ -0,0 +1,68 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -mcpu=emeraldrapids -O3 %s -o - | FileCheck %s
+
+; Test that extractelement from a spilled v2i64 vector lane 1 is folded into a
+; direct memory load from the spill slot at offset 8, instead of reloading the
+; full vector and extracting with vpextrq.
+
+declare <2 x i64> @llvm.masked.load.v2i64.p0(ptr, i32 immarg, <2 x i1>, <2 x i64>)
+declare void @clobber()
+declare i64 @llvm.fshl.i64(i64, i64, i64)
+
+define void @repro(ptr %src, ptr %dst, i64 %sh) {
+; CHECK-LABEL: repro:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    pushq %r15
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    pushq %r14
+; CHECK-NEXT:    .cfi_def_cfa_offset 24
+; CHECK-NEXT:    pushq %rbx
+; CHECK-NEXT:    .cfi_def_cfa_offset 32
+; CHECK-NEXT:    subq $16, %rsp
+; CHECK-NEXT:    .cfi_def_cfa_offset 48
+; CHECK-NEXT:    .cfi_offset %rbx, -32
+; CHECK-NEXT:    .cfi_offset %r14, -24
+; CHECK-NEXT:    .cfi_offset %r15, -16
+; CHECK-NEXT:    movq %rdx, %rbx
+; CHECK-NEXT:    movq %rsi, %r14
+; CHECK-NEXT:    movl $3, %eax
+; CHECK-NEXT:    kmovd %eax, %k1
+; CHECK-NEXT:    vmovdqu64 (%rdi), %xmm0 {%k1} {z}
+; CHECK-NEXT:    vmovdqa %xmm0, (%rsp) # 16-byte Spill
+; CHECK-NEXT:    vmovq %xmm0, %r15
+; CHECK-NEXT:    vmovq %xmm0, (%rsi)
+; CHECK-NEXT:    callq clobber at PLT
+; CHECK-NEXT:    movq {{[-0-9]+}}(%r{{[sb]}}p), %rax # 16-byte Reload
+; CHECK-NEXT:    movl %ebx, %ecx
+; CHECK-NEXT:    shldq %cl, %r15, %rax
+; CHECK-NEXT:    movq %rax, 8(%r14)
+; CHECK-NEXT:    addq $16, %rsp
+; CHECK-NEXT:    .cfi_def_cfa_offset 32
+; CHECK-NEXT:    popq %rbx
+; CHECK-NEXT:    .cfi_def_cfa_offset 24
+; CHECK-NEXT:    popq %r14
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    popq %r15
+; CHECK-NEXT:    .cfi_def_cfa_offset 8
+; CHECK-NEXT:    retq
+entry:
+  %mask32 = bitcast i32 3 to <32 x i1>
+  %mask = shufflevector <32 x i1> %mask32, <32 x i1> poison, <2 x i32> <i32 0, i32 1>
+
+  %vec = call <2 x i64> @llvm.masked.load.v2i64.p0(
+      ptr align 8 %src,
+      i32 8,
+      <2 x i1> %mask,
+      <2 x i64> zeroinitializer)
+
+  %lo = extractelement <2 x i64> %vec, i64 0
+  store i64 %lo, ptr %dst, align 8
+
+  call void @clobber()
+
+  %hi = extractelement <2 x i64> %vec, i64 1
+  %r = call i64 @llvm.fshl.i64(i64 %hi, i64 %lo, i64 %sh)
+  %out = getelementptr i8, ptr %dst, i64 8
+  store i64 %r, ptr %out, align 8
+  ret void
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/203339