From c85776af9592aefb8f65be11a48d26acea8ac5ee Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper@sifive.com>
Date: Wed, 23 Apr 2025 08:46:10 -0700
Subject: [PATCH] [RISCV] Check the extension type for atomic loads in isel
 patterns.

Previously we ignored the extension type and only used the memory
type. The extension type on RISC-V today can only be nonextload,
extload, or sextload. It is ok to treat extload as the same as sextload
so ignoring the extension type is fine.

For #136502, we want to support zextload as well so we will need
to disambiguate based on the extension type.

I wanted to use IsAtomic/IsZeroExtLoad/IsSignExtLoad/IsAnyExtLoad
flags from PatFrags to autogenerate the predicates, but those aren't
hooked up properly in tablegen for ISD::ATOMIC_LOAD. Fixing that
will impact other targets as almost all of them also ignore the
extension type.
---
 llvm/lib/Target/RISCV/RISCVInstrInfoA.td      | 33 +++++++++++++++++--
 llvm/lib/Target/RISCV/RISCVInstrInfoZalasr.td | 26 +++++++++------
 2 files changed, 46 insertions(+), 13 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVInstrInfoA.td b/llvm/lib/Target/RISCV/RISCVInstrInfoA.td
index 0575e17c72287..6600b33d638c3 100644
--- a/llvm/lib/Target/RISCV/RISCVInstrInfoA.td
+++ b/llvm/lib/Target/RISCV/RISCVInstrInfoA.td
@@ -118,6 +118,29 @@ defm AMOMAXU_D  : AMO_rr_aq_rl<0b11100, 0b011, "amomaxu.d">,
 // Pseudo-instructions and codegen patterns
 //===----------------------------------------------------------------------===//
 
