[SLSR] consider &B[S << i] as &B[(1 << i) * S]

[jingyue]

Summary

This reduces handling &B[(1 << i) * s] to handling &B[i * S].

[sanjoy]

Super-pedantic nit on the comment message: I think it should be "This reduces handling &B[(i << 1) * s] to handling &B[i * S]."

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

372 ↗

(On Diff #23263)

A << B is nsw does not imply A * (1 << B) is nsw. For i8 a counterexample is A = -1 and B = 7 -- -1 << 7 does not sign-overflow but -1 * -128 sign-overflows. In general, for iN I think the only counterexample has B = N-1; if that can be proved (I haven't tried) then it should be sufficient to guard for just that.

[sanjoy]

Please disregard my last comment on the nsw bit. I checked, the langref defines nsw in terms of the multiplication that would be equivalent to the shift, and your code is correct as it is.

In case you're curious, I was going by the following definition of sign-overflow -- A OP B sign-overflows if "sext(A) OP sext(B) != sext(A OP B)". That definition does not really make sense in the context of left shifts.

[jingyue]

A minor fix to the commit message. Thanks, Sanjoy!

[meheff]

LGTM. Thanks!

llvm/trunk/lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

void StraightLineStrengthReduce::factorArrayIndex(Value *ArrayIdx,
                                                   uint64_t ElementSize,
                                                   GetElementPtrInst *GEP) {
   // At least, ArrayIdx = ArrayIdx *s 1.
   allocateCandidateAndFindBasisForGEP(
       Base, ConstantInt::get(cast<IntegerType>(ArrayIdx->getType()), 1),
       ArrayIdx, ElementSize, GEP);
   Value *LHS = nullptr;
   ConstantInt *RHS = nullptr;
-  // TODO: handle shl. e.g., we could treat (S << 2) as (S * 4).
-  //
   // One alternative is matching the SCEV of ArrayIdx instead of ArrayIdx
   // itself. This would allow us to handle the shl case for free. However,
   // matching SCEVs has two issues:
   //
   // 1. this would complicate rewriting because the rewriting procedure
   // would have to translate SCEVs back to IR instructions. This translation
   // is difficult when LHS is further evaluated to a composite SCEV.
   //
   // 2. ScalarEvolution is designed to be control-flow oblivious. It tends
   // to strip nsw/nuw flags which are critical for SLSR to trace into
   // sext'ed multiplication.
   if (match(ArrayIdx, m_NSWMul(m_Value(LHS), m_ConstantInt(RHS)))) {
     // SLSR is currently unsafe if i * S may overflow.
     // GEP = Base + sext(LHS *nsw RHS) * ElementSize
     allocateCandidateAndFindBasisForGEP(Base, RHS, LHS, ElementSize, GEP);
+  } else if (match(ArrayIdx, m_NSWShl(m_Value(LHS), m_ConstantInt(RHS)))) {
+    // GEP = Base + sext(LHS <<nsw RHS) * ElementSize
+    //     = Base + sext(LHS *nsw (1 << RHS)) * ElementSize
+    APInt One(RHS->getBitWidth(), 1);
+    ConstantInt *PowerOf2 =
+        ConstantInt::get(RHS->getContext(), One << RHS->getValue());
+    allocateCandidateAndFindBasisForGEP(Base, PowerOf2, LHS, ElementSize, GEP);
   }
 }
 
 void StraightLineStrengthReduce::allocateCandidateAndFindBasisForGEP(
     GetElementPtrInst *GEP) {
   // TODO: handle vector GEPs
   if (GEP->getType()->isVectorTy())
     return;

llvm/trunk/test/Transforms/StraightLineStrengthReduce/slsr-gep.ll

 ; RUN: opt < %s -slsr -gvn -dce -S | FileCheck %s
 
 target datalayout = "e-i64:64-v16:16-v32:32-n16:32:64"
 
 define i32 @slsr_gep(i32* %input, i64 %s) {
 ; CHECK-LABEL: @slsr_gep(
   ; v0 = input[0];
   %p0 = getelementptr inbounds i32, i32* %input, i64 0
   %v0 = load i32, i32* %p0
 
   ; v1 = input[s];
   %p1 = getelementptr inbounds i32, i32* %input, i64 %s
 ; CHECK: %p1 = getelementptr inbounds i32, i32* %input, i64 %s
   %v1 = load i32, i32* %p1
 
   ; v2 = input[s * 2];
-  %s2 = mul nsw i64 %s, 2
+  %s2 = shl nsw i64 %s, 1
   %p2 = getelementptr inbounds i32, i32* %input, i64 %s2
 ; CHECK: %p2 = getelementptr inbounds i32, i32* %p1, i64 %s
   %v2 = load i32, i32* %p2
 
   ; return v0 + v1 + v2;
   %1 = add i32 %v0, %v1
   %2 = add i32 %1, %v2
   ret i32 %2
 }
 
 define i32 @slsr_gep_sext(i32* %input, i32 %s) {
 ; CHECK-LABEL: @slsr_gep_sext(
   ; v0 = input[0];
   %p0 = getelementptr inbounds i32, i32* %input, i64 0
   %v0 = load i32, i32* %p0
 
   ; v1 = input[(long)s];
   %t = sext i32 %s to i64
   %p1 = getelementptr inbounds i32, i32* %input, i64 %t
 ; CHECK: %p1 = getelementptr inbounds i32, i32* %input, i64 %t
   %v1 = load i32, i32* %p1
 
   ; v2 = input[(long)(s * 2)];
-  %s2 = mul nsw i32 %s, 2
+  %s2 = shl nsw i32 %s, 1
   %t2 = sext i32 %s2 to i64
   %p2 = getelementptr inbounds i32, i32* %input, i64 %t2
 ; CHECK: %p2 = getelementptr inbounds i32, i32* %p1, i64 %t
   %v2 = load i32, i32* %p2
 
   ; return v0 + v1 + v2;
   %1 = add i32 %v0, %v1
   %2 = add i32 %1, %v2
   ret i32 %2
 }
 
 define i32 @slsr_gep_2d([10 x [5 x i32]]* %input, i64 %s, i64 %t) {
 ; CHECK-LABEL: @slsr_gep_2d(
   ; v0 = input[s][t];
   %p0 = getelementptr inbounds [10 x [5 x i32]], [10 x [5 x i32]]* %input, i64 0, i64 %s, i64 %t
   %v0 = load i32, i32* %p0
 
   ; v1 = input[s * 2][t];
-  %s2 = mul nsw i64 %s, 2
+  %s2 = shl nsw i64 %s, 1
 ; CHECK: [[BUMP:%[a-zA-Z0-9]+]] = mul i64 %s, 5
   %p1 = getelementptr inbounds [10 x [5 x i32]], [10 x [5 x i32]]* %input, i64 0, i64 %s2, i64 %t
 ; CHECK: %p1 = getelementptr inbounds i32, i32* %p0, i64 [[BUMP]]
   %v1 = load i32, i32* %p1
 
   ; v2 = input[s * 3][t];
   %s3 = mul nsw i64 %s, 3
   %p2 = getelementptr inbounds [10 x [5 x i32]], [10 x [5 x i32]]* %input, i64 0, i64 %s3, i64 %t

 ; @slsr_gep_2d.
 define i64 @slsr_gep_uglygep([10 x [5 x %struct.S]]* %input, i64 %s, i64 %t) {
 ; CHECK-LABEL: @slsr_gep_uglygep(
   ; v0 = input[s][t].f1;
   %p0 = getelementptr inbounds [10 x [5 x %struct.S]], [10 x [5 x %struct.S]]* %input, i64 0, i64 %s, i64 %t, i32 0
   %v0 = load i64, i64* %p0
 
   ; v1 = input[s * 2][t].f1;
-  %s2 = mul nsw i64 %s, 2
+  %s2 = shl nsw i64 %s, 1
 ; CHECK: [[BUMP:%[a-zA-Z0-9]+]] = mul i64 %s, 60
   %p1 = getelementptr inbounds [10 x [5 x %struct.S]], [10 x [5 x %struct.S]]* %input, i64 0, i64 %s2, i64 %t, i32 0
 ; CHECK: getelementptr inbounds i8, i8* %{{[0-9]+}}, i64 [[BUMP]]
   %v1 = load i64, i64* %p1
 
   ; v2 = input[s * 3][t].f1;
   %s3 = mul nsw i64 %s, 3
   %p2 = getelementptr inbounds [10 x [5 x %struct.S]], [10 x [5 x %struct.S]]* %input, i64 0, i64 %s3, i64 %t, i32 0
; CHECK-LABEL: @slsr_out_of_bounds_gep(
 
   ; v1 = input[(long)s];
   %t = sext i32 %s to i64
   %p1 = getelementptr i32, i32* %input, i64 %t
 ; CHECK: %p1 = getelementptr i32, i32* %input, i64 %t
   %v1 = load i32, i32* %p1
 
   ; v2 = input[(long)(s * 2)];
-  %s2 = mul nsw i32 %s, 2
+  %s2 = shl nsw i32 %s, 1
   %t2 = sext i32 %s2 to i64
   %p2 = getelementptr i32, i32* %input, i64 %t2
 ; CHECK: %p2 = getelementptr i32, i32* %p1, i64 %t
   %v2 = load i32, i32* %p2
 
   ; return v0 + v1 + v2;
   %1 = add i32 %v0, %v1
   %2 = add i32 %1, %v2
   ret i32 %2
 }