[llvm] a98c294 - [X86] Handle multiple use freeze(undef) in LowerAVXCONCAT_VECTORS as zero vectors (#128830)

[via llvm-commits]

Author: Phoebe Wang
Date: 2025-02-26T18:53:08+08:00
New Revision: a98c2940dbc04bf84de95cb1893694cdcbc4f5fe

URL: https://github.com/llvm/llvm-project/commit/a98c2940dbc04bf84de95cb1893694cdcbc4f5fe
DIFF: https://github.com/llvm/llvm-project/commit/a98c2940dbc04bf84de95cb1893694cdcbc4f5fe.diff

LOG: [X86] Handle multiple use freeze(undef) in LowerAVXCONCAT_VECTORS as zero vectors (#128830)

Follow up of
https://github.com/llvm/llvm-project/commit/ee52af74d8e5e3083cf5195d11c92f8df95b8072
Handles the multiple use come from different vectors:
https://godbolt.org/z/GMb3Endhr

Added: 
    

Modified: 
    llvm/lib/Target/X86/X86ISelLowering.cpp
    llvm/test/CodeGen/X86/avx2-arith.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index 2e88a27bbd2c6..84aaf86550842 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -9486,16 +9486,19 @@ static SDValue LowerAVXCONCAT_VECTORS(SDValue Op, SelectionDAG &DAG,
   unsigned NumZero = 0;
   unsigned NumNonZero = 0;
   unsigned NonZeros = 0;
+  SmallSet<SDValue, 4> Undefs;
   for (unsigned i = 0; i != NumOperands; ++i) {
     SDValue SubVec = Op.getOperand(i);
     if (SubVec.isUndef())
       continue;
     if (ISD::isFreezeUndef(SubVec.getNode())) {
         // If the freeze(undef) has multiple uses then we must fold to zero.
-        if (SubVec.hasOneUse())
+        if (SubVec.hasOneUse()) {
           ++NumFreezeUndef;
-        else
+        } else {
           ++NumZero;
+          Undefs.insert(SubVec);
+        }
     }
     else if (ISD::isBuildVectorAllZeros(SubVec.getNode()))
       ++NumZero;
@@ -9522,6 +9525,11 @@ static SDValue LowerAVXCONCAT_VECTORS(SDValue Op, SelectionDAG &DAG,
                         : (NumFreezeUndef ? DAG.getFreeze(DAG.getUNDEF(ResVT))
                                           : DAG.getUNDEF(ResVT));
 
+  // Replace Undef operands with ZeroVector.
+  for (SDValue U : Undefs)
+    DAG.ReplaceAllUsesWith(
+        U, getZeroVector(U.getSimpleValueType(), Subtarget, DAG, dl));
+
   MVT SubVT = Op.getOperand(0).getSimpleValueType();
   unsigned NumSubElems = SubVT.getVectorNumElements();
   for (unsigned i = 0; i != NumOperands; ++i) {

diff  --git a/llvm/test/CodeGen/X86/avx2-arith.ll b/llvm/test/CodeGen/X86/avx2-arith.ll
index 44ab33ad67f27..70b3b99b46ce9 100644
--- a/llvm/test/CodeGen/X86/avx2-arith.ll
+++ b/llvm/test/CodeGen/X86/avx2-arith.ll
@@ -260,3 +260,31 @@ define <4 x i32> @mul_const11(<4 x i32> %x) {
   %m = mul <4 x i32> %x, <i32 2155905152, i32 2155905152, i32 2155905152, i32 2155905152>
   ret <4 x i32> %m
 }
+
+; check we will zero both vectors.
+define void @multi_freeze(<2 x double> %x, <2 x double> %y) nounwind {
+; X86-LABEL: multi_freeze:
+; X86:       # %bb.0:
+; X86-NEXT:    vmovaps %xmm0, %xmm0
+; X86-NEXT:    vmovaps %xmm1, %xmm1
+; X86-NEXT:    calll foo at PLT
+; X86-NEXT:    vzeroupper
+; X86-NEXT:    retl
+;
+; X64-LABEL: multi_freeze:
+; X64:       # %bb.0:
+; X64-NEXT:    pushq %rax
+; X64-NEXT:    vmovaps %xmm0, %xmm0
+; X64-NEXT:    vmovaps %xmm1, %xmm1
+; X64-NEXT:    callq foo at PLT
+; X64-NEXT:    popq %rax
+; X64-NEXT:    vzeroupper
+; X64-NEXT:    retq
+  %1 = freeze <2 x double> poison
+  %2 = shufflevector <2 x double> %x, <2 x double> %1, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+  %3 = shufflevector <2 x double> %y, <2 x double> %1, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+  call void @foo(<4 x double> %2, <4 x double> %3)
+  ret void
+}
+
+declare void @foo(<4 x double>, <4 x double>)