+def riscv_atomic_asextload : PatFrag<(ops node:$ptr), (atomic_load node:$ptr), [{
+  ISD::LoadExtType ETy = cast<AtomicSDNode>(N)->getExtensionType();
+  return ETy == ISD::EXTLOAD || ETy == ISD::SEXTLOAD;
+}]>;
+
+def riscv_atomic_asextload_8 : PatFrag<(ops node:$ptr),
+                                       (riscv_atomic_asextload node:$ptr)> {
+  let IsAtomic = true;
+  let MemoryVT = i8;
+}
+
+def riscv_atomic_asextload_16 : PatFrag<(ops node:$ptr),
+                                        (riscv_atomic_asextload node:$ptr)> {
+  let IsAtomic = true;
+  let MemoryVT = i16;
+}
+
+def riscv_atomic_asextload_32 : PatFrag<(ops node:$ptr),
+                                        (riscv_atomic_asextload node:$ptr)> {
+  let IsAtomic = true;
+  let MemoryVT = i32;
+}
+
 let IsAtomic = 1 in {
 // An atomic load operation that does not need either acquire or release
 // semantics.
@@ -165,16 +188,20 @@ class seq_cst_store<PatFrag base>
 // any ordering. This is necessary because AtomicExpandPass has added fences to
 // atomic load/stores and changed them to unordered ones.
 let Predicates = [HasAtomicLdSt] in {
-  def : LdPat<relaxed_load<atomic_load_8>,  LB>;
-  def : LdPat<relaxed_load<atomic_load_16>, LH>;
-  def : LdPat<relaxed_load<atomic_load_32>, LW>;
+  def : LdPat<relaxed_load<riscv_atomic_asextload_8>,  LB>;
+  def : LdPat<relaxed_load<riscv_atomic_asextload_16>, LH>;
 
   def : StPat<relaxed_store<atomic_store_8>,  SB, GPR, XLenVT>;
   def : StPat<relaxed_store<atomic_store_16>, SH, GPR, XLenVT>;
   def : StPat<relaxed_store<atomic_store_32>, SW, GPR, XLenVT>;
 }
 
+let Predicates = [HasAtomicLdSt, IsRV32] in {
+  def : LdPat<relaxed_load<atomic_load_32>, LW>;
+}
+
 let Predicates = [HasAtomicLdSt, IsRV64] in {
+  def : LdPat<relaxed_load<riscv_atomic_asextload_32>, LW>;
   def : LdPat<relaxed_load<atomic_load_64>, LD, i64>;
   def : StPat<relaxed_store<atomic_store_64>, SD, GPR, i64>;
 }
diff --git a/llvm/lib/Target/RISCV/RISCVInstrInfoZalasr.td b/llvm/lib/Target/RISCV/RISCVInstrInfoZalasr.td
index 085353ab88306..f42352d1716b0 100644
--- a/llvm/lib/Target/RISCV/RISCVInstrInfoZalasr.td
+++ b/llvm/lib/Target/RISCV/RISCVInstrInfoZalasr.td
@@ -70,25 +70,22 @@ class PatLAQ<SDPatternOperator OpNode, RVInst Inst, ValueType vt = XLenVT>
 //  while atomic_store has data, addr
 class PatSRL<SDPatternOperator OpNode, RVInst Inst, ValueType vt = XLenVT>
     : Pat<(OpNode (vt GPR:$rs2), (vt GPRMemZeroOffset:$rs1)),
-          (Inst GPRMemZeroOffset:$rs1, GPR:$rs2)>; 
-                              
+          (Inst GPRMemZeroOffset:$rs1, GPR:$rs2)>;
+
 
 let Predicates = [HasStdExtZalasr] in {
   // the sequentially consistent loads use
   //  .aq instead of .aqrl to match the psABI/A.7
-  def : PatLAQ<acquiring_load<atomic_load_8>, LB_AQ>;
-  def : PatLAQ<seq_cst_load<atomic_load_8>, LB_AQ>; 
+  def : PatLAQ<acquiring_load<riscv_atomic_asextload_8>, LB_AQ>;
+  def : PatLAQ<seq_cst_load<riscv_atomic_asextload_8>, LB_AQ>;
 
-  def : PatLAQ<acquiring_load<atomic_load_16>, LH_AQ>;
-  def : PatLAQ<seq_cst_load<atomic_load_16>, LH_AQ>;
-
-  def : PatLAQ<acquiring_load<atomic_load_32>, LW_AQ>;
-  def : PatLAQ<seq_cst_load<atomic_load_32>, LW_AQ>;
+  def : PatLAQ<acquiring_load<riscv_atomic_asextload_16>, LH_AQ>;
+  def : PatLAQ<seq_cst_load<riscv_atomic_asextload_16>, LH_AQ>;
 
   // the sequentially consistent stores use
   //  .rl instead of .aqrl to match the psABI/A.7
   def : PatSRL<releasing_store<atomic_store_8>, SB_RL>;
-  def : PatSRL<seq_cst_store<atomic_store_8>, SB_RL>; 
+  def : PatSRL<seq_cst_store<atomic_store_8>, SB_RL>;
 
   def : PatSRL<releasing_store<atomic_store_16>, SH_RL>;
   def : PatSRL<seq_cst_store<atomic_store_16>, SH_RL>;
@@ -97,7 +94,16 @@ let Predicates = [HasStdExtZalasr] in {
   def : PatSRL<seq_cst_store<atomic_store_32>, SW_RL>;
 } // Predicates = [HasStdExtZalasr]
 
+let Predicates = [HasStdExtZalasr, IsRV32] in {
+  def : PatLAQ<acquiring_load<atomic_load_32>, LW_AQ>;
+  def : PatLAQ<seq_cst_load<atomic_load_32>, LW_AQ>;
+
+} // Predicates = [HasStdExtZalasr, IsRV64]
+
 let Predicates = [HasStdExtZalasr, IsRV64] in {
+  def : PatLAQ<acquiring_load<riscv_atomic_asextload_32>, LW_AQ>;
+  def : PatLAQ<seq_cst_load<riscv_atomic_asextload_32>, LW_AQ>;
+
   def : PatLAQ<acquiring_load<atomic_load_64>, LD_AQ>;
   def : PatLAQ<seq_cst_load<atomic_load_64>, LD_AQ>;