[LoongArch] Implement TargetLowering::isLegalAddressingMode() hook

[gonglingqin]

Summary

Use the exact addressing mode information instead of the default to
make better use of offsets in instructions.

[xen0n]

llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp

3181

Isn't k*r+r entirely expressible with alsl in case of k in (2, 4, 8, 16 if LA64)? You may want to implement it this time, but I'm fine either way.

[xen0n]

llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp

3181

Hmm, I haven't dug any deeper but if this predicate only governs actions that end up with a ld/st then please ignore my comment.

[gonglingqin]

llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp

3181

Thank you for your review! As explained in https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/CodeGen/TargetLowering.h#L2600-L2607, this predicate only governs actions that end up with a load/store.

[xen0n]

LGTM with the nit addressed, thanks!

llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp

3153–3155

nit: "reg + 14-bit signed offset left-shifted by 2" could be a lot more concise.

[gonglingqin]

llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp

3153–3155

Thanks! I will modify it.

[gonglingqin]

Address xen0n's comments.

llvm/lib/Target/LoongArch/LoongArchISelLowering.h

Register getRegisterByName(const char *RegName, LLT VT,
                              const MachineFunction &MF) const override;
   bool mayBeEmittedAsTailCall(const CallInst *CI) const override;
 
   bool decomposeMulByConstant(LLVMContext &Context, EVT VT,
                               SDValue C) const override;
 
   bool isUsedByReturnOnly(SDNode *N, SDValue &Chain) const override;
 
+  bool isLegalAddressingMode(const DataLayout &DL, const AddrMode &AM, Type *Ty,
+                             unsigned AS,
+                             Instruction *I = nullptr) const override;
+
 private:
   /// Target-specific function used to lower LoongArch calling conventions.
   typedef bool LoongArchCCAssignFn(const DataLayout &DL, LoongArchABI::ABI ABI,
                                    unsigned ValNo, MVT ValVT,
                                    CCValAssign::LocInfo LocInfo,
                                    ISD::ArgFlagsTy ArgFlags, CCState &State,
                                    bool IsFixed, bool IsReg, Type *OrigTy);
 

llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp

if (auto *ConstNode = dyn_cast<ConstantSDNode>(C.getNode())) {
     const APInt &Imm = ConstNode->getAPIntValue();
     if ((Imm + 1).isPowerOf2() || (Imm - 1).isPowerOf2() ||
         (1 - Imm).isPowerOf2() || (-1 - Imm).isPowerOf2())
       return true;
   }
 
   return false;
 }
+
+bool LoongArchTargetLowering::isLegalAddressingMode(const DataLayout &DL,
+                                                    const AddrMode &AM,
+                                                    Type *Ty, unsigned AS,
+                                                    Instruction *I) const {
+  // LoongArch has four basic addressing modes:
+  //  1. reg
+  //  2. reg + 12-bit signed offset
+  //  3. reg + 14-bit signed offset left-shifted by 2
+  //  4. reg1 + reg2
+  // TODO: Add more checks after support vector extension.
+
+  // No global is ever allowed as a base.
+  if (AM.BaseGV)
+    return false;
+
+  // Require a 12 or 14 bit signed offset.
+  if (!isInt<12>(AM.BaseOffs) || !isShiftedInt<14, 2>(AM.BaseOffs))
+    return false;
+
+  switch (AM.Scale) {
+  case 0:
+    // "i" is not allowed.
+    if (!AM.HasBaseReg)
+      return false;
+    // Otherwise we have "r+i".
+    break;
+  case 1:
+    // "r+r+i" is not allowed.
+    if (AM.HasBaseReg && AM.BaseOffs != 0)
+      return false;
+    // Otherwise we have "r+r" or "r+i".
+    break;
+  case 2:
+    // "2*r+r" or "2*r+i" is not allowed.
+    if (AM.HasBaseReg || AM.BaseOffs)
+      return false;
+    // Otherwise we have "r+r".
+    break;
+  default:
+    return false;
+  }
+
+  return true;
+}

llvm/test/CodeGen/LoongArch/ldptr.ll

