[llvm] [LoongArch] Fix incorrect logic in isLegalAddressingMode() (PR #88694)

[via llvm-commits]

https://github.com/wangleiat created https://github.com/llvm/llvm-project/pull/88694

This will adress issue:
https://github.com/ClangBuiltLinux/linux/issues/2014

>From b7dd423db9bf56b89e0bbc60b5fc4f5bdd5a9d57 Mon Sep 17 00:00:00 2001
From: wanglei <wanglei at loongson.cn>
Date: Mon, 15 Apr 2024 14:38:28 +0800
Subject: [PATCH] [LoongArch] Fix incorrect logic in isLegalAddressingMode()

This will adress issue:
https://github.com/ClangBuiltLinux/linux/issues/2014
---
 .../LoongArch/LoongArchISelLowering.cpp       | 15 +++---
 llvm/test/CodeGen/LoongArch/gep-imm.ll        | 48 +++++++++++++++++++
 2 files changed, 55 insertions(+), 8 deletions(-)
 create mode 100644 llvm/test/CodeGen/LoongArch/gep-imm.ll

diff --git a/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp b/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp
index c13b10a320f836..285d5c2a63b2da 100644
--- a/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp
+++ b/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp
@@ -4876,20 +4876,19 @@ bool LoongArchTargetLowering::isLegalAddressingMode(const DataLayout &DL,
   if (AM.BaseGV)
     return false;
 
-  // Require a 12 or 14 bit signed offset.
-  if (!isInt<12>(AM.BaseOffs) || !isShiftedInt<14, 2>(AM.BaseOffs))
+  // Require a 12-bit signed offset or 14-bit signed offset left-shifted by 2
+  // with `UAL` feature.
+  if (!isInt<12>(AM.BaseOffs) &&
+      !(isShiftedInt<14, 2>(AM.BaseOffs) && Subtarget.hasUAL()))
     return false;
 
   switch (AM.Scale) {
   case 0:
-    // "i" is not allowed.
-    if (!AM.HasBaseReg)
-      return false;
-    // Otherwise we have "r+i".
+    // "r+i" or just "i", depending on HasBaseReg.
     break;
   case 1:
     // "r+r+i" is not allowed.
-    if (AM.HasBaseReg && AM.BaseOffs != 0)
+    if (AM.HasBaseReg && AM.BaseOffs)
       return false;
     // Otherwise we have "r+r" or "r+i".
     break;
@@ -4897,7 +4896,7 @@ bool LoongArchTargetLowering::isLegalAddressingMode(const DataLayout &DL,
     // "2*r+r" or "2*r+i" is not allowed.
     if (AM.HasBaseReg || AM.BaseOffs)
       return false;
-    // Otherwise we have "r+r".
+    // Allow "2*r" as "r+r".
     break;
   default:
     return false;
diff --git a/llvm/test/CodeGen/LoongArch/gep-imm.ll b/llvm/test/CodeGen/LoongArch/gep-imm.ll
new file mode 100644
index 00000000000000..0eef7e4517f3d8
--- /dev/null
+++ b/llvm/test/CodeGen/LoongArch/gep-imm.ll
@@ -0,0 +1,48 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
+; RUN: llc --mtriple=loongarch64 < %s | FileCheck %s
+
+define void @test(ptr %sp, ptr %t, i32 %n) {
+; CHECK-LABEL: test:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    ld.d $a0, $a0, 0
+; CHECK-NEXT:    move $a3, $zero
+; CHECK-NEXT:    addi.w $a2, $a2, 0
+; CHECK-NEXT:    addi.w $a4, $a3, 0
+; CHECK-NEXT:    bge $a4, $a2, .LBB0_2
+; CHECK-NEXT:    .p2align 4, , 16
+; CHECK-NEXT:  .LBB0_1: # %while_body
+; CHECK-NEXT:    # =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    addi.d $a4, $a3, 1
+; CHECK-NEXT:    stptr.w $a4, $a0, 8000
+; CHECK-NEXT:    stptr.w $a3, $a0, 8004
+; CHECK-NEXT:    stptr.w $a4, $a1, 8000
+; CHECK-NEXT:    stptr.w $a3, $a1, 8004
+; CHECK-NEXT:    move $a3, $a4
+; CHECK-NEXT:    addi.w $a4, $a3, 0
+; CHECK-NEXT:    blt $a4, $a2, .LBB0_1
+; CHECK-NEXT:  .LBB0_2: # %while_end
+; CHECK-NEXT:    ret
+entry:
+  %s = load ptr, ptr %sp
+  br label %while_cond
+
+while_cond:
+  %phi = phi i32 [ 0, %entry ], [ %i, %while_body ]
+  %gep0 = getelementptr [65536 x i32], ptr %s, i64 0, i64 2000
+  %gep1 = getelementptr [65536 x i32], ptr %s, i64 0, i64 2001
+  %gep2 = getelementptr [65536 x i32], ptr %t, i64 0, i64 2000
+  %gep3 = getelementptr [65536 x i32], ptr %t, i64 0, i64 2001
+  %cmp = icmp slt i32 %phi, %n
+  br i1 %cmp, label %while_body, label %while_end
+
+while_body:
+  %i = add i32 %phi, 1
+  store i32 %i, ptr %gep0
+  store i32 %phi, ptr %gep1
+  store i32 %i, ptr %gep2
+  store i32 %phi, ptr %gep3
+  br label %while_cond
+
+while_end:
+  ret void
+}