[llvm] [llvm][LoongArch] Get the code model from the global object (PR #72079)

[via llvm-commits]

https://github.com/heiher updated https://github.com/llvm/llvm-project/pull/72079

>From 9d2f1fd8d56114d904909cbd93ca07a74cdaf4ab Mon Sep 17 00:00:00 2001
From: WANG Rui <wangrui at loongson.cn>
Date: Fri, 10 Nov 2023 21:07:50 -0600
Subject: [PATCH] [llvm][LoongArch] Get the code model from the global object

This gets the code model from global variable if it has, otherwise
the target's will be used.

Signed-off-by: WANG Rui <wangrui at loongson.cn>
---
 .../LoongArch/LoongArchISelLowering.cpp       | 19 +++++---
 .../Target/LoongArch/LoongArchISelLowering.h  |  3 +-
 .../LoongArch/global-variable-code-model.ll   | 44 +++++++++++++++++++
 3 files changed, 60 insertions(+), 6 deletions(-)
 create mode 100644 llvm/test/CodeGen/LoongArch/global-variable-code-model.ll

diff --git a/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp b/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp
index f3f72e74ef085a2..71be018381204ae 100644
--- a/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp
+++ b/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp
@@ -599,12 +599,13 @@ static SDValue getTargetNode(JumpTableSDNode *N, SDLoc DL, EVT Ty,
 
 template <class NodeTy>
 SDValue LoongArchTargetLowering::getAddr(NodeTy *N, SelectionDAG &DAG,
+                                         CodeModel::Model M,
                                          bool IsLocal) const {
   SDLoc DL(N);
   EVT Ty = getPointerTy(DAG.getDataLayout());
   SDValue Addr = getTargetNode(N, DL, Ty, DAG, 0);
 
-  switch (DAG.getTarget().getCodeModel()) {
+  switch (M) {
   default:
     report_fatal_error("Unsupported code model");
 
@@ -645,24 +646,32 @@ SDValue LoongArchTargetLowering::getAddr(NodeTy *N, SelectionDAG &DAG,
 
 SDValue LoongArchTargetLowering::lowerBlockAddress(SDValue Op,
                                                    SelectionDAG &DAG) const {
-  return getAddr(cast<BlockAddressSDNode>(Op), DAG);
+  return getAddr(cast<BlockAddressSDNode>(Op), DAG,
+                 DAG.getTarget().getCodeModel());
 }
 
 SDValue LoongArchTargetLowering::lowerJumpTable(SDValue Op,
                                                 SelectionDAG &DAG) const {
-  return getAddr(cast<JumpTableSDNode>(Op), DAG);
+  return getAddr(cast<JumpTableSDNode>(Op), DAG,
+                 DAG.getTarget().getCodeModel());
 }
 
 SDValue LoongArchTargetLowering::lowerConstantPool(SDValue Op,
                                                    SelectionDAG &DAG) const {
-  return getAddr(cast<ConstantPoolSDNode>(Op), DAG);
+  return getAddr(cast<ConstantPoolSDNode>(Op), DAG,
+                 DAG.getTarget().getCodeModel());
 }
 
 SDValue LoongArchTargetLowering::lowerGlobalAddress(SDValue Op,
                                                     SelectionDAG &DAG) const {
   GlobalAddressSDNode *N = cast<GlobalAddressSDNode>(Op);
   assert(N->getOffset() == 0 && "unexpected offset in global node");
-  return getAddr(N, DAG, N->getGlobal()->isDSOLocal());
+  auto CM = DAG.getTarget().getCodeModel();
+  const GlobalValue *GV = N->getGlobal();
+  const auto *GO = dyn_cast<GlobalObject>(GV);
+  if (GV->isDSOLocal() && GO && GO->hasCodeModel())
+    CM = GO->getCodeModel();
+  return getAddr(N, DAG, CM, GV->isDSOLocal());
 }
 
 SDValue LoongArchTargetLowering::getStaticTLSAddr(GlobalAddressSDNode *N,
diff --git a/llvm/lib/Target/LoongArch/LoongArchISelLowering.h b/llvm/lib/Target/LoongArch/LoongArchISelLowering.h
index 3141286671055dd..d7e28608890ac1a 100644
--- a/llvm/lib/Target/LoongArch/LoongArchISelLowering.h
+++ b/llvm/lib/Target/LoongArch/LoongArchISelLowering.h
@@ -247,7 +247,8 @@ class LoongArchTargetLowering : public TargetLowering {
                          LoongArchCCAssignFn Fn) const;
 
   template <class NodeTy>
-  SDValue getAddr(NodeTy *N, SelectionDAG &DAG, bool IsLocal = true) const;
+  SDValue getAddr(NodeTy *N, SelectionDAG &DAG, CodeModel::Model M,
+                  bool IsLocal = true) const;
   SDValue getStaticTLSAddr(GlobalAddressSDNode *N, SelectionDAG &DAG,
                            unsigned Opc, bool Large = false) const;
   SDValue getDynamicTLSAddr(GlobalAddressSDNode *N, SelectionDAG &DAG,
diff --git a/llvm/test/CodeGen/LoongArch/global-variable-code-model.ll b/llvm/test/CodeGen/LoongArch/global-variable-code-model.ll
new file mode 100644
index 000000000000000..db2feb5c0acdb38
--- /dev/null
+++ b/llvm/test/CodeGen/LoongArch/global-variable-code-model.ll
@@ -0,0 +1,44 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc --mtriple=loongarch64 < %s | FileCheck %s
+
+ at a= external dso_local global i32, code_model "small", align 4
+
+define dso_local signext i32 @local_small() #0 {
+; CHECK-LABEL: local_small:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    pcalau12i $a0, %pc_hi20(a)
+; CHECK-NEXT:    addi.d $a0, $a0, %pc_lo12(a)
+; CHECK-NEXT:    ld.w $a0, $a0, 0
+; CHECK-NEXT:    ret
+  %1 = load i32, ptr @a, align 4
+  ret i32 %1
+}
+
+ at b= external dso_local global i32, code_model "large", align 4
+
+define dso_local signext i32 @local_large() #0 {
+; CHECK-LABEL: local_large:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    pcalau12i $a0, %pc_hi20(b)
+; CHECK-NEXT:    addi.d $a1, $zero, %pc_lo12(b)
+; CHECK-NEXT:    lu32i.d $a1, %pc64_lo20(b)
+; CHECK-NEXT:    lu52i.d $a1, $a1, %pc64_hi12(b)
+; CHECK-NEXT:    add.d $a0, $a1, $a0
+; CHECK-NEXT:    ld.w $a0, $a0, 0
+; CHECK-NEXT:    ret
+  %1 = load i32, ptr @b, align 4
+  ret i32 %1
+}
+
+ at c= external global i32, code_model "large", align 4
+
+define dso_local signext i32 @non_local_large() #0 {
+; CHECK-LABEL: non_local_large:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    pcalau12i $a0, %got_pc_hi20(c)
+; CHECK-NEXT:    ld.d $a0, $a0, %got_pc_lo12(c)
+; CHECK-NEXT:    ld.w $a0, $a0, 0
+; CHECK-NEXT:    ret
+  %1 = load i32, ptr @c, align 4
+  ret i32 %1
+}