[VectorCombine] Support AND/UREM indices that require freezing.

[fhahn]

Summary

38b098be6605 limited scalarization to indices that are known non-poison.
For certain patterns that restrict the range of an index, we can insert
a freeze of the original value, to prevent propagation of poison.

[spatel]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

779

Other than the name, this looks like it could be a generic freeze utility. Should we put it in some header for other potential callers? cc @aqjune @hyeongyukim

812

Use InsertPointGuard so caller doesn't have to manually reset it.

861

stray punctuation

[aqjune]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

779

Not sure whether this class can be reused, but I agree that having a header containing utilities for dealing with freeze is a good idea (or probably reuse Local.h?).
It's because we might want to keep freezes in a specific pattern for easier optimization. Currently, it is InstCombine's work to keep freeze-including expressions in a specific pattern, but it is kind of hidden.
In the case of LoopUnswitch, naively freezing the branch condition caused many assembly diffs.
To resolve this regression, we are experimenting with (1) pushing freeze into icmp arguments (which could resolve assembly diffs a lot: https://reviews.llvm.org/D106041#2922020) (2) replacing all uses of a branch condition with the frozen one.
Once the scheme is set, I'd like to have FreezeBrCond(BranchInst) that performs what LoopUnswitch is supposed to do.

[fhahn]

Addressed comments, thanks: 1) Use insert point guard and 2) remove stray ;

[fhahn]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

779

Other than the name, this looks like it could be a generic freeze utility. Should we put it in some header for other potential callers? cc @aqjune @hyeongyukim

I tried to keep things as general as possible here. But for now it serves a very specific use case, so maybe it would make sense to wait with moving this to a header until there are similar use cases?

[lebedev.ri]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

808–810

I think you want to assert that User is an user of ToFreeze.

810

810

Why do you have to go through all the uses of ToFreeze, even though you only want to fix User?
Can't you just go through the operands of User?

[nikic]

Is these some specific use case for handling this special case?

[fhahn]

Added assertion as suggested and also just update operands of UserI, thanks!

In D107580#2944538, @nikic wrote:

Is these some specific use case for handling this special case?

LLVM doesn't do a great job when lowering operations on largish vectors in general, e.g. originating from uses of extended vectors in C/C++ or other frontends. The patch in particular improves the lowering of code like below

using vec_ty = float __attribute__((ext_vector_type(32)));

void foo(vec_ty *ptr, int idx) {
  (*ptr)[idx&31] = 0.0;
}

https://godbolt.org/z/cq97c4KG5

[fhahn]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

810

Can't you just go through the operands of User?

For the current version of the API yes! Updated, thanks!

[lebedev.ri]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

790

Assertion message

811

is_contained(ToFreeze->users(), &UserI)?

[fhahn]

Use is_contained and add assertion message, thanks!

[lebedev.ri]

Seems fine to me, thanks.

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

785–786

maybe

816

843–844

@nlopes I wanted to check what can we do for potentially OOB insertions,
and i guess this is an alive2 problem:
https://alive2.llvm.org/ce/z/VhUXF-
https://alive2.llvm.org/ce/z/qHk9Wv
?

989–991

This should either contain a FIXME comment, or a comment saying that here freeze won't help.

[nlopes]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

843–844

Both tests look correct to me.
The 2nd one is funny because the 'and' doesn't do anything. A proper masking should be 'and 7'.

[fhahn]

Address outstanding comments. I'll land this version in a bit.

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

785–786

I think for constructors it is quite common to have names matching the directly initialized fields for arguments. To make things consistent for the constructor, I renamed to S argument to Status.

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

static bool isMemModifiedBetween(BasicBlock::iterator Begin,
                                  const MemoryLocation &Loc, AAResults &AA) {
   unsigned NumScanned = 0;
   return std::any_of(Begin, End, [&](const Instruction &Instr) {
     return isModSet(AA.getModRefInfo(&Instr, Loc)) ||
            ++NumScanned > MaxInstrsToScan;
   });
 }
 
