[llvm] [X86] Don't fold non-temporal loads if we have an instruction for them (PR #86839)

[via llvm-commits]

https://github.com/AtariDreams created https://github.com/llvm/llvm-project/pull/86839

None

>From a4b04e4c480f85df074fc95ac9277b3c942eb5db Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Wed, 27 Mar 2024 13:33:36 -0400
Subject: [PATCH] [X86] Don't fold non-temporal loads if we have an instruction
 for them

---
 llvm/lib/Target/X86/X86ISelLowering.cpp | 31 +++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index 9d98d31b31df0b..a9e707cb0857cd 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -2586,6 +2586,32 @@ X86TargetLowering::createFastISel(FunctionLoweringInfo &funcInfo,
 //                           Other Lowering Hooks
 //===----------------------------------------------------------------------===//
 
+// Indicates we should prefer to use a non-temporal load for this load.
+// TODO: Should be merged with useNonTemporalLoad in X86ISelDAGToDAG.cpp
+static bool useNonTemporalLoad(LoadSDNode *N, const X86Subtarget &Subtarget) {
+  if (!N->isNonTemporal())
+    return false;
+
+  unsigned StoreSize = N->getMemoryVT().getStoreSize();
+
+  if (N->getAlign().value() < StoreSize)
+    return false;
+
+  switch (StoreSize) {
+  default:
+    llvm_unreachable("Unsupported store size");
+  case 4:
+  case 8:
+    return false;
+  case 16:
+    return Subtarget.hasSSE41();
+  case 32:
+    return Subtarget.hasAVX2();
+  case 64:
+    return Subtarget.hasAVX512();
+  }
+}
+
 bool X86::mayFoldLoad(SDValue Op, const X86Subtarget &Subtarget,
                       bool AssumeSingleUse) {
   if (!AssumeSingleUse && !Op.hasOneUse())
@@ -2599,8 +2625,9 @@ bool X86::mayFoldLoad(SDValue Op, const X86Subtarget &Subtarget,
       Ld->getValueSizeInBits(0) == 128 && Ld->getAlign() < Align(16))
     return false;
 
-  // TODO: If this is a non-temporal load and the target has an instruction
-  //       for it, it should not be folded. See "useNonTemporalLoad()".
+  // Don't fold non-temporal loads if we have an instruction for them.
+  if (useNonTemporalLoad(Ld, Subtarget))
+    return false;
 
   return true;
 }