[AArch64] Check FeatureAddrLSLFast to fold base address with scale

[Allen]

Summary

We Split LSLFast into Addr and ALU parts on D157982, then
add the related subtarget features FeatureAddrLSLFast to
check the contant foldable in the shift of the addressing mode.

FeatureAddrLSLFast: Shifts of 1, 2 and 3 (Scales of 2, 4 and 8) are cheap.

[dmgreen]

Can you give more details about what this helps with? We could probably argue both ways about whether these extends are free or not. Thanks

[Allen]

hi @dmgreen, thanks for your attention.
I didn't find that the offset of a memory access instructions can be scaled 16 times, while this attempts to align with the comment for FeatureAddrLSLFast, which hint logical shift number is up to 3.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

       // Get the shift amount based on the scaling factor:
       // log2(sizeof(IdxTy)) - log2(8).
       if (IdxTy->isScalableTy())
         return false;
       uint64_t ShiftAmt =
           llvm::countr_zero(DL.getTypeStoreSizeInBits(IdxTy).getFixedValue()) -
           3;
       // Is the constant foldable in the shift of the addressing mode?
-      // I.e., shift amount is between 1 and 4 inclusive.
-      if (ShiftAmt == 0 || ShiftAmt > 4)
+      // I.e., shift amount is between 1 and 3 inclusive.
+      if (!Subtarget->hasFeature(AArch64::FeatureAddrLSLFast) ||
+          ShiftAmt == 0 || ShiftAmt > 3)
         return false;
       break;
     }
     case Instruction::Trunc:
       // Check if this is a noop.
       // trunc(sext ty1 to ty2) to ty1.
       if (Instr->getType() == Ext->getOperand(0)->getType())
         continue;

llvm/test/CodeGen/AArch64/arm64-codegen-prepare-extload.ll

-; RUN: opt -codegenprepare < %s -mtriple=aarch64-apple-ios -S | FileCheck -enable-var-scope %s --check-prefix=OPTALL --check-prefix=OPT --check-prefix=NONSTRESS
-; RUN: opt -codegenprepare < %s -mtriple=aarch64-apple-ios -S -stress-cgp-ext-ld-promotion | FileCheck -enable-var-scope %s --check-prefix=OPTALL --check-prefix=OPT --check-prefix=STRESS
-; RUN: opt -codegenprepare < %s -mtriple=aarch64-apple-ios -S -disable-cgp-ext-ld-promotion | FileCheck -enable-var-scope %s --check-prefix=OPTALL --check-prefix=DISABLE
+; RUN: opt -codegenprepare < %s -mtriple=aarch64-apple-ios -S -mattr=+addr-lsl-fast | FileCheck -enable-var-scope %s --check-prefix=OPTALL --check-prefix=OPT --check-prefix=NONSTRESS
+; RUN: opt -codegenprepare < %s -mtriple=aarch64-apple-ios -S -stress-cgp-ext-ld-promotion -mattr=+addr-lsl-fast | FileCheck -enable-var-scope %s --check-prefix=OPTALL --check-prefix=OPT --check-prefix=STRESS
+; RUN: opt -codegenprepare < %s -mtriple=aarch64-apple-ios -S -disable-cgp-ext-ld-promotion -mattr=+addr-lsl-fast | FileCheck -enable-var-scope %s --check-prefix=OPTALL --check-prefix=DISABLE
 
 ; CodeGenPrepare should move the zext into the block with the load
 ; so that SelectionDAG can select it with the load.
 ;
 ; OPTALL-LABEL: @foo
 ; OPTALL: [[LD:%[a-zA-Z_0-9-]+]] = load i8, ptr %p
 ; OPTALL-NEXT: [[ZEXT:%[a-zA-Z_0-9-]+]] = zext i8 [[LD]] to i32
 ; OPTALL: store i32 [[ZEXT]], ptr %q

   store i64 %stuff, ptr %staddr
   ret void
 }
 
 ; Check that we do not increase the cost of the code.
 ; The input has one free zext and one free sext. If we would have promoted
 ; all the way through the load we would end up with a free zext and a
 ; non-free sext (of %b).