entry:
   %val = load i64, ptr %addr, align 8
   ret i64 %val
 }
 
 ;; Check that ldptr.d is emitted for applicable offsets.
 define i64 @ldptr_d(ptr %p) nounwind {
 ; LA32-LABEL: ldptr_d:
 ; LA32:       # %bb.0: # %entry
-; LA32-NEXT:    ori $a1, $zero, 2052
+; LA32-NEXT:    ori $a1, $zero, 2048
 ; LA32-NEXT:    add.w $a1, $a0, $a1
-; LA32-NEXT:    ori $a2, $zero, 2048
-; LA32-NEXT:    add.w $a0, $a0, $a2
-; LA32-NEXT:    ld.w $a0, $a0, 0
-; LA32-NEXT:    ld.w $a1, $a1, 0
+; LA32-NEXT:    ld.w $a0, $a1, 0
+; LA32-NEXT:    ld.w $a1, $a1, 4
 ; LA32-NEXT:    ret
 ;
 ; LA64-LABEL: ldptr_d:
 ; LA64:       # %bb.0: # %entry
 ; LA64-NEXT:    ldptr.d $a0, $a0, 2048
 ; LA64-NEXT:    ret
 entry:
   %addr = getelementptr inbounds i64, ptr %p, i64 256
   %val = load i64, ptr %addr, align 8
   ret i64 %val
 }
 
 ;; Check that ldptr.d is not emitted for out-of-range offsets.
 define i64 @ldptr_d_too_big_offset(ptr %p) nounwind {
 ; LA32-LABEL: ldptr_d_too_big_offset:
 ; LA32:       # %bb.0: # %entry
 ; LA32-NEXT:    lu12i.w $a1, 8
-; LA32-NEXT:    ori $a2, $a1, 4
-; LA32-NEXT:    add.w $a2, $a0, $a2
-; LA32-NEXT:    add.w $a0, $a0, $a1
-; LA32-NEXT:    ld.w $a0, $a0, 0
-; LA32-NEXT:    ld.w $a1, $a2, 0
+; LA32-NEXT:    add.w $a1, $a0, $a1
+; LA32-NEXT:    ld.w $a0, $a1, 0
+; LA32-NEXT:    ld.w $a1, $a1, 4
 ; LA32-NEXT:    ret
 ;
 ; LA64-LABEL: ldptr_d_too_big_offset:
 ; LA64:       # %bb.0: # %entry
 ; LA64-NEXT:    lu12i.w $a1, 8
 ; LA64-NEXT:    ldx.d $a0, $a0, $a1
 ; LA64-NEXT:    ret
 entry:
   %addr = getelementptr inbounds i64, ptr %p, i64 4096
   %val = load i64, ptr %addr, align 8
   ret i64 %val
 }

llvm/test/CodeGen/LoongArch/stptr.ll

; LA64-NEXT:    ret
   store i64 %val, ptr %addr, align 8
   ret void
 }
 
 ;; Check that stptr.d is emitted for applicable offsets.
 define void @stptr_d(ptr %p, i64 %val) nounwind {
 ; LA32-LABEL: stptr_d:
 ; LA32:       # %bb.0:
-; LA32-NEXT:    ori $a3, $zero, 2052
-; LA32-NEXT:    add.w $a3, $a0, $a3
-; LA32-NEXT:    st.w $a2, $a3, 0
-; LA32-NEXT:    ori $a2, $zero, 2048
-; LA32-NEXT:    add.w $a0, $a0, $a2
+; LA32-NEXT:    ori $a3, $zero, 2048
+; LA32-NEXT:    add.w $a0, $a0, $a3
+; LA32-NEXT:    st.w $a2, $a0, 4
 ; LA32-NEXT:    st.w $a1, $a0, 0
 ; LA32-NEXT:    ret
 ;
 ; LA64-LABEL: stptr_d:
 ; LA64:       # %bb.0:
 ; LA64-NEXT:    stptr.d $a1, $a0, 2048
 ; LA64-NEXT:    ret
   %addr = getelementptr inbounds i64, ptr %p, i64 256
   store i64 %val, ptr %addr, align 8
   ret void
 }
 
 ;; Check that stptr.d is not emitted for out-of-range offsets.
 define void @stptr_d_too_big_offset(ptr %p, i64 %val) nounwind {
 ; LA32-LABEL: stptr_d_too_big_offset:
 ; LA32:       # %bb.0:
 ; LA32-NEXT:    lu12i.w $a3, 8
-; LA32-NEXT:    add.w $a4, $a0, $a3
-; LA32-NEXT:    st.w $a1, $a4, 0
-; LA32-NEXT:    ori $a1, $a3, 4
-; LA32-NEXT:    add.w $a0, $a0, $a1
-; LA32-NEXT:    st.w $a2, $a0, 0
+; LA32-NEXT:    add.w $a0, $a0, $a3
+; LA32-NEXT:    st.w $a2, $a0, 4
+; LA32-NEXT:    st.w $a1, $a0, 0
 ; LA32-NEXT:    ret
 ;
 ; LA64-LABEL: stptr_d_too_big_offset:
 ; LA64:       # %bb.0:
 ; LA64-NEXT:    lu12i.w $a2, 8
 ; LA64-NEXT:    stx.d $a1, $a0, $a2
 ; LA64-NEXT:    ret
   %addr = getelementptr inbounds i64, ptr %p, i64 4096
   store i64 %val, ptr %addr, align 8
   ret void
 }