+/// Helper class to indicate whether a vector index can be safely scalarized and
+/// if a freeze needs to be inserted.
+class ScalarizationResult {
+  enum class StatusTy { Unsafe, Safe, SafeWithFreeze };
+
+  StatusTy Status;
+  Value *ToFreeze;
+
+  ScalarizationResult(StatusTy S, Value *ToFreeze = nullptr)
+      : Status(S), ToFreeze(ToFreeze) {}
+
+public:
+  ScalarizationResult(const ScalarizationResult &Other) = default;
+  ~ScalarizationResult() { assert(!ToFreeze); }
+
+  static ScalarizationResult unsafe() { return {StatusTy::Unsafe}; }
+  static ScalarizationResult safe() { return {StatusTy::Safe}; }
+  static ScalarizationResult safeWithFreeze(Value *ToFreeze) {
+    return {StatusTy::SafeWithFreeze, ToFreeze};
+  }
+
+  /// Returns true if the index can be scalarize without requiring a freeze.
+  bool isSafe() const { return Status == StatusTy::Safe; }
+  /// Returns true if the index cannot be scalarized.
+  bool isUnsafe() const { return Status == StatusTy::Unsafe; }
+  /// Returns true if the index can be scalarize, but requires inserting a
+  /// freeze.
+  bool isSafeWithFreeze() const { return Status == StatusTy::SafeWithFreeze; }
+
+  /// Freeze the ToFreeze and update the use in \p User to use it.
+  void freeze(IRBuilder<> &Builder, Instruction &UserI) {
+    assert(isSafeWithFreeze() &&
+           "should only be used when freezing is required");
+    assert(
+        any_of(ToFreeze->users(), [&UserI](User *U) { return &UserI == U; }) &&
+        "UserI must be a user of ToFreeze");
+    IRBuilder<>::InsertPointGuard Guard(Builder);
+    Builder.SetInsertPoint(cast<Instruction>(&UserI));
+    Value *Frozen = Builder.CreateFreeze(ToFreeze);
+    for (Use &U : make_early_inc_range((UserI.operands()))) {
+      if (U.get() == ToFreeze)
+        U.set(Frozen);
+    }
+
+    ToFreeze = nullptr;
+  }
+};
+
 /// Check if it is legal to scalarize a memory access to \p VecTy at index \p
 /// Idx. \p Idx must access a valid vector element.