-; OPTALL-LABEL: @doNotPromoteFreeSExtFromAddrMode128
+; OPTALL-LABEL: @promoteFreeSExtFromAddrMode128
 ; OPTALL: [[LD:%[a-zA-Z_0-9-]+]] = load i8, ptr %p
 ;
 ; STRESS-NEXT: [[ZEXT64:%[a-zA-Z_0-9-]+]] = zext i8 [[LD]] to i64
 ; STRESS-NEXT: [[SEXTB:%[a-zA-Z_0-9-]+]] = sext i32 %b to i64
 ; STRESS-NEXT: [[IDX64:%[a-zA-Z_0-9-]+]] = add nsw i64 [[ZEXT64]], [[SEXTB]]
 ;
-; NONSTRESS-NEXT: [[ZEXT32:%[a-zA-Z_0-9-]+]] = zext i8 [[LD]] to i32
-; NONSTRESS-NEXT: [[RES32:%[a-zA-Z_0-9-]+]] = add nsw i32 [[ZEXT32]], %b
-; NONSTRESS-NEXT: [[IDX64:%[a-zA-Z_0-9-]+]] = sext i32 [[RES32]] to i64
+; NONSTRESS-NEXT: [[ZEXT64:%[a-zA-Z_0-9-]+]] = zext i8 [[LD]] to i64
+; NONSTRESS-NEXT: [[SEXT64:%[a-zA-Z_0-9-]+]] = sext i32 %b to i64
+; NONSTRESS-NEXT: [[IDX64:%[a-zA-Z_0-9-]+]] = add nsw i64 [[ZEXT64]], [[SEXT64]]
 ;
 ; DISABLE-NEXT: [[ZEXT32:%[a-zA-Z_0-9-]+]] = zext i8 [[LD]] to i32
 ; DISABLE-NEXT: [[RES32:%[a-zA-Z_0-9-]+]] = add nsw i32 [[ZEXT32]], %b
 ; DISABLE-NEXT: [[IDX64:%[a-zA-Z_0-9-]+]] = sext i32 [[RES32]] to i64
 ;
 ; OPTALL-NEXT: [[GEP:%[a-zA-Z_0-9-]+]] = getelementptr inbounds i128, ptr %addr, i64 [[IDX64]]
 ; OPTALL-NEXT: store i128 %stuff, ptr [[GEP]]
 ; OPTALL-NEXT: ret void
-define void @doNotPromoteFreeSExtFromAddrMode128(ptr %p, i32 %b, ptr %addr, i128 %stuff) {
+define void @promoteFreeSExtFromAddrMode128(ptr %p, i32 %b, ptr %addr, i128 %stuff) {
 entry:
   %t = load i8, ptr %p
   %zextt = zext i8 %t to i32
   %add = add nsw i32 %zextt, %b
   %idx64 = sext i32 %add to i64
   %staddr = getelementptr inbounds i128, ptr %addr, i64 %idx64
   store i128 %stuff, ptr %staddr
   ret void

llvm/test/Transforms/Inline/AArch64/ext.ll

 ; REQUIRES: asserts
 ; RUN: opt -passes=inline -mtriple=aarch64--linux-gnu -S -debug-only=inline-cost < %s 2>&1 | FileCheck %s
 
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64--linux-gnu"
 
 define i32 @outer1(ptr %ptr, i32 %i) {
   %C = call i32 @inner1(ptr %ptr, i32 %i)
   ret i32 %C
 }
 
 ; sext can be folded into gep.
 ; CHECK: Analyzing call of inner1
-; CHECK: NumInstructionsSimplified: 3
+; CHECK: NumInstructionsSimplified: 2
 ; CHECK: NumInstructions: 4
 define i32 @inner1(ptr %ptr, i32 %i) {
   %E = sext i32 %i to i64
   %G = getelementptr inbounds i32, ptr %ptr, i64 %E
   %L = load i32, ptr %G
   ret i32 %L
 }
 

llvm/test/Transforms/SLPVectorizer/AArch64/ext-trunc.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt -S -passes=slp-vectorizer -mtriple=aarch64--linux-gnu < %s | FileCheck %s
+; RUN: opt -S -passes=slp-vectorizer -mtriple=aarch64--linux-gnu -mattr=+addr-lsl-fast < %s | FileCheck %s
 
 target datalayout = "e-m:e-i32:64-i128:128-n32:64-S128"
 
 declare void @foo(i64, i64, i64, i64)
 
 define void @test1(<4 x i16> %a, <4 x i16> %b, ptr %p) {
 ; Make sure types of sub and its sources are not extended.
 ; CHECK-LABEL: @test1(