[SLSR] handle (B | i) * S

[jingyue]

Summary

Consider (B | i) * S as (B + i) * S if B and i have no bits set in
common.

[meheff]

LGTM

lib/Transforms/Scalar/StraightLineStrengthReduce.cpp

 // - When (i' - i) is constant but i and i' are not, we could still perform
 //   SLSR.
 #include <vector>
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/FoldingSet.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/Local.h"
void StraightLineStrengthReduce::allocateCandidatesAndFindBasisForAdd(
   } else {
     // At least, I = LHS + 1 * RHS
     ConstantInt *One = ConstantInt::get(cast<IntegerType>(I->getType()), 1);
     allocateCandidatesAndFindBasis(Candidate::Add, SE->getSCEV(LHS), One, RHS,
                                    I);
   }
 }
 
+// Returns true if A matches B + C where C is constant.
+static bool matchesAdd(Value *A, Value *&B, ConstantInt *&C) {
+  return (match(A, m_Add(m_Value(B), m_ConstantInt(C))) ||
+          match(A, m_Add(m_ConstantInt(C), m_Value(B))));
+}
+
+// Returns true if A matches B | C where C is constant.
+static bool matchesOr(Value *A, Value *&B, ConstantInt *&C) {
+  return (match(A, m_Or(m_Value(B), m_ConstantInt(C))) ||
+          match(A, m_Or(m_ConstantInt(C), m_Value(B))));
+}
+
 void StraightLineStrengthReduce::allocateCandidatesAndFindBasisForMul(
     Value *LHS, Value *RHS, Instruction *I) {
   Value *B = nullptr;
   ConstantInt *Idx = nullptr;
-  if (match(LHS, m_Add(m_Value(B), m_ConstantInt(Idx))) ||
-      match(LHS, m_Add(m_ConstantInt(Idx), m_Value(B)))) {
+  if (matchesAdd(LHS, B, Idx)) {
     // If LHS is in the form of "Base + Index", then I is in the form of
     // "(Base + Index) * RHS".
     allocateCandidatesAndFindBasis(Candidate::Mul, SE->getSCEV(B), Idx, RHS, I);
+  } else if (matchesOr(LHS, B, Idx) && haveNoCommonBitsSet(B, Idx, *DL)) {
+    // If LHS is in the form of "Base | Index" and Base and Index have no common
+    // bits set, then
+    //   Base | Index = Base + Index
+    // and I is thus in the form of "(Base + Index) * RHS".
+    allocateCandidatesAndFindBasis(Candidate::Mul, SE->getSCEV(B), Idx, RHS, I);
   } else {
     // Otherwise, at least try the form (LHS + 0) * RHS.
     ConstantInt *Zero = ConstantInt::get(cast<IntegerType>(I->getType()), 0);
     allocateCandidatesAndFindBasis(Candidate::Mul, SE->getSCEV(LHS), Zero, RHS,
                                    I);
   }
 }
 
 void StraightLineStrengthReduce::allocateCandidatesAndFindBasisForMul(
     Instruction *I) {
   // Try matching (B + i) * S.
   // TODO: we could extend SLSR to float and vector types.
   if (!isa<IntegerType>(I->getType()))

test/Transforms/StraightLineStrengthReduce/slsr-mul.ll

; CHECK-NOT: mul i32
   ; foo((2 + b) * s);
   %b2 = add i32 2, %b
   %mul2 = mul i32 %b2, %s
   call void @foo(i32 %mul2)
 
   ret void
 }
 
+define void @or(i32 %a, i32 %s) {
+  %b = shl i32 %a, 1
+; CHECK-LABEL: @or(
+  ; foo(b * s);
+  %mul0 = mul i32 %b, %s
+; CHECK: [[base:[^ ]+]] = mul i32
+  call void @foo(i32 %mul0)
+
+  ; foo((b | 1) * s);
+  %b1 = or i32 %b, 1
+  %mul1 = mul i32 %b1, %s
+; CHECK: add i32 [[base]], %s
+  call void @foo(i32 %mul1)
+
+  ; foo((b | 2) * s);
+  %b2 = or i32 %b, 2
+  %mul2 = mul i32 %b2, %s
+; CHECK: mul i32 %b2, %s
+  call void @foo(i32 %mul2)
+
+  ret void
+}
+
 ; foo(a * b)
 ; foo((a + 1) * b)
 ; foo(a * (b + 1))
 ; foo((a + 1) * (b + 1))
 define void @slsr2(i32 %a, i32 %b) {
 ; CHECK-LABEL: @slsr2(
   %a1 = add i32 %a, 1
   %b1 = add i32 %b, 1