-static bool canScalarizeAccess(FixedVectorType *VecTy, Value *Idx,
-                               Instruction *CtxI, AssumptionCache &AC) {
-  if (auto *C = dyn_cast<ConstantInt>(Idx))
-    return C->getValue().ult(VecTy->getNumElements());
-
-  if (!isGuaranteedNotToBePoison(Idx, &AC))
-    return false;
+static ScalarizationResult canScalarizeAccess(FixedVectorType *VecTy,
+                                              Value *Idx, Instruction *CtxI,
+                                              AssumptionCache &AC) {
+  if (auto *C = dyn_cast<ConstantInt>(Idx)) {
+    if (C->getValue().ult(VecTy->getNumElements()))
+      return ScalarizationResult::safe();
+    return ScalarizationResult::unsafe();
+  }
 
-  APInt Zero(Idx->getType()->getScalarSizeInBits(), 0);
-  APInt MaxElts(Idx->getType()->getScalarSizeInBits(), VecTy->getNumElements());
+  unsigned IntWidth = Idx->getType()->getScalarSizeInBits();
+  APInt Zero(IntWidth, 0);
+  APInt MaxElts(IntWidth, VecTy->getNumElements());
   ConstantRange ValidIndices(Zero, MaxElts);
-  ConstantRange IdxRange = computeConstantRange(Idx, true, &AC, CtxI, 0);
-  return ValidIndices.contains(IdxRange);
+  ConstantRange IdxRange(IntWidth, true);
+
+  if (isGuaranteedNotToBePoison(Idx, &AC)) {
+    if (ValidIndices.contains(computeConstantRange(Idx, true, &AC, CtxI, 0)))
+      return ScalarizationResult::safe();
+    return ScalarizationResult::unsafe();
+  }
+
+  // If the index may be poison, check if we can insert a freeze before the
+  // range of the index is restricted.
+  Value *IdxBase;
+  ConstantInt *CI;
+  if (match(Idx, m_And(m_Value(IdxBase), m_ConstantInt(CI)))) {
+    IdxRange = IdxRange.binaryAnd(CI->getValue());
+  } else if (match(Idx, m_URem(m_Value(IdxBase), m_ConstantInt(CI)))) {
+    IdxRange = IdxRange.urem(CI->getValue());
+  }
+
+  if (ValidIndices.contains(IdxRange))
+    return ScalarizationResult::safeWithFreeze(IdxBase);
+  return ScalarizationResult::unsafe();
 }
 
 /// The memory operation on a vector of \p ScalarType had alignment of
 /// \p VectorAlignment. Compute the maximal, but conservatively correct,
 /// alignment that will be valid for the memory operation on a single scalar
 /// element of the same type with index \p Idx.
 static Align computeAlignmentAfterScalarization(Align VectorAlignment,
                                                 Type *ScalarType, Value *Idx,
                                                 const DataLayout &DL) {
bool VectorCombine::foldSingleElementStore(Instruction &I) {
   if (auto *Load = dyn_cast<LoadInst>(Source)) {
     auto VecTy = cast<FixedVectorType>(SI->getValueOperand()->getType());
     const DataLayout &DL = I.getModule()->getDataLayout();
     Value *SrcAddr = Load->getPointerOperand()->stripPointerCasts();
     // Don't optimize for atomic/volatile load or store. Ensure memory is not
     // modified between, vector type matches store size, and index is inbounds.
     if (!Load->isSimple() || Load->getParent() != SI->getParent() ||
         !DL.typeSizeEqualsStoreSize(Load->getType()) ||
-        !canScalarizeAccess(VecTy, Idx, Load, AC) ||
-        SrcAddr != SI->getPointerOperand()->stripPointerCasts() ||
+        SrcAddr != SI->getPointerOperand()->stripPointerCasts())
+      return false;
+
+    auto ScalarizableIdx = canScalarizeAccess(VecTy, Idx, Load, AC);
+    if (ScalarizableIdx.isUnsafe() ||
         isMemModifiedBetween(Load->getIterator(), SI->getIterator(),
                              MemoryLocation::get(SI), AA))
       return false;
 
+    if (ScalarizableIdx.isSafeWithFreeze())
+      ScalarizableIdx.freeze(Builder, *cast<Instruction>(Idx));
     Value *GEP = Builder.CreateInBoundsGEP(
         SI->getValueOperand()->getType(), SI->getPointerOperand(),
         {ConstantInt::get(Idx->getType(), 0), Idx});
     StoreInst *NSI = Builder.CreateStore(NewElement, GEP);
     NSI->copyMetadata(*SI);
     Align ScalarOpAlignment = computeAlignmentAfterScalarization(
         std::max(SI->getAlign(), Load->getAlign()), NewElement->getType(), Idx,
         DL);
if (LastCheckedInst->comesBefore(UI)) {
       }
     }
 
     if (!LastCheckedInst)
       LastCheckedInst = UI;
     else if (LastCheckedInst->comesBefore(UI))
       LastCheckedInst = UI;
 
-    if (!canScalarizeAccess(FixedVT, UI->getOperand(1), &I, AC))
+    auto ScalarIdx = canScalarizeAccess(FixedVT, UI->getOperand(1), &I, AC);
+    if (!ScalarIdx.isSafe())
       return false;
 
     auto *Index = dyn_cast<ConstantInt>(UI->getOperand(1));
     OriginalCost +=
         TTI.getVectorInstrCost(Instruction::ExtractElement, LI->getType(),
                                Index ? Index->getZExtValue() : -1);
     ScalarizedCost +=
         TTI.getMemoryOpCost(Instruction::Load, FixedVT->getElementType(),
                             Align(1), LI->getPointerAddressSpace());

llvm/test/Transforms/VectorCombine/load-insert-store.ll

entry:
   %vecins = insertelement <16 x i8> %0, i8 %s, i32 %idx.clamped.frozen
   store <16 x i8> %vecins, <16 x i8>* %q
   ret void
 }
 
 define void @insert_store_nonconst_index_known_valid_by_and_but_may_be_poison(<16 x i8>* %q, i8 zeroext %s, i32 %idx) {
 ; CHECK-LABEL: @insert_store_nonconst_index_known_valid_by_and_but_may_be_poison(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = load <16 x i8>, <16 x i8>* [[Q:%.*]], align 16
-; CHECK-NEXT:    [[IDX_CLAMPED:%.*]] = and i32 [[IDX:%.*]], 7
-; CHECK-NEXT:    [[VECINS:%.*]] = insertelement <16 x i8> [[TMP0]], i8 [[S:%.*]], i32 [[IDX_CLAMPED]]
-; CHECK-NEXT:    store <16 x i8> [[VECINS]], <16 x i8>* [[Q]], align 16
+; CHECK-NEXT:    [[TMP0:%.*]] = freeze i32 [[IDX:%.*]]
+; CHECK-NEXT:    [[IDX_CLAMPED:%.*]] = and i32 [[TMP0]], 7
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds <16 x i8>, <16 x i8>* [[Q:%.*]], i32 0, i32 [[IDX_CLAMPED]]
+; CHECK-NEXT:    store i8 [[S:%.*]], i8* [[TMP1]], align 1
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %0 = load <16 x i8>, <16 x i8>* %q
   %idx.clamped = and i32 %idx, 7
   %vecins = insertelement <16 x i8> %0, i8 %s, i32 %idx.clamped
   store <16 x i8> %vecins, <16 x i8>* %q
   ret void
entry:
   %vecins = insertelement <16 x i8> %0, i8 %s, i32 %idx.clamped.frozen
   store <16 x i8> %vecins, <16 x i8>* %q
   ret void
 }
 
 define void @insert_store_nonconst_index_known_valid_by_urem_but_may_be_poison(<16 x i8>* %q, i8 zeroext %s, i32 %idx) {
 ; CHECK-LABEL: @insert_store_nonconst_index_known_valid_by_urem_but_may_be_poison(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = load <16 x i8>, <16 x i8>* [[Q:%.*]], align 16
-; CHECK-NEXT:    [[IDX_CLAMPED:%.*]] = urem i32 [[IDX:%.*]], 16
-; CHECK-NEXT:    [[VECINS:%.*]] = insertelement <16 x i8> [[TMP0]], i8 [[S:%.*]], i32 [[IDX_CLAMPED]]
-; CHECK-NEXT:    store <16 x i8> [[VECINS]], <16 x i8>* [[Q]], align 16
+; CHECK-NEXT:    [[TMP0:%.*]] = freeze i32 [[IDX:%.*]]
+; CHECK-NEXT:    [[IDX_CLAMPED:%.*]] = urem i32 [[TMP0]], 16
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds <16 x i8>, <16 x i8>* [[Q:%.*]], i32 0, i32 [[IDX_CLAMPED]]
+; CHECK-NEXT:    store i8 [[S:%.*]], i8* [[TMP1]], align 1
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %0 = load <16 x i8>, <16 x i8>* %q
   %idx.clamped = urem i32 %idx, 16
   %vecins = insertelement <16 x i8> %0, i8 %s, i32 %idx.clamped
   store <16 x i8> %vecins, <16 x i8>* %q
   ret void