[LoopVectorize] Support reductions that store intermediary result

[igor.kirillov]

Summary

Adds ability to vectorize loops containing a store to a loop-invariant address as part of a reduction that isn't converted to SSA form due to lack of aliasing info. Runtime checks are generated to ensure the store does not alias any other accesses in the loop.
Ordered fadd reductions are not yet supported.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

941

What about loads to the same address in the loop? At the moment, LAA cannot analyze dependences with invariant addresses. But if this limitation gets removed the code here may become incorrect, because it relies on this limitation IIUC?

[david-arm]

llvm/include/llvm/Analysis/LoopAccessAnalysis.h

643

I assume here that InvariantStores refers to stores to an invariant (i.e. uniform?) address, rather than storing an invariant value to variant address. If so, perhaps it could be named StoresToInvariantAddress or something like that?

[david-arm]

llvm/test/Transforms/LoopVectorize/reduction.ll

471 ↗

(On Diff #378503)

Is it possible to add a simple floating point test with "fadd fast"?

[igor.kirillov]

It seems like the main problem is that we potentially bail out too early at the moment when checking for reductions due to the store, but once we generate runtime checks, sinking the store may become legal (see inline comment about loads to the same address)? If that's the case, ideally we'd just sink any such loads/stores before detecting reductions once we know they can be sunk due to runtime checks, but unfortunately I do not think that's possible with the current structure/ordering.

Is it worth for me to explore if it possible to do other way around or we should work on with the solution from this merge?

llvm/include/llvm/Analysis/LoopAccessAnalysis.h

643

@fhah Just to make sure we are on the same line - for me uniform and invariant are just synonyms, isn't it so?

llvm/include/llvm/Transforms/Utils/LoopUtils.h

402 ↗

(On Diff #378503)

What do you think if I make this function a public member of class ScalarEvolution? Or is there a better place for it?

llvm/lib/Analysis/IVDescriptors.cpp

322

I checked and only LoopInterchangeLegality is using this function and it is not affected by stores. Anyway, I can add a parameter to RecurrenceDescriptor::isReductionPHI or a member to RecurrenceDescriptor allowing or not to handle stores. What do you think about it?

The comment is to be updated, yes.

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

941

Loads from or stores to the same address in the loop? I'm sorry could you clarify what the problem is. As it is I don't understand the message.

[david-arm]

Hi @igor.kirillov, is it also possible to get this working for ordered reductions, i.e.

float sum = 0;
for(i=0..N) {
  sum += src[i];
  dst[42] = sum;
}

when building with -O3? I think it might mean updating checkOrderedReductions to look through the store. If it looks too difficult to do as part of this patch we can always follow-up with a patch later.

[igor.kirillov]

Update commit message and comments
Move storeToSameAddress to ScalarEvolution
Add fadd fast test
Do not apply patch for ordered fadd reductions

[igor.kirillov]

In D110235#3072725, @david-arm wrote:

Hi @igor.kirillov, is it also possible to get this working for ordered reductions, i.e.

float sum = 0;
for(i=0..N) {
  sum += src[i];
  dst[42] = sum;
}

when building with -O3? I think it might mean updating checkOrderedReductions to look through the store. If it looks too difficult to do as part of this patch we can always follow-up with a patch later.

Yes, I added a check to LoopVectorizationLegality::canVectorizeFPMath so as not to allow this reduction when math is strict. Enabling it requires some work and it is better to do it separately.

llvm/test/Transforms/LoopVectorize/reduction.ll

471 ↗

(On Diff #378503)

Added! see reduc_store_fadd_fast function

[david-arm]

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

1011

Maybe this can be folded into the all_of case below, i.e.

return (all_of(getReductionVars(), [&](auto &Reduction) -> bool {
  const RecurrenceDescriptor &RdxDesc = Reduction.second;
  return !RdxDesc.hasExactFPMath() ||
         (RdxDesc.isOrdered() && !RdxDesc.IntermediateStore);
}));

Also, the problem with this code at the moment is that you could have a mixture of fast and ordered reductions in the same loop. There could be an intermediate store for one of the fast reductions, but not for the ordered ones. At the moment with your code we will just bail out in this case.

[igor.kirillov]

Update intermediate store check for ordered fadd vectorization

[igor.kirillov]

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

1011

You are right! It can be done much simpler and the check is needed only when ordered reduction is present.

[david-arm]

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

943

Hi @igor.kirillov, I think you're missing a test case here. Suppose we have two stores in the loop to the same invariant address. If the first store is predicated, but the second isn't we should still vectorise because the second store wins and the first one can be removed. At least the code here suggests that. I couldn't find any test that exercised this code path.

951

I think you might be missing a test case for this. Can you make sure this code path is exercised please by your existing tests please?

llvm/test/Transforms/LoopVectorize/reduction.ll

477 ↗

(On Diff #385810)

Can you put all the CHECK lines in the same place near the top of the function please to be consistent with the existing tests?

517 ↗

(On Diff #385810)

Hi @igor.kirillov, this doesn't look right. The test is called @reduc_store_fadd_fast, but there is not fast keyword on the fadd instruction. I don't think we should even be vectorising this because it requires ordered reductions, which you haven't enabled for this test. Can you change the name of this to @reduc_store_fadd_ordered and investigate why we are vectorising this?

Can you also add a separate test called @reduc_store_fadd_fast that actually has the fast keyword too?

[igor.kirillov]

Update tests

[igor.kirillov]

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

943

Added reduc_store_final_store_predicated test that address both requests

951

It is now covered by new reduc_store_final_store_predicated test and an old test reduc_store_inside_unrolled also executes this path

llvm/test/Transforms/LoopVectorize/reduction.ll

517 ↗

(On Diff #385810)

I missed the fast keyword there. As for why this code gets vectorized - it happens because of Hints->allowReordering() returning true in LoopVectorizationLegality::canVectorizeFPMath. As you can see the test specifies vector width (-force-vector-width=4) and llvm allows to process ordered instructions in unordered manner (see also LoopVectorizeHints::allowReordering function) in that case.

I added a new test with fast keyword to AArch64/strict-fadd.ll and the loop is not vectorized there.

[david-arm]

llvm/test/Transforms/LoopVectorize/reduction.ll

656 ↗

(On Diff #388126)

Can you also add a test where the first store is predicated and the second one isn't? According to the code changes in this patch we should vectorise this case because the second one overrides the first.

[igor.kirillov]

Add test

llvm/test/Transforms/LoopVectorize/reduction.ll

656 ↗

(On Diff #388126)

Added! See reduc_store_middle_store_predicated

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

941

The case I was thinking about was something like the snippet below, where we have a load of the invariant address in the loop (%lv = load... in the example below).

define void @reduc_store(i32* %dst, i32* readonly %src, i32* noalias %dst.2) {
entry:
  %arrayidx = getelementptr inbounds i32, i32* %dst, i64 42
  store i32 0, i32* %arrayidx, align 4
  br label %for.body
for.body:
  %0 = phi i32 [ 0, %entry ], [ %add, %for.body ]
  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
  %arrayidx1 = getelementptr inbounds i32, i32* %src, i64 %indvars.iv
  %1 = load i32, i32* %arrayidx1, align 4
  %add = add nsw i32 %0, %1
  %lv = load i32, i32* %arrayidx
  store i32 %add, i32* %arrayidx, align 4
  
  %gep.dst.2  = getelementptr inbounds i32, i32* %dst.2, i64 %indvars.iv
  store i32 %lv, i32* %gep.dst.2,

  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
  %exitcond = icmp eq i64 %indvars.iv.next, 1000
  br i1 %exitcond, label %for.cond.cleanup, label %for.body

for.cond.cleanup:
  ret void
}

llvm/test/Transforms/LoopVectorize/reduction.ll

468 ↗

(On Diff #388148)

Could you add a brief textual explanation of what the test covers?

482 ↗

(On Diff #388148)

perhaps rename to %sum or something like that, to make it a bit easier to read the test?

487 ↗

(On Diff #388148)

The names for the 2 different GEPs in the tests are very similar. Could you rename them to make it easier to distinguish them (e.g. something like %gep.src/%gep.dst)..

[david-arm]

Hi @igor.kirillov, thanks for all the changes and the patch looks good now! I just had a couple of minor comments.

llvm/lib/Analysis/IVDescriptors.cpp

501

nit: Could you remove an extra level of indentation here before merging the patch?

i.e. this can be written like this:

if (isa<PHINode>(UI))
  PHIs.push_back(UI);
else if (auto *SI = dyn_cast<StoreInst>(UI)) {
  ...
} else
  NonPHIs.push_back(UI);

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

1021

nit: I think you can you might be able to simplify the code here by removing FoundMatchingRecurrence . Then, lower down instead of the break you can just do

if (DSI && (DSI == SI)) {
  *IsPredicated = blockNeedsPredication(DSI->getParent());
  return true;
}

and at the bottom of the function just do:

return false;

[igor.kirillov]

Refactoring a bit

[igor.kirillov]

llvm/lib/Analysis/IVDescriptors.cpp

501

I made an extra step and simplified it even more

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

941

In that case we do not vectorize, the rejection happens in LoopAccessInfo::analyzeLoop when loads ape processed:

  for (LoadInst *LD : Loads) {
    Value *Ptr = LD->getPointerOperand();
...
    // See if there is an unsafe dependency between a load to a uniform address and
    // store to the same uniform address.
    if (UniformStores.count(Ptr)) {
      LLVM_DEBUG(dbgs() << "LAA: Found an unsafe dependency between a uniform "
                           "load and uniform store to the same address!\n");
      HasDependenceInvolvingLoopInvariantAddress = true;
    }
...

I added reduc_store_load test anyway

[igor.kirillov]

Add scalable reduction test

[igor.kirillov]

Little extension in a new test

[david-arm]

LGTM! Thanks for making all the changes @igor.kirillov.

llvm/test/Transforms/LoopVectorize/reduction.ll

466 ↗

(On Diff #388188)

nit: I think this should be invariant

[igor.kirillov]

Fix typo

[fhahn]

llvm/include/llvm/Analysis/IVDescriptors.h

268

nit: should be a doc-comment?

llvm/include/llvm/Analysis/LoopAccessAnalysis.h

579

nit: could return ArrayRef, not leaking any details about the underlying container.

643

nit: Is there a reason for choosing 5 for the size? if not, nowadays SmallVector can pick a good size automatically.

llvm/include/llvm/Analysis/ScalarEvolution.h

1120 ↗

(On Diff #388906)

Do you anticipate this to be used outside Transforms/Vectorize? If not I'm not sure if it makes sense to live here, and extending the API interface of ScalarEvolution. Can this instead be defined somewhere in Transforms/Vectorize?

llvm/lib/Analysis/IVDescriptors.cpp

250

needs documentation?

322

I guess it would be good to have at least a test case for loop-interchange to make sure it can deal with the change properly?

530

nit: for consistency with other code avoid using != nullptr. This is more in line with the style use in LLVM in general (and you use ==/!= nullptr in the adjacent code here).

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

941

Yeah but this is only due to some limitations that currently existing in LAA, right? I think we should at least make this clear somewhere, e.g. in a comment.

947

nit: no need to use llvm::

1024

nit: redundant ().

1025

could we instead just do the check at the call site or is there a benefit of doing it here?

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

2788

We effectively sink the store outside the loop. In that case, I don't think we should create a recipe and also we should not consider its cost.

[igor.kirillov]

Lots of updates related to the recent comments

[igor.kirillov]

llvm/lib/Analysis/IVDescriptors.cpp

322

You are right. We actually should not do loop interchange is an invariant address is present. Otherwise it may introduce incorrect result.

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

941

Added comment

[igor.kirillov]

Hi, @fhahn! Since there have been several quite serious changes in the patch after your last review, I would be happy to receive you approval before merge (even though status is accepted now).

[fhahn]

In D110235#3174384, @igor.kirillov wrote:

Hi, @fhahn! Since there have been several quite serious changes in the patch after your last review, I would be happy to receive you approval before merge (even though status is accepted now).

sure, I'll try to take another look by end-of-day Monday

[igor.kirillov]

@fhahn ping

[fhahn]

Thanks for the latest update! I left some more additional comments. It might be good to pre-commit the tests and only include the diff in the patch. Another thing to consider is moving the reduction-store tests to a separate file, as reduction.ll is already quite large.

One thing to note on the overall approach is that it's a bit of a workaround some current limitations in our modeling, but I don't think there's an alternative in the short term and this is clearly a desirable case to support, so that seems fine to me in general. There's work in progress to model the pre-header and exit block in VPlan, which allows us to do the sinking of the store more easily, so we can improve things further then :)

llvm/include/llvm/Analysis/IVDescriptors.h

175–176

From an interface perspective, I think it would be better to make SE an optional argument and only perform the the store analysis when it is passed. This would also require users to explicitly opt-in, which would at least partially guarding against people ignoring IntermediateStore.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8978

either deleted or go outside the loop sounds a bit unclear. Aren't they moved to the vector exit block and store the final reduction value?

8980

I think it would also be good to include at least some information in the VPlan that the store is handled as part of the reduction. Perhaps VPReductionRecipe should print if the result is stored after the loop? Please add a test case to vplan-printing.ll

llvm/test/Transforms/LoopVectorize/reduction.ll

469 ↗

(On Diff #390414)

FWIW for such compact IR test cases, the pseudo code doesn't add much value in my personal opinion. Better to strive to make the tests as readable/compact as possible and have a comment explaining what it tests when needed.

546 ↗

(On Diff #390414)

nit: newline.

628 ↗

(On Diff #390414)

Here (and for the other tests it would be good to check at least the vector reduction sequence and that the correct value is stored.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

936

It would probably be good to make clear what is handled here exactly and why we can handle those stores. IIUC this applies only to invariant stores that store reduction results and is safe because runtime checks guarantee that it won't alias with other objects. The store won't get vectorized, but sunk to the exit block during codegen.

llvm/test/Transforms/LoopVectorize/reduction.ll

582 ↗

(On Diff #390414)

not needed?

585 ↗

(On Diff #390414)

nit: is this needed? Can just pass %n as i64

672 ↗

(On Diff #390414)

move exit block to end of function?

711 ↗

(On Diff #390414)

move exit block to end of function?

753 ↗

(On Diff #390414)

move exit block to end of function?

803 ↗

(On Diff #390414)

move exit block to end of function?

[lebedev.ri]

Not sure if you saw, there's a somewhat related bugreport: https://github.com/llvm/llvm-project/issues/50286
Not sure if this already supports that pattern though.

[igor.kirillov]

Update tests, move invariant store tests into separate file, add more comments

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8980

I'm not sure how to display this information properly. This is how output of VReductionRecipe::print looks like now:

REDUCE ir<%red.next> = ir<%red> + fast reduce.fadd (ir<%lv>)

I could add something to the end but it doesn't seem to fit there. Also I have not found any proper V.*Recipe where I could place something like DELETE store .*.

llvm/test/Transforms/LoopVectorize/reduction.ll

469 ↗

(On Diff #390414)

When I look at pseudo-code I immediately understand what it is about, whereas looking at IR takes at least 30 seconds of cognitive exertions. And it is also easier to see the difference between and purpose of all those quite similar tests. But I delete it, of course, if you insist :)

[igor.kirillov]

@fhahn I created the review with tests only here - https://reviews.llvm.org/D117213. Once it is merged I'll update this one.

[igor.kirillov]

@lebedev.ri This patch addresses a different problem. Here we try to handle when a reduction value is stored in some address and if this address is invariant plus other lucky conditions are satisfied we manage to vectorize. In your example address where value is stored is not invariant at all. Nevertheless, the case is interesting.

[igor.kirillov]

@fhahn ping

[igor.kirillov]

Update tests

[igor.kirillov]

Add invariant store information to VReductionRecipe::print output

[igor.kirillov]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8980

@fhahn I added some info to VReductionRecipe::print and a relevant test. What do you think about it?

[david-arm]

Hi @igor.kirillov, the patch looks good now! I just had one question about an odd-looking CHECK line in one of the test files.

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

455

nit: I think it's fine (and cleaner) to just do

%[[PHI:%.*]] =
%[[ADDR:%.*]] =

here. Also, can you add the incoming value to the phi just to make sure it's the correct value? I think it should just be

phi i32 {{ [[TMP]], %middle.block }}

483

This looks a bit odd. I'm not sure why it's needed? If there is a CHECK-LABEL for every function I don't think this should happen.

[fhahn]

Thanks for the update!

llvm/lib/Analysis/IVDescriptors.cpp

365

Do we have to make sure that the stored value is feeding the phi again and not an earlier value in the cycle?

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

179

I think it would be good to check the full reduction cycle & store here. I don't think this test case is handled correctly at the moment as the IR seems out of sync with the pseudo code (which is why personally I think the C pseudo code is a bit distracting)

Note that in the IR the reduction cycle is %sum -> sum.1 -> %sum and not %sum -> %sum.1 -> %sum.2 -> %sum. When this gets vectorized, the final value written to dst[42] is the final value of %sum.1, not %sum.2 as it should be I think.

[igor.kirillov]

Add fix that prevents vectorization for the case when not a final reduction value stored in a loop invariant address.
Add test for this case.
Remove some unused code and rename a couple of IR variables in the reduction-with-invariant-store.ll test.
Fix incorrect phi node in reduc_store_inside_unrolled and reduc_double_invariant_store tests.

[igor.kirillov]

llvm/lib/Analysis/IVDescriptors.cpp

365

Yes! I added this check and also a test case (reduc_store_not_final_value)

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

179

The IR was incorrect actually because of the error I made during renaming and the pseudo-code was showing my original plot for the test-case. Luckily it helped to expose a bug, so I added an extra check to IVDescriptors.cpp and now reduc_store_not_final_value function is testing this case (when we have IntermediateStore but no ExitInstruction and value stored is not actually a final value).

[fhahn]

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

179

I think it would still be good to test that the full reduction cycle is generated correctly (probably worth checking the full vector body + middle.block), at least for some of the tests in the file, to make sure the full sequence is generated correctly.

[fhahn]

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

179

The IR was incorrect actually because of the error I made during renaming and the pseudo-code was showing my original plot for the test-case.

This seems to indicate that the pseudo code may be distracting (:

[igor.kirillov]

Add full vector.body and middle.block checks for reduc_store and reduc_store_inside_unrolled tests

[igor.kirillov]

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

179

@fhahn Added full checks for reduc_store and this function. The code generated is overscalarized for reduc_store_inside_unrolled but I checked and without invariant stores it is the same.

[fhahn]

Thanks for the update with the tests! I think there might still be a correctness issue. More details inline.

llvm/include/llvm/Analysis/IVDescriptors.h

177

nit: addresses ?

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

965

I'm not sure if the comment matches the code. There's no guarantee that all stores to the same address come before the store of the reduction AFAICT. E.g. the test below has a store to the same address after the store of the reduction. I think this may get mis-compiled, as the store to 0 gets replaced with the store of the final value of the reduction.

define void @reduc_store(i32* %dst, i32* readonly %src) {
entry:
  %gep.dst = getelementptr inbounds i32, i32* %dst, i64 42
  store i32 1, i32* %gep.dst, align 4
  br label %for.body

for.body:
  %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
  %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
  %gep.src = getelementptr inbounds i32, i32* %src, i64 %iv
  %0 = load i32, i32* %gep.src, align 4
  %add = add nsw i32 %sum, %0
  store i32 %add, i32* %gep.dst, align 4
  store i32 0, i32* %gep.dst, align 4
  %iv.next = add nuw nsw i64 %iv, 1
  %exitcond = icmp eq i64 %iv.next, 1000
  br i1 %exitcond, label %exit, label %for.body

exit:
  ret void
}

1023

nit: can drop DSI, as SI is guaranteed to be non-null?

[igor.kirillov]

Fix incorrectly processed case when reduction value stored in invariant could be overwritten inside loop

[igor.kirillov]

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

965

Yeah, it was processed incorrectly. I added this test along with the fix.

[igor.kirillov]

@fhahn, gentle ping :)

[fhahn]

Thanks for the update! I'll do a bit more testing and will take another look soon, but I *think* it should be good now.

[igor.kirillov]

@fhahn, @david-arm ping. Also I run lnt test-suite and everything is fine and 17 more loops were vectorized according to -stats

[david-arm]

LGTM! Thanks for addressing all the comments and fixing bugs, adding new tests, etc. I have a couple of nits you can address before merging. I think given that you've run the LLVM test suite without seeing failures and we know more loops are vectorising then that gives us a good level of confidence. :)

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

256

nit: Should be i++, not i+=2.

303

nit: Looks like an unnecessary change from %sum.1 -> %sum.2.

[fhahn]

LGTM, thanks! I added a few additional comments that can be addressed directly before committing.

FYI work has started to model the exit block in VPlan as well: D123457, D123537. fixReduction and the store sinking can and should now be migrated to be modeled explicitly in VPlan.

llvm/include/llvm/Analysis/IVDescriptors.h

273

nit: would be good to also document what this means, what the properties of such stores are.

llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h

311

nit: could be more specific, the increment of the reduction is used as stored operand, right?

Also, they only handle reductions right? Then using Reduction instead of Recurrence seems more descriptive, .e.g. isInvariantStoreOfReduction? (same for isRecurringInvariantAddress

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

453

With opaque pointers, this code may behave differently.

It's possible to have

store i32 0, ptr %x
store i8 0, ptr %x

In that case the pointer operands will be the same, but different store widths. I *think* we should also check that the types of the stored values match for now, as we use this to remove earlier stores. This should only be correct if the later store writes at least as many bits as the earlier stores.

1020

nit: could be any_of

1030

nit: could be any_of.

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

412

Would it be possible to add runtime tests for the mis-compiles fixed to https://github.com/llvm/llvm-test-suite/tree/main/SingleSource/UnitTests/Vectorizer?

[igor.kirillov]

Lots of updates related to the recent comments

[igor.kirillov]

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

453

I think if this check is added here, then the purpose of the function would be different (address is still the same even if values have different size). So, instead of that I added a check to a place where those pointers are really processed - see LoopVectorizationLegality.cpp:975. There we now make sure that values stored in an invariant address are of the same type.

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

303

sum.1 is sum + src[i] and sum.2 is sum + src[i] + src[i+1]. Looks like everything is fine for me.

[igor.kirillov]

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

412

Added a simple test here - https://reviews.llvm.org/D124609
I don't feel like any other test adds more robustness, but if you have an idea what could be also added I'll gladly implement it!

[uabelho]

Hello,

I wrote
https://github.com/llvm/llvm-project/issues/57572
about a verifier complaint/crash that started happening with this patch.

llvm/include/llvm/Analysis/IVDescriptors.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_IVDESCRIPTORS_H
 #define LLVM_ANALYSIS_IVDESCRIPTORS_H
 
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/ValueHandle.h"
 
 namespace llvm {
 
 class AssumptionCache;
 class DemandedBits;
 class DominatorTree;
 class Instruction;
 class Loop;
 class PredicatedScalarEvolution;
 class ScalarEvolution;
 class SCEV;
+class StoreInst;
 
 /// These are the kinds of recurrences that we support.
 enum class RecurKind {
   None,       ///< Not a recurrence.
   Add,        ///< Sum of integers.
   Mul,        ///< Product of integers.
   Or,         ///< Bitwise or logical OR of integers.
   And,        ///< Bitwise or logical AND of integers.

 /// special case of chains of recurrences (CR). See ScalarEvolution for CR
 /// references.
 
 /// This struct holds information about recurrence variables.
 class RecurrenceDescriptor {
 public:
   RecurrenceDescriptor() = default;
 
-  RecurrenceDescriptor(Value *Start, Instruction *Exit, RecurKind K,
-                       FastMathFlags FMF, Instruction *ExactFP, Type *RT,
-                       bool Signed, bool Ordered,
+  RecurrenceDescriptor(Value *Start, Instruction *Exit, StoreInst *Store,
+                       RecurKind K, FastMathFlags FMF, Instruction *ExactFP,
+                       Type *RT, bool Signed, bool Ordered,
                        SmallPtrSetImpl<Instruction *> &CI,
                        unsigned MinWidthCastToRecurTy)
-      : StartValue(Start), LoopExitInstr(Exit), Kind(K), FMF(FMF),
-        ExactFPMathInst(ExactFP), RecurrenceType(RT), IsSigned(Signed),
-        IsOrdered(Ordered),
+      : IntermediateStore(Store), StartValue(Start), LoopExitInstr(Exit),
+        Kind(K), FMF(FMF), ExactFPMathInst(ExactFP), RecurrenceType(RT),
+        IsSigned(Signed), IsOrdered(Ordered),
         MinWidthCastToRecurrenceType(MinWidthCastToRecurTy) {
     CastInsts.insert(CI.begin(), CI.end());
   }
 
   /// This POD struct holds information about a potential recurrence operation.
   class InstDesc {
   public:
     InstDesc(bool IsRecur, Instruction *I, Instruction *ExactFP = nullptr)
public:
 
   /// Returns the opcode corresponding to the RecurrenceKind.
   static unsigned getOpcode(RecurKind Kind);
 
   /// Returns true if Phi is a reduction of type Kind and adds it to the
   /// RecurrenceDescriptor. If either \p DB is non-null or \p AC and \p DT are
   /// non-null, the minimal bit width needed to compute the reduction will be
   /// computed.
-  static bool AddReductionVar(PHINode *Phi, RecurKind Kind, Loop *TheLoop,
-                              FastMathFlags FuncFMF,
-                              RecurrenceDescriptor &RedDes,
-                              DemandedBits *DB = nullptr,
-                              AssumptionCache *AC = nullptr,
-                              DominatorTree *DT = nullptr);
+  static bool
+  AddReductionVar(PHINode *Phi, RecurKind Kind, Loop *TheLoop,
+                  FastMathFlags FuncFMF, RecurrenceDescriptor &RedDes,
+                  DemandedBits *DB = nullptr, AssumptionCache *AC = nullptr,
+                  DominatorTree *DT = nullptr, ScalarEvolution *SE = nullptr);
 
   /// Returns true if Phi is a reduction in TheLoop. The RecurrenceDescriptor
   /// is returned in RedDes. If either \p DB is non-null or \p AC and \p DT are
   /// non-null, the minimal bit width needed to compute the reduction will be
-  /// computed.
-  static bool isReductionPHI(PHINode *Phi, Loop *TheLoop,
-                             RecurrenceDescriptor &RedDes,
-                             DemandedBits *DB = nullptr,
-                             AssumptionCache *AC = nullptr,
-                             DominatorTree *DT = nullptr);
+  /// computed. If \p SE is non-null, store instructions to loop invariant
+  /// addresses are processed.
+  static bool
+  isReductionPHI(PHINode *Phi, Loop *TheLoop, RecurrenceDescriptor &RedDes,
+                 DemandedBits *DB = nullptr, AssumptionCache *AC = nullptr,
+                 DominatorTree *DT = nullptr, ScalarEvolution *SE = nullptr);
 
   /// Returns true if Phi is a first-order recurrence. A first-order recurrence
   /// is a non-reduction recurrence relation in which the value of the
   /// recurrence in the current loop iteration equals a value defined in the
   /// previous iteration. \p SinkAfter includes pairs of instructions where the
   /// first will be rescheduled to appear after the second if/when the loop is
   /// vectorized. It may be augmented with additional pairs if needed in order
   /// to handle Phi as a first-order recurrence.
public:
   bool isOrdered() const { return IsOrdered; }
 
   /// Attempts to find a chain of operations from Phi to LoopExitInst that can
   /// be treated as a set of reductions instructions for in-loop reductions.
   SmallVector<Instruction *, 4> getReductionOpChain(PHINode *Phi,
                                                     Loop *L) const;
 
   /// Returns true if the instruction is a call to the llvm.fmuladd intrinsic.
   static bool isFMulAddIntrinsic(Instruction *I) {
     return isa<IntrinsicInst>(I) &&
            cast<IntrinsicInst>(I)->getIntrinsicID() == Intrinsic::fmuladd;
   }
 
+  /// Intermediate store of the reduction
+  StoreInst *IntermediateStore = nullptr;
+
 private:
   // The starting value of the recurrence.
   // It does not have to be zero!
   TrackingVH<Value> StartValue;
   // The instruction who's value is used outside the loop.
   Instruction *LoopExitInstr = nullptr;
   // The kind of the recurrence.
   RecurKind Kind = RecurKind::None;

llvm/include/llvm/Analysis/LoopAccessAnalysis.h

public:
   void print(raw_ostream &OS, unsigned Depth = 0) const;
 
   /// If the loop has memory dependence involving an invariant address, i.e. two
   /// stores or a store and a load, then return true, else return false.
   bool hasDependenceInvolvingLoopInvariantAddress() const {
     return HasDependenceInvolvingLoopInvariantAddress;
   }
 
+  /// Return the list of stores to invariant addresses.
+  const ArrayRef<StoreInst *> getStoresToInvariantAddresses() const {
+    return StoresToInvariantAddresses;
+  }
+
   /// Used to add runtime SCEV checks. Simplifies SCEV expressions and converts
   /// them to a more usable form.  All SCEV expressions during the analysis
   /// should be re-written (and therefore simplified) according to PSE.
   /// A user of LoopAccessAnalysis will need to emit the runtime checks
   /// associated with this predicate.
   const PredicatedScalarEvolution &getPSE() const { return *PSE; }
 
 private:
private:
 
   /// Cache the result of analyzeLoop.
   bool CanVecMem = false;
   bool HasConvergentOp = false;
 
   /// Indicator that there are non vectorizable stores to a uniform address.
   bool HasDependenceInvolvingLoopInvariantAddress = false;
 
+  /// List of stores to invariant addresses.
+  SmallVector<StoreInst *> StoresToInvariantAddresses;
+
   /// The diagnostics report generated for the analysis.  E.g. why we
   /// couldn't analyze the loop.
   std::unique_ptr<OptimizationRemarkAnalysis> Report;
 
   /// If an access has a symbolic strides, this maps the pointer value to
   /// the stride symbol.
   ValueToValueMap SymbolicStrides;
 

llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h

public:
   RecurrenceSet &getFirstOrderRecurrences() { return FirstOrderRecurrences; }
 
   /// Return the set of instructions to sink to handle first-order recurrences.
   MapVector<Instruction *, Instruction *> &getSinkAfter() { return SinkAfter; }
 
   /// Returns the widest induction type.
   Type *getWidestInductionType() { return WidestIndTy; }
 
+  /// Returns True if given invariant store uses recurrent expression
+  bool isRecurringInvariantStore(StoreInst *SI);
+
+  /// Returns True if given address is invariant and is used to store recurrent
+  /// expression
+  bool isRecurringInvariantAddress(Value *V);
+
   /// Returns True if V is a Phi node of an induction variable in this loop.
   bool isInductionPhi(const Value *V) const;
 
   /// Returns a pointer to the induction descriptor, if \p Phi is an integer or
   /// floating point induction.
   const InductionDescriptor *getIntOrFpInductionDescriptor(PHINode *Phi) const;
 
   /// Returns True if V is a cast that is part of an induction def-use chain,

llvm/lib/Analysis/IVDescriptors.cpp

if (Kind == RecurKind::FMulAdd && Exit->getOperand(2) != Phi)
     return false;
 
   LLVM_DEBUG(dbgs() << "LV: Found an ordered reduction: Phi: " << *Phi
                     << ", ExitInst: " << *Exit << "\n");
 
   return true;
 }
 
-bool RecurrenceDescriptor::AddReductionVar(PHINode *Phi, RecurKind Kind,
-                                           Loop *TheLoop, FastMathFlags FuncFMF,
-                                           RecurrenceDescriptor &RedDes,
-                                           DemandedBits *DB,
-                                           AssumptionCache *AC,
-                                           DominatorTree *DT) {
+bool RecurrenceDescriptor::AddReductionVar(
+    PHINode *Phi, RecurKind Kind, Loop *TheLoop, FastMathFlags FuncFMF,
+    RecurrenceDescriptor &RedDes, DemandedBits *DB, AssumptionCache *AC,
+    DominatorTree *DT, ScalarEvolution *SE) {
   if (Phi->getNumIncomingValues() != 2)
     return false;
 
   // Reduction variables are only found in the loop header block.
   if (Phi->getParent() != TheLoop->getHeader())
     return false;
 
   // Obtain the reduction start value from the value that comes from the loop
   // preheader.
   Value *RdxStart = Phi->getIncomingValueForBlock(TheLoop->getLoopPreheader());
 
   // ExitInstruction is the single value which is used outside the loop.
   // We only allow for a single reduction value to be used outside the loop.
   // This includes users of the reduction, variables (which form a cycle
   // which ends in the phi node).
   Instruction *ExitInstruction = nullptr;
+
+  // Variable to keep last visited store instruction. By the end of the
+  // algorithm this variable will be either empty or having intermediate
+  // reduction value stored in invariant address.
+  StoreInst *IntermediateStore = nullptr;
+
   // Indicates that we found a reduction operation in our scan.
   bool FoundReduxOp = false;
 
   // We start with the PHI node and scan for all of the users of this
   // instruction. All users must be instructions that can be used as reduction
   // variables (such as ADD). We must have a single out-of-block user. The cycle
   // must include the original PHI.
   bool FoundStartPHI = false;
bool RecurrenceDescriptor::AddReductionVar(
   //        - One use of reduction value (safe).
   //        - Multiple use of reduction value (not safe).
   //      - PHI:
   //        - All uses of the PHI must be the reduction (safe).
   //        - Otherwise, not safe.
   //  - By instructions outside of the loop (safe).
   //      * One value may have several outside users, but all outside
   //        uses must be of the same value.
+  //  - By store instructions with a loop invariant address (safe with
+  //    the following restrictions):
+  //      * If there are several stores, all must have the same address.
+  //      * Final value should be stored in that loop invariant address.
   //  - By an instruction that is not part of the reduction (not safe).
   //    This is either:
   //      * An instruction type other than PHI or the reduction operation.
   //      * A PHI in the header other than the initial PHI.
   while (!Worklist.empty()) {
     Instruction *Cur = Worklist.pop_back_val();
 
+    // Store instructions are allowed iff it is the store of the reduction
+    // value to the same loop invariant memory location.
+    if (auto *SI = dyn_cast<StoreInst>(Cur)) {
+      if (!SE) {
+        LLVM_DEBUG(dbgs() << "Store instructions are not processed without "
+                          << "Scalar Evolution Analysis\n");
+        return false;
+      }
+
+      const SCEV *PtrScev = SE->getSCEV(SI->getPointerOperand());
+      // Check it is the same address as previous stores
+      if (IntermediateStore) {
+        const SCEV *OtherScev =
+            SE->getSCEV(IntermediateStore->getPointerOperand());
+
+        if (OtherScev != PtrScev) {
+          LLVM_DEBUG(dbgs() << "Storing reduction value to different addresses "
+                            << "inside the loop: " << *SI->getPointerOperand()
+                            << " and "
+                            << *IntermediateStore->getPointerOperand() << '\n');
+          return false;
+        }
+      }
+
+      // Check the pointer is loop invariant
+      if (!SE->isLoopInvariant(PtrScev, TheLoop)) {
+        LLVM_DEBUG(dbgs() << "Storing reduction value to non-uniform address "
+                          << "inside the loop: " << *SI->getPointerOperand()
+                          << '\n');
+        return false;
+      }
+
+      // IntermediateStore is always the last store in the loop.
+      IntermediateStore = SI;
+      continue;
+    }
+
     // No Users.
     // If the instruction has no users then this is a broken chain and can't be
     // a reduction variable.
     if (Cur->use_empty())
       return false;
 
     bool IsAPhi = isa<PHINode>(Cur);
 
for (User *U : Cur->users()) {
         continue;
       }
 
       // Process instructions only once (termination). Each reduction cycle
       // value must only be used once, except by phi nodes and min/max
       // reductions which are represented as a cmp followed by a select.
       InstDesc IgnoredVal(false, nullptr);
       if (VisitedInsts.insert(UI).second) {
-        if (isa<PHINode>(UI))
+        if (isa<PHINode>(UI)) {
           PHIs.push_back(UI);
-        else
+        } else {
+          StoreInst *SI = dyn_cast<StoreInst>(UI);
+          if (SI && SI->getPointerOperand() == Cur) {
+            // Reduction variable chain can only be stored somewhere but it
+            // can't be used as an address.
+            return false;
+          }
           NonPHIs.push_back(UI);
+        }
       } else if (!isa<PHINode>(UI) &&
                  ((!isa<FCmpInst>(UI) && !isa<ICmpInst>(UI) &&
                    !isa<SelectInst>(UI)) ||
                   (!isConditionalRdxPattern(Kind, UI).isRecurrence() &&
                    !isSelectCmpPattern(TheLoop, Phi, UI, IgnoredVal)
                         .isRecurrence() &&
                    !isMinMaxPattern(UI, Kind, IgnoredVal).isRecurrence())))
         return false;
 
       // Remember that we completed the cycle.
       if (UI == Phi)
         FoundStartPHI = true;
     }
     Worklist.append(PHIs.begin(), PHIs.end());
     Worklist.append(NonPHIs.begin(), NonPHIs.end());
   }
 
   // This means we have seen one but not the other instruction of the
   // pattern or more than just a select and cmp. Zero implies that we saw a
   // llvm.min/max instrinsic, which is always OK.
   if (isMinMaxRecurrenceKind(Kind) && NumCmpSelectPatternInst != 2 &&
       NumCmpSelectPatternInst != 0)
     return false;
 
   if (isSelectCmpRecurrenceKind(Kind) && NumCmpSelectPatternInst != 1)
     return false;
 
+  if (IntermediateStore) {
+    // Check that stored value goes to the phi node again. This way we make sure
+    // that the value stored in IntermediateStore is indeed the final reduction
+    // value.
+    if (!is_contained(Phi->operands(), IntermediateStore->getValueOperand())) {
+      LLVM_DEBUG(dbgs() << "Not a final reduction value stored: "
+                        << *IntermediateStore << '\n');
+      return false;
+    }
+
+    // If there is an exit instruction it's value should be stored in
+    // IntermediateStore
+    if (ExitInstruction &&
+        IntermediateStore->getValueOperand() != ExitInstruction) {
+      LLVM_DEBUG(dbgs() << "Last store Instruction of reduction value does not "
+                           "store last calculated value of the reduction: "
+                        << *IntermediateStore << '\n');
+      return false;
+    }
+
+    // If all uses are inside the loop (intermediate stores), then the
+    // reduction value after the loop will be the one used in the last store.
+    if (!ExitInstruction)
+      ExitInstruction = cast<Instruction>(IntermediateStore->getValueOperand());
+  }
+
   if (!FoundStartPHI || !FoundReduxOp || !ExitInstruction)
     return false;
 
   const bool IsOrdered =
       checkOrderedReduction(Kind, ExactFPMathInst, ExitInstruction, Phi);
 
   if (Start != Phi) {
     // If the starting value is not the same as the phi node, we speculatively
bool RecurrenceDescriptor::AddReductionVar(
 
   // We found a reduction var if we have reached the original phi node and we
   // only have a single instruction with out-of-loop users.
 
   // The ExitInstruction(Instruction which is allowed to have out-of-loop users)
   // is saved as part of the RecurrenceDescriptor.
 
   // Save the description of this reduction variable.
-  RecurrenceDescriptor RD(RdxStart, ExitInstruction, Kind, FMF, ExactFPMathInst,
-                          RecurrenceType, IsSigned, IsOrdered, CastInsts,
-                          MinWidthCastToRecurrenceType);
+  RecurrenceDescriptor RD(RdxStart, ExitInstruction, IntermediateStore, Kind,
+                          FMF, ExactFPMathInst, RecurrenceType, IsSigned,
+                          IsOrdered, CastInsts, MinWidthCastToRecurrenceType);
   RedDes = RD;
 
   return true;
 }
 
 // We are looking for loops that do something like this:
 //   int r = 0;
 //   for (int i = 0; i < n; i++) {
bool RecurrenceDescriptor::hasMultipleUsesOf(
   }
 
   return false;
 }
 
 bool RecurrenceDescriptor::isReductionPHI(PHINode *Phi, Loop *TheLoop,
                                           RecurrenceDescriptor &RedDes,
                                           DemandedBits *DB, AssumptionCache *AC,
-                                          DominatorTree *DT) {
+                                          DominatorTree *DT,
+                                          ScalarEvolution *SE) {
   BasicBlock *Header = TheLoop->getHeader();
   Function &F = *Header->getParent();
   FastMathFlags FMF;
   FMF.setNoNaNs(
       F.getFnAttribute("no-nans-fp-math").getValueAsBool());
   FMF.setNoSignedZeros(
       F.getFnAttribute("no-signed-zeros-fp-math").getValueAsBool());
 
-  if (AddReductionVar(Phi, RecurKind::Add, TheLoop, FMF, RedDes, DB, AC, DT)) {
+  if (AddReductionVar(Phi, RecurKind::Add, TheLoop, FMF, RedDes, DB, AC, DT,
+                      SE)) {
     LLVM_DEBUG(dbgs() << "Found an ADD reduction PHI." << *Phi << "\n");
     return true;
   }
-  if (AddReductionVar(Phi, RecurKind::Mul, TheLoop, FMF, RedDes, DB, AC, DT)) {
+  if (AddReductionVar(Phi, RecurKind::Mul, TheLoop, FMF, RedDes, DB, AC, DT,
+                      SE)) {
     LLVM_DEBUG(dbgs() << "Found a MUL reduction PHI." << *Phi << "\n");
     return true;
   }
-  if (AddReductionVar(Phi, RecurKind::Or, TheLoop, FMF, RedDes, DB, AC, DT)) {
+  if (AddReductionVar(Phi, RecurKind::Or, TheLoop, FMF, RedDes, DB, AC, DT,
+                      SE)) {
     LLVM_DEBUG(dbgs() << "Found an OR reduction PHI." << *Phi << "\n");
     return true;
   }
-  if (AddReductionVar(Phi, RecurKind::And, TheLoop, FMF, RedDes, DB, AC, DT)) {
+  if (AddReductionVar(Phi, RecurKind::And, TheLoop, FMF, RedDes, DB, AC, DT,
+                      SE)) {
     LLVM_DEBUG(dbgs() << "Found an AND reduction PHI." << *Phi << "\n");
     return true;
   }
-  if (AddReductionVar(Phi, RecurKind::Xor, TheLoop, FMF, RedDes, DB, AC, DT)) {
+  if (AddReductionVar(Phi, RecurKind::Xor, TheLoop, FMF, RedDes, DB, AC, DT,
+                      SE)) {
     LLVM_DEBUG(dbgs() << "Found a XOR reduction PHI." << *Phi << "\n");
     return true;
   }
-  if (AddReductionVar(Phi, RecurKind::SMax, TheLoop, FMF, RedDes, DB, AC, DT)) {
+  if (AddReductionVar(Phi, RecurKind::SMax, TheLoop, FMF, RedDes, DB, AC, DT,
+                      SE)) {
     LLVM_DEBUG(dbgs() << "Found a SMAX reduction PHI." << *Phi << "\n");
     return true;
   }
-  if (AddReductionVar(Phi, RecurKind::SMin, TheLoop, FMF, RedDes, DB, AC, DT)) {
+  if (AddReductionVar(Phi, RecurKind::SMin, TheLoop, FMF, RedDes, DB, AC, DT,
+                      SE)) {
     LLVM_DEBUG(dbgs() << "Found a SMIN reduction PHI." << *Phi << "\n");
     return true;
   }
-  if (AddReductionVar(Phi, RecurKind::UMax, TheLoop, FMF, RedDes, DB, AC, DT)) {
+  if (AddReductionVar(Phi, RecurKind::UMax, TheLoop, FMF, RedDes, DB, AC, DT,
+                      SE)) {
     LLVM_DEBUG(dbgs() << "Found a UMAX reduction PHI." << *Phi << "\n");
     return true;
   }
-  if (AddReductionVar(Phi, RecurKind::UMin, TheLoop, FMF, RedDes, DB, AC, DT)) {
+  if (AddReductionVar(Phi, RecurKind::UMin, TheLoop, FMF, RedDes, DB, AC, DT,
+                      SE)) {
     LLVM_DEBUG(dbgs() << "Found a UMIN reduction PHI." << *Phi << "\n");
     return true;
   }
   if (AddReductionVar(Phi, RecurKind::SelectICmp, TheLoop, FMF, RedDes, DB, AC,
-                      DT)) {
+                      DT, SE)) {
     LLVM_DEBUG(dbgs() << "Found an integer conditional select reduction PHI."
                       << *Phi << "\n");
     return true;
   }
-  if (AddReductionVar(Phi, RecurKind::FMul, TheLoop, FMF, RedDes, DB, AC, DT)) {
+  if (AddReductionVar(Phi, RecurKind::FMul, TheLoop, FMF, RedDes, DB, AC, DT,
+                      SE)) {
     LLVM_DEBUG(dbgs() << "Found an FMult reduction PHI." << *Phi << "\n");
     return true;
   }
-  if (AddReductionVar(Phi, RecurKind::FAdd, TheLoop, FMF, RedDes, DB, AC, DT)) {
+  if (AddReductionVar(Phi, RecurKind::FAdd, TheLoop, FMF, RedDes, DB, AC, DT,
+                      SE)) {
     LLVM_DEBUG(dbgs() << "Found an FAdd reduction PHI." << *Phi << "\n");
     return true;
   }
-  if (AddReductionVar(Phi, RecurKind::FMax, TheLoop, FMF, RedDes, DB, AC, DT)) {
+  if (AddReductionVar(Phi, RecurKind::FMax, TheLoop, FMF, RedDes, DB, AC, DT,
+                      SE)) {
     LLVM_DEBUG(dbgs() << "Found a float MAX reduction PHI." << *Phi << "\n");
     return true;
   }
-  if (AddReductionVar(Phi, RecurKind::FMin, TheLoop, FMF, RedDes, DB, AC, DT)) {
+  if (AddReductionVar(Phi, RecurKind::FMin, TheLoop, FMF, RedDes, DB, AC, DT,
+                      SE)) {
     LLVM_DEBUG(dbgs() << "Found a float MIN reduction PHI." << *Phi << "\n");
     return true;
   }
   if (AddReductionVar(Phi, RecurKind::SelectFCmp, TheLoop, FMF, RedDes, DB, AC,
-                      DT)) {
+                      DT, SE)) {
     LLVM_DEBUG(dbgs() << "Found a float conditional select reduction PHI."
                       << " PHI." << *Phi << "\n");
     return true;
   }
-  if (AddReductionVar(Phi, RecurKind::FMulAdd, TheLoop, FMF, RedDes, DB, AC,
-                      DT)) {
+  if (AddReductionVar(Phi, RecurKind::FMulAdd, TheLoop, FMF, RedDes, DB, AC, DT,
+                      SE)) {
     LLVM_DEBUG(dbgs() << "Found an FMulAdd reduction PHI." << *Phi << "\n");
     return true;
   }
   // Not a reduction of known type.
   return false;
 }
 
 bool RecurrenceDescriptor::isFirstOrderRecurrence(

llvm/lib/Analysis/LoopAccessAnalysis.cpp

void LoopAccessInfo::analyzeLoop(AAResults *AA, LoopInfo *LI,
 
   // Record uniform store addresses to identify if we have multiple stores
   // to the same address.
   SmallPtrSet<Value *, 16> UniformStores;
 
   for (StoreInst *ST : Stores) {
     Value *Ptr = ST->getPointerOperand();
 
-    if (isUniform(Ptr))
+    if (isUniform(Ptr)) {
+      // Record store instructions to loop invariant addresses
+      StoresToInvariantAddresses.push_back(ST);
       HasDependenceInvolvingLoopInvariantAddress |=
           !UniformStores.insert(Ptr).second;
+    }
 
     // If we did *not* see this pointer before, insert it to  the read-write
     // list. At this phase it is only a 'write' list.
     Type *AccessTy = getLoadStoreType(ST);
     if (Seen.insert({Ptr, AccessTy}).second) {
       ++NumReadWrites;
 
       MemoryLocation Loc = MemoryLocation::get(ST);

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

for (User *U : Inst->users()) {
       if (!TheLoop->contains(UI)) {
         LLVM_DEBUG(dbgs() << "LV: Found an outside user for : " << *UI << '\n');
         return true;
       }
     }
   return false;
 }
 
+/// Returns true if A and B have same pointer operands or same SCEVs addresses
+static bool storeToSameAddress(ScalarEvolution *SE, StoreInst *A,
+                               StoreInst *B) {
+  // Compare store
+  if (A == B)
+    return true;
+
+  // Otherwise Compare pointers
+  Value *APtr = A->getPointerOperand();
+  Value *BPtr = B->getPointerOperand();
+  if (APtr == BPtr)
+    return true;
+
+  // Otherwise compare address SCEVs
+  if (SE->getSCEV(APtr) == SE->getSCEV(BPtr))
+    return true;
+
+  return false;
+}
+
 int LoopVectorizationLegality::isConsecutivePtr(Type *AccessTy,
                                                 Value *Ptr) const {
   const ValueToValueMap &Strides =
       getSymbolicStrides() ? *getSymbolicStrides() : ValueToValueMap();
 
   Function *F = TheLoop->getHeader()->getParent();
   bool OptForSize = F->hasOptSize() ||
                     llvm::shouldOptimizeForSize(TheLoop->getHeader(), PSI, BFI,
for (Instruction &I : *BB) {
           reportVectorizationFailure("Found an invalid PHI",
               "loop control flow is not understood by vectorizer",
               "CFGNotUnderstood", ORE, TheLoop, Phi);
           return false;
         }
 
         RecurrenceDescriptor RedDes;
         if (RecurrenceDescriptor::isReductionPHI(Phi, TheLoop, RedDes, DB, AC,
-                                                 DT)) {
+                                                 DT, PSE.getSE())) {
           Requirements->addExactFPMathInst(RedDes.getExactFPMathInst());
           AllowedExit.insert(RedDes.getLoopExitInstr());
           Reductions[Phi] = RedDes;
           continue;
         }
 
         // TODO: Instead of recording the AllowedExit, it would be good to record the
         // complementary set: NotAllowedExit. These include (but may not be
ORE->emit([&]() {
       return OptimizationRemarkAnalysis(Hints->vectorizeAnalysisPassName(),
                                         "loop not vectorized: ", *LAR);
     });
   }
 
   if (!LAI->canVectorizeMemory())
     return false;
 
+  // We can vectorize stores to invariant address when final reduction value is
+  // guaranteed to be stored at the end of the loop. Also, if decision to
+  // vectorize loop is made, runtime checks are added so as to make sure that
+  // invariant address won't alias with any other objects.
+  if (!LAI->getStoresToInvariantAddresses().empty()) {
+    // For each invariant address, check its last stored value is unconditional.
+    for (StoreInst *SI : LAI->getStoresToInvariantAddresses()) {
+      if (isRecurringInvariantStore(SI) &&
+          blockNeedsPredication(SI->getParent())) {
+        reportVectorizationFailure(
+            "We don't allow storing to uniform addresses",
+            "write of conditional recurring variant value to a loop "
+            "invariant address could not be vectorized",
+            "CantVectorizeStoreToLoopInvariantAddress", ORE, TheLoop);
+        return false;
+      }
+    }
+
     if (LAI->hasDependenceInvolvingLoopInvariantAddress()) {
-    reportVectorizationFailure("Stores to a uniform address",
-        "write to a loop invariant address could not be vectorized",
+      // For each invariant address, check its last stored value is the result
+      // of one of our reductions.
+      //
+      // We do not check if dependence with loads exists because they are
+      // currently rejected earlier in LoopAccessInfo::analyzeLoop. In case this
+      // behaviour changes we have to modify this code.
+      ScalarEvolution *SE = PSE.getSE();
+      SmallVector<StoreInst *, 4> UnhandledStores;
+      for (StoreInst *SI : LAI->getStoresToInvariantAddresses()) {
+        if (isRecurringInvariantStore(SI)) {
+          // Earlier stores to this address are effectively deadcode.
+          erase_if(UnhandledStores, [SE, SI](StoreInst *I) {
+            return storeToSameAddress(SE, SI, I);
+          });
+          continue;
+        }
+        UnhandledStores.push_back(SI);
+      }
+
+      bool IsOK = UnhandledStores.empty();
+      // TODO: we should also validate against InvariantMemSets.
+      if (!IsOK) {
+        reportVectorizationFailure(
+            "We don't allow storing to uniform addresses",
+            "write to a loop invariant address could not "
+            "be vectorized",
             "CantVectorizeStoreToLoopInvariantAddress", ORE, TheLoop);
         return false;
       }
+    }
+  }
 
   Requirements->addRuntimePointerChecks(LAI->getNumRuntimePointerChecks());
   PSE.addPredicate(LAI->getPSE().getPredicate());
   return true;
 }
 
 bool LoopVectorizationLegality::canVectorizeFPMath(
     bool EnableStrictReductions) {
if (!EnableStrictReductions ||
       any_of(getInductionVars(), [&](auto &Induction) -> bool {
         InductionDescriptor IndDesc = Induction.second;
         return IndDesc.getExactFPMathInst();
       }))
     return false;
 
   // We can now only vectorize if all reductions with Exact FP math also
   // have the isOrdered flag set, which indicates that we can move the
-  // reduction operations in-loop.
+  // reduction operations in-loop, and do not have intermediate store.
   return (all_of(getReductionVars(), [&](auto &Reduction) -> bool {
     const RecurrenceDescriptor &RdxDesc = Reduction.second;
-    return !RdxDesc.hasExactFPMath() || RdxDesc.isOrdered();
+    return !RdxDesc.hasExactFPMath() ||
+           (RdxDesc.isOrdered() && !RdxDesc.IntermediateStore);
   }));
 }
 
+bool LoopVectorizationLegality::isRecurringInvariantStore(StoreInst *SI) {
+  for (auto &Reduction : Reductions) {
+    RecurrenceDescriptor &DS = Reduction.second;
+    if (DS.IntermediateStore == SI)
+      return true;
+  }
+  return false;
+}
+
+bool LoopVectorizationLegality::isRecurringInvariantAddress(Value *V) {
+  ScalarEvolution *SE = PSE.getSE();
+  for (auto &Reduction : Reductions) {
+    RecurrenceDescriptor &DS = Reduction.second;
+    if (!DS.IntermediateStore)
+      continue;
+    Value *InvariantAddress = DS.IntermediateStore->getPointerOperand();
+    if (V == InvariantAddress ||
+        SE->getSCEV(V) == SE->getSCEV(InvariantAddress))
+      return true;
+  }
+  return false;
+}
+
 bool LoopVectorizationLegality::isInductionPhi(const Value *V) const {
   Value *In0 = const_cast<Value *>(V);
   PHINode *PN = dyn_cast_or_null<PHINode>(In0);
   if (!PN)
     return false;
 
   return Inductions.count(PN);
 }

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

void InnerLoopVectorizer::scalarizeInstruction(Instruction *Instr,
                                                VPReplicateRecipe *RepRecipe,
                                                const VPIteration &Instance,
                                                bool IfPredicateInstr,
                                                VPTransformState &State) {
   assert(!Instr->getType()->isAggregateType() && "Can't handle vectors");
 
   // llvm.experimental.noalias.scope.decl intrinsics must only be duplicated for
   // the first lane and part.
   if (isa<NoAliasScopeDeclInst>(Instr))
     if (!Instance.isFirstIteration())
       return;
 
   setDebugLocFromInst(Instr);
 
   // Does this instruction return a value ?
   bool IsVoidRetTy = Instr->getType()->isVoidTy();
 
else if (ResumePhi && llvm::is_contained(ResumePhi->blocks(), Incoming))
                               Incoming);
     else
       BCBlockPhi->addIncoming(ReductionStartValue, Incoming);
   }
 
   // Set the resume value for this reduction
   ReductionResumeValues.insert({&RdxDesc, BCBlockPhi});
 
+  // If there were stores of the reduction value to a uniform memory address
+  // inside the loop, create the final store here.
+  if (StoreInst *SI = RdxDesc.IntermediateStore) {
+    StoreInst *NewSI =
+        Builder.CreateStore(ReducedPartRdx, SI->getPointerOperand());
+    propagateMetadata(NewSI, SI);
+
+    // If the reduction value is used in other places,
+    // then let the code below create PHI's for that.
+  }
+
   // Now, we need to fix the users of the reduction variable
   // inside and outside of the scalar remainder loop.
 
   // We know that the loop is in LCSSA form. We need to update the PHI nodes
   // in the exit blocks.  See comment on analogous loop in
   // fixFirstOrderRecurrence for a more complete explaination of the logic.
   if (!Cost->requiresScalarEpilogue(VF))
     for (PHINode &LCSSAPhi : LoopExitBlock->phis())
if (auto *Ptr = getLoadStorePointerOperand(Inst))
     return Legal->isConsecutivePtr(getLoadStoreType(Inst), Ptr);
   return false;
 }
 
 void LoopVectorizationCostModel::collectValuesToIgnore() {
   // Ignore ephemeral values.
   CodeMetrics::collectEphemeralValues(TheLoop, AC, ValuesToIgnore);
 
+  // Find all stores to invariant variables. Since they are going to sink
+  // outside the loop we do not need calculate cost for them.
+  for (BasicBlock *BB : TheLoop->blocks())
+    for (Instruction &I : *BB) {
+      StoreInst *SI;
+      if ((SI = dyn_cast<StoreInst>(&I)) &&
+          Legal->isRecurringInvariantAddress(SI->getPointerOperand()))
+        ValuesToIgnore.insert(&I);
+    }
+
   // Ignore type-promoting instructions we identified during reduction
   // detection.
   for (auto &Reduction : Legal->getReductionVars()) {
     const RecurrenceDescriptor &RedDes = Reduction.second;
     const SmallPtrSetImpl<Instruction *> &Casts = RedDes.getCastInsts();
     VecValuesToIgnore.insert(Casts.begin(), Casts.end());
   }
   // Ignore type-casting instructions we identified during induction
for (Instruction &I : BB->instructionsWithoutDebug()) {
           InductionsToMove.push_back(
               cast<VPWidenIntOrFpInductionRecipe>(Recipe));
         }
         RecipeBuilder.setRecipe(Instr, Recipe);
         VPBB->appendRecipe(Recipe);
         continue;
       }
 
+      // Invariant stores inside loop will be deleted and a single store
+      // with the final reduction value will be added to the exit block
+      StoreInst *SI;
+      if ((SI = dyn_cast<StoreInst>(&I)) &&
+          Legal->isRecurringInvariantAddress(SI->getPointerOperand()))
+        continue;
+
       // Otherwise, if all widening options failed, Instruction is to be
       // replicated. This may create a successor for VPBB.
       VPBasicBlock *NextVPBB =
           RecipeBuilder.handleReplication(Instr, Range, VPBB, Plan);
       if (NextVPBB != VPBB) {
         VPBB = NextVPBB;
         VPBB->setName(BB->hasName() ? BB->getName() + "." + Twine(VPBBsForBB++)
                                     : "");

llvm/lib/Transforms/Vectorize/VPlan.cpp

if (isa<FPMathOperator>(getUnderlyingInstr()))
     O << getUnderlyingInstr()->getFastMathFlags();
   O << " reduce." << Instruction::getOpcodeName(RdxDesc->getOpcode()) << " (";
   getVecOp()->printAsOperand(O, SlotTracker);
   if (getCondOp()) {
     O << ", ";
     getCondOp()->printAsOperand(O, SlotTracker);
   }
   O << ")";
+  if (RdxDesc->IntermediateStore)
+    O << " (with final reduction value stored in invariant address sank "
+         "outside of loop)";
 }
 
 void VPReplicateRecipe::print(raw_ostream &O, const Twine &Indent,
                               VPSlotTracker &SlotTracker) const {
   O << Indent << (IsUniform ? "CLONE " : "REPLICATE ");
 
   if (!getUnderlyingInstr()->getType()->isVoidTy()) {
     printAsOperand(O, SlotTracker);

llvm/test/Transforms/LoopVectorize/AArch64/scalable-reductions.ll

for.body:
   br i1 %exitcond.not, label %for.end, label %for.body, !llvm.loop !0
 
 for.end:
   ret float %.sroa.speculated
 }
 
 ; ADD (with reduction stored in invariant address)
 
-; CHECK-REMARK: loop not vectorized: value that could not be identified as reduction is used outside the loop
+; CHECK-REMARK: vectorized loop (vectorization width: vscale x 4, interleaved count: 2)
 define void @invariant_store(i32* %dst, i32* readonly %src) {
 ; CHECK-LABEL: @invariant_store
-; CHECK-NOT: vector.body
+; CHECK: vector.body:
+; CHECK: %[[LOAD1:.*]] = load <vscale x 4 x i32>
+; CHECK: %[[LOAD2:.*]] = load <vscale x 4 x i32>
+; CHECK: %[[ADD1:.*]] = add <vscale x 4 x i32> %{{.*}}, %[[LOAD1]]
+; CHECK: %[[ADD2:.*]] = add <vscale x 4 x i32> %{{.*}}, %[[LOAD2]]
+; CHECK: call void @llvm.masked.scatter.nxv4i32.nxv4p0i32(<vscale x 4 x i32> %[[ADD1]]
+; CHECK: call void @llvm.masked.scatter.nxv4i32.nxv4p0i32(<vscale x 4 x i32> %[[ADD2]]
+; CHECK: middle.block:
+; CHECK: %[[ADD:.*]] = add <vscale x 4 x i32> %[[ADD2]], %[[ADD1]]
+; CHECK-NEXT: %[[SUM:.*]] = call i32 @llvm.vector.reduce.add.nxv4i32(<vscale x 4 x i32> %[[ADD]])
+; CHECK-NEXT: store i32 %[[SUM]], i32* %gep.dst, align 4
 entry:
   %gep.dst = getelementptr inbounds i32, i32* %dst, i64 42
   store i32 0, i32* %gep.dst, align 4
   br label %for.body
 for.body:
   %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
   %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
   %gep.src = getelementptr inbounds i32, i32* %src, i64 %indvars.iv

llvm/test/Transforms/LoopVectorize/reduction-with-invariant-store.ll

 ; RUN: opt < %s -passes="loop-vectorize" -force-vector-interleave=1 -force-vector-width=4 -S | FileCheck %s
 
 target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
 
 ; This test checks that we can vectorize loop with reduction variable
 ; stored in an invariant address.
 ;
 ; int sum = 0;
 ; for(i=0..N) {
 ;   sum += src[i];
 ;   dst[42] = sum;
 ; }
 ; CHECK-LABEL: @reduc_store
-; CHECK-NOT: vector.body
+; CHECK:       vector.body:
+; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH:%.*]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY:%.*]] ]
+; CHECK-NEXT:    [[VEC_PHI:%.*]] = phi <4 x i32> [ zeroinitializer, [[VECTOR_PH]] ], [ [[TMP4:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT:    [[TMP0:%.*]] = add i64 [[INDEX]], 0
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i32, i32* [[SRC:%.*]], i64 [[TMP0]]
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds i32, i32* [[TMP1]], i32 0
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast i32* [[TMP2]] to <4 x i32>*
+; CHECK-NEXT:    [[WIDE_LOAD:%.*]] = load <4 x i32>, <4 x i32>* [[TMP3]], align 4, !alias.scope !0
+; CHECK-NEXT:    [[TMP4]] = add <4 x i32> [[VEC_PHI]], [[WIDE_LOAD]]
+; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 4
+; CHECK-NEXT:    [[TMP5:%.*]] = icmp eq i64 [[INDEX_NEXT]], 1000
+; CHECK-NEXT:    br i1 [[TMP5]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP3:![0-9]+]]
+; CHECK:       middle.block:
+; CHECK-NEXT:    [[TMP6:%.*]] = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> [[TMP4]])
+; CHECK-NEXT:    store i32 [[TMP6]], i32* [[GEP_DST:%.*]], align 4
+; CHECK-NEXT:    [[CMP_N:%.*]] = icmp eq i64 1000, 1000
+; CHECK-NEXT:    br i1 [[CMP_N]], label [[EXIT:%.*]], label [[SCALAR_PH:%.*]]
 define void @reduc_store(i32* %dst, i32* readonly %src) {
 entry:
   %gep.dst = getelementptr inbounds i32, i32* %dst, i64 42
   store i32 0, i32* %gep.dst, align 4
   br label %for.body
 
 for.body:
   %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]

 ; Same as above but with floating point numbers instead.
 ;
 ; float sum = 0;
 ; for(i=0..N) {
 ;   sum += src[i];
 ;   dst[42] = sum;
 ; }
 ; CHECK-LABEL: @reduc_store_fadd_fast
-; CHECK-NOT: vector.body
+; CHECK: vector.body:
+; CHECK: phi <4 x float>
+; CHECK: load <4 x float>
+; CHECK: fadd fast <4 x float>
+; CHECK-NOT: store float %{{[0-9]+}}, float* %gep.dst
+; CHECK: middle.block:
+; CHECK-NEXT: [[TMP:%.*]] = call fast float @llvm.vector.reduce.fadd.v4f32
+; CHECK-NEXT: store float %{{[0-9]+}}, float* %gep.dst
 define void @reduc_store_fadd_fast(float* %dst, float* readonly %src) {
 entry:
   %gep.dst = getelementptr inbounds float, float* %dst, i64 42
   store float 0.000000e+00, float* %gep.dst, align 4
   br label %for.body
 
 for.body:
   %sum = phi float [ 0.000000e+00, %entry ], [ %add, %for.body ]

 ;  int sum = 0;
 ;  for(int i=0; i < 1000; i+=2) {
 ;    sum += src[i];
 ;    dst[42] = sum;
 ;    sum += src[i+1];
 ;    dst[42] = sum;
 ;  }
 ; CHECK-LABEL: @reduc_store_inside_unrolled
-; CHECK-NOT: vector.body
+; CHECK:       vector.body:
+; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH:%.*]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY:%.*]] ]
+; CHECK-NEXT:    [[VEC_IND:%.*]] = phi <4 x i64> [ <i64 0, i64 2, i64 4, i64 6>, [[VECTOR_PH]] ], [ [[VEC_IND_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT:    [[VEC_PHI:%.*]] = phi <4 x i32> [ zeroinitializer, [[VECTOR_PH]] ], [ [[TMP34:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT:    [[OFFSET_IDX:%.*]] = mul i64 [[INDEX]], 2
+; CHECK-NEXT:    [[TMP0:%.*]] = add i64 [[OFFSET_IDX]], 0
+; CHECK-NEXT:    [[TMP1:%.*]] = add i64 [[OFFSET_IDX]], 2
+; CHECK-NEXT:    [[TMP2:%.*]] = add i64 [[OFFSET_IDX]], 4
+; CHECK-NEXT:    [[TMP3:%.*]] = add i64 [[OFFSET_IDX]], 6
+; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr inbounds i32, i32* [[SRC:%.*]], i64 [[TMP0]]
+; CHECK-NEXT:    [[TMP5:%.*]] = getelementptr inbounds i32, i32* [[SRC]], i64 [[TMP1]]
+; CHECK-NEXT:    [[TMP6:%.*]] = getelementptr inbounds i32, i32* [[SRC]], i64 [[TMP2]]
+; CHECK-NEXT:    [[TMP7:%.*]] = getelementptr inbounds i32, i32* [[SRC]], i64 [[TMP3]]
+; CHECK-NEXT:    [[TMP8:%.*]] = load i32, i32* [[TMP4]], align 4, !alias.scope !11
+; CHECK-NEXT:    [[TMP9:%.*]] = load i32, i32* [[TMP5]], align 4, !alias.scope !11
+; CHECK-NEXT:    [[TMP10:%.*]] = load i32, i32* [[TMP6]], align 4, !alias.scope !11
+; CHECK-NEXT:    [[TMP11:%.*]] = load i32, i32* [[TMP7]], align 4, !alias.scope !11
+; CHECK-NEXT:    [[TMP12:%.*]] = insertelement <4 x i32> poison, i32 [[TMP8]], i32 0
+; CHECK-NEXT:    [[TMP13:%.*]] = insertelement <4 x i32> [[TMP12]], i32 [[TMP9]], i32 1
+; CHECK-NEXT:    [[TMP14:%.*]] = insertelement <4 x i32> [[TMP13]], i32 [[TMP10]], i32 2
+; CHECK-NEXT:    [[TMP15:%.*]] = insertelement <4 x i32> [[TMP14]], i32 [[TMP11]], i32 3
+; CHECK-NEXT:    [[TMP16:%.*]] = add <4 x i32> [[TMP15]], [[VEC_PHI]]
+; CHECK-NEXT:    [[TMP17:%.*]] = or <4 x i64> [[VEC_IND]], <i64 1, i64 1, i64 1, i64 1>
+; CHECK-NEXT:    [[TMP18:%.*]] = extractelement <4 x i64> [[TMP17]], i32 0
+; CHECK-NEXT:    [[TMP19:%.*]] = getelementptr inbounds i32, i32* [[SRC]], i64 [[TMP18]]
+; CHECK-NEXT:    [[TMP20:%.*]] = extractelement <4 x i64> [[TMP17]], i32 1
+; CHECK-NEXT:    [[TMP21:%.*]] = getelementptr inbounds i32, i32* [[SRC]], i64 [[TMP20]]
+; CHECK-NEXT:    [[TMP22:%.*]] = extractelement <4 x i64> [[TMP17]], i32 2
+; CHECK-NEXT:    [[TMP23:%.*]] = getelementptr inbounds i32, i32* [[SRC]], i64 [[TMP22]]
+; CHECK-NEXT:    [[TMP24:%.*]] = extractelement <4 x i64> [[TMP17]], i32 3
+; CHECK-NEXT:    [[TMP25:%.*]] = getelementptr inbounds i32, i32* [[SRC]], i64 [[TMP24]]
+; CHECK-NEXT:    [[TMP26:%.*]] = load i32, i32* [[TMP19]], align 4, !alias.scope !11
+; CHECK-NEXT:    [[TMP27:%.*]] = load i32, i32* [[TMP21]], align 4, !alias.scope !11
+; CHECK-NEXT:    [[TMP28:%.*]] = load i32, i32* [[TMP23]], align 4, !alias.scope !11
+; CHECK-NEXT:    [[TMP29:%.*]] = load i32, i32* [[TMP25]], align 4, !alias.scope !11
+; CHECK-NEXT:    [[TMP30:%.*]] = insertelement <4 x i32> poison, i32 [[TMP26]], i32 0
+; CHECK-NEXT:    [[TMP31:%.*]] = insertelement <4 x i32> [[TMP30]], i32 [[TMP27]], i32 1
+; CHECK-NEXT:    [[TMP32:%.*]] = insertelement <4 x i32> [[TMP31]], i32 [[TMP28]], i32 2
+; CHECK-NEXT:    [[TMP33:%.*]] = insertelement <4 x i32> [[TMP32]], i32 [[TMP29]], i32 3
+; CHECK-NEXT:    [[TMP34]] = add <4 x i32> [[TMP33]], [[TMP16]]
+; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 4
+; CHECK-NEXT:    [[VEC_IND_NEXT]] = add <4 x i64> [[VEC_IND]], <i64 8, i64 8, i64 8, i64 8>
+; CHECK-NEXT:    [[TMP35:%.*]] = icmp eq i64 [[INDEX_NEXT]], 500
+; CHECK-NEXT:    br i1 [[TMP35]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP14:![0-9]+]]
+; CHECK:       middle.block:
+; CHECK-NEXT:    [[TMP36:%.*]] = call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> [[TMP34]])
+; CHECK-NEXT:    store i32 [[TMP36]], i32* [[GEP_DST:%.*]], align 4
+; CHECK-NEXT:    [[CMP_N:%.*]] = icmp eq i64 500, 500
+; CHECK-NEXT:    br i1 [[CMP_N]], label [[EXIT:%.*]], label [[SCALAR_PH:%.*]]
 define void @reduc_store_inside_unrolled(i32* %dst, i32* readonly %src) {
 entry:
   %gep.dst = getelementptr inbounds i32, i32* %dst, i64 42
   br label %for.body
 
 for.body:
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
-  %sum = phi i32 [ 0, %entry ], [ %sum.1, %for.body ]
+  %sum = phi i32 [ 0, %entry ], [ %sum.2, %for.body ]
   %gep.src = getelementptr inbounds i32, i32* %src, i64 %iv
   %0 = load i32, i32* %gep.src, align 4
   %sum.1 = add nsw i32 %0, %sum
   store i32 %sum.1, i32* %gep.dst, align 4
   %1 = or i64 %iv, 1
   %gep.src.1 = getelementptr inbounds i32, i32* %src, i64 %1
   %2 = load i32, i32* %gep.src.1, align 4
   %sum.2 = add nsw i32 %2, %sum.1
   store i32 %sum.2, i32* %gep.dst, align 4
   %iv.next = add nuw nsw i64 %iv, 2
   %cmp = icmp slt i64 %iv.next, 1000
   br i1 %cmp, label %for.body, label %exit
 
 exit:
   ret void
 }
 
+; Check that we cannot vectorize code if stored value is not the final reduction
+; value
+;
+;  int sum = 0;
+;  for(int i=0; i < 1000; i+=2) {
+;    sum += src[i];
+;    dst[42] = sum + 1;
+;  }
+; CHECK-LABEL: @reduc_store_not_final_value
+; CHECK-NOT: vector.body:
+define void @reduc_store_not_final_value(i32* %dst, i32* readonly %src) {
+entry:
+  %gep.dst = getelementptr inbounds i32, i32* %dst, i64 42
+  store i32 0, i32* %gep.dst, align 4
+  br label %for.body
+
+for.body:
+  %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+  %gep.src = getelementptr inbounds i32, i32* %src, i64 %iv
+  %0 = load i32, i32* %gep.src, align 4
+  %add = add nsw i32 %sum, %0
+  %sum_plus_one = add i32 %add, 1
+  store i32 %sum_plus_one, i32* %gep.dst, align 4
+  %iv.next = add nuw nsw i64 %iv, 1
+  %exitcond = icmp eq i64 %iv.next, 1000
+  br i1 %exitcond, label %exit, label %for.body
+
+exit:
+  ret void
+}
+
 ; We cannot vectorize if two (or more) invariant stores exist in a loop.
 ;
 ;  int sum = 0;
 ;  for(int i=0; i < 1000; i+=2) {
 ;    sum += src[i];
 ;    dst[42] = sum;
 ;    sum += src[i+1];
 ;    other_dst[42] = sum;
 ;  }
 ; CHECK-LABEL: @reduc_double_invariant_store
 ; CHECK-NOT: vector.body:
 define void @reduc_double_invariant_store(i32* %dst, i32* %other_dst, i32* readonly %src) {
 entry:
   %gep.dst = getelementptr inbounds i32, i32* %dst, i64 42
   %gep.other_dst = getelementptr inbounds i32, i32* %other_dst, i64 42
   br label %for.body
 
 for.body:
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
-  %sum = phi i32 [ 0, %entry ], [ %sum.1, %for.body ]
+  %sum = phi i32 [ 0, %entry ], [ %sum.2, %for.body ]
   %arrayidx = getelementptr inbounds i32, i32* %src, i64 %iv
   %0 = load i32, i32* %arrayidx, align 4
   %sum.1 = add nsw i32 %0, %sum
   store i32 %sum.1, i32* %gep.dst, align 4
   %1 = or i64 %iv, 1
   %arrayidx4 = getelementptr inbounds i32, i32* %src, i64 %1
   %2 = load i32, i32* %arrayidx4, align 4
   %sum.2 = add nsw i32 %2, %sum.1

 ;  for(int i=0; i < 1000; i+=2) {
 ;    sum += src[i];
 ;    if (src[i+1] > 0)
 ;      dst[42] = sum;
 ;    sum += src[i+1];
 ;    dst[42] = sum;
 ;  }
 ; CHECK-LABEL: @reduc_store_middle_store_predicated
-; CHECK-NOT: vector.body
+; CHECK: vector.body:
+; CHECK-NOT: store i32 %{{[0-9]+}}, i32* %gep.dst
+; CHECK: middle.block:
+; CHECK-NEXT: [[TMP:%.*]] = call i32 @llvm.vector.reduce.add.v4i32
+; CHECK-NEXT: store i32 [[TMP]], i32* %gep.dst
+; CHECK: ret void
 define void @reduc_store_middle_store_predicated(i32* %dst, i32* readonly %src) {
 entry:
   %gep.dst = getelementptr inbounds i32, i32* %dst, i64 42
   br label %for.body
 
 for.body:                                         ; preds = %latch, %entry
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %latch ]
   %sum = phi i32 [ 0, %entry ], [ %sum.2, %latch ]
latch:                                            ; preds = %predicated, %for.body
   %iv.next = add nuw nsw i64 %iv, 2
   %cmp = icmp slt i64 %iv.next, 1000
   br i1 %cmp, label %for.body, label %exit
 
 exit:                                 ; preds = %latch
   ret void
 }
 
+; Final reduction value is overwritten inside loop
+;
+; for(int i=0; i < 1000; i++) {
+;   sum += src[i];
+;   dst[42] = sum;
+;   dst[42] = 0;
+; }
+; CHECK-LABEL: @reduc_store_final_store_overwritten
+; CHECK-NOT: vector.body:
+define void @reduc_store_final_store_overwritten(i32* %dst, i32* readonly %src) {
+entry:
+  %gep.dst = getelementptr inbounds i32, i32* %dst, i64 42
+  br label %for.body
+
+for.body:
+  %sum = phi i32 [ 0, %entry ], [ %add, %for.body ]
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+  %gep.src = getelementptr inbounds i32, i32* %src, i64 %iv
+  %0 = load i32, i32* %gep.src, align 4
+  %add = add nsw i32 %sum, %0
+  store i32 %add, i32* %gep.dst, align 4
+  store i32 0, i32* %gep.dst, align 4
+  %iv.next = add nuw nsw i64 %iv, 1
+  %exitcond = icmp eq i64 %iv.next, 1000
+  br i1 %exitcond, label %exit, label %for.body
+
+exit:
+  ret void
+}
+
 ; Final value used outside of loop does not prevent vectorization
 ;
 ; int sum = 0;
 ; for(int i=0; i < 1000; i++) {
 ;   sum += src[i];
 ;   dst[42] = sum;
 ; }
 ; dst[43] = sum;
 ; CHECK-LABEL: @reduc_store_inoutside
-; CHECK-NOT: vector.body
+; CHECK: vector.body:
+; CHECK-NOT: store i32 %{{[0-9]+}}, i32* %gep.src
+; CHECK: middle.block:
+; CHECK-NEXT: [[TMP:%.*]] = call i32 @llvm.vector.reduce.add.v4i32
+; CHECK-NEXT: store i32 [[TMP]], i32* %gep.dst
+; CHECK: exit:
+; CHECK: [[PHI:%.*]] = phi i32 [ [[TMP1:%.*]], %for.body ], [ [[TMP2:%.*]], %middle.block ]
+; CHECK: [[ADDR:%.*]] = getelementptr inbounds i32, i32* %dst, i64 43
+; CHECK: store i32 [[PHI]], i32* [[ADDR]]
+; CHECK: ret void
 define void @reduc_store_inoutside(i32* %dst, i32* readonly %src) {
 entry:
   %gep.dst = getelementptr inbounds i32, i32* %dst, i64 42
   br label %for.body
 
 for.body:
   %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
   %sum = phi i32 [ 0, %entry ], [ %sum.1, %for.body ]
   %arrayidx = getelementptr inbounds i32, i32* %src, i64 %iv
   %0 = load i32, i32* %arrayidx, align 4
   %sum.1 = add nsw i32 %0, %sum
   store i32 %sum.1, i32* %gep.dst, align 4
   %iv.next = add nuw nsw i64 %iv, 1
   %exitcond = icmp eq i64 %iv.next, 1000
   br i1 %exitcond, label %exit, label %for.body
 
 exit:
   %sum.lcssa = phi i32 [ %sum.1, %for.body ]
   %gep.dst.1 = getelementptr inbounds i32, i32* %dst, i64 43
   store i32 %sum.lcssa, i32* %gep.dst.1, align 4
   ret void
 }

llvm/test/Transforms/LoopVectorize/vplan-printing.ll

for.body:                                         ; preds = %entry, %for.body
   %iv.next = add i64 %iv, 1
   %exitcond = icmp eq i64 %iv.next, %n
   br i1 %exitcond, label %for.end, label %for.body
 
 for.end:                                          ; preds = %for.body, %entry
   ret float %red.next
 }
 
+define void @print_reduction_with_invariant_store(i64 %n, float* noalias %y, float* noalias %dst) {
+; CHECK-LABEL: Checking a loop in 'print_reduction_with_invariant_store'
+; CHECK:      VPlan 'Initial VPlan for VF={4},UF>=1' {
+; CHECK-NEXT: Live-in vp<[[VEC_TC:%.+]]> = vector-trip-count
+; CHECK-EMPTY:
+; CHECK-NEXT: <x1> vector loop: {
+; CHECK-NEXT: for.body:
+; CHECK-NEXT:   EMIT vp<[[CAN_IV:%.+]]> = CANONICAL-INDUCTION
+; CHECK-NEXT:   WIDEN-REDUCTION-PHI ir<%red> = phi ir<0.000000e+00>, ir<%red.next>
+; CHECK-NEXT:   vp<[[IV:%.+]]> = SCALAR-STEPS vp<[[CAN_IV]]>, ir<0>, ir<1>
+; CHECK-NEXT:   CLONE ir<%arrayidx> = getelementptr ir<%y>, vp<[[IV]]>
+; CHECK-NEXT:   WIDEN ir<%lv> = load ir<%arrayidx>
+; CHECK-NEXT:   REDUCE ir<%red.next> = ir<%red> + fast reduce.fadd (ir<%lv>) (with final reduction value stored in invariant address sank outside of loop)
+; CHECK-NEXT:   EMIT vp<[[CAN_IV_NEXT:%.+]]> = VF * UF +(nuw) vp<[[CAN_IV]]>
+; CHECK-NEXT:   EMIT branch-on-count vp<[[CAN_IV_NEXT]]> vp<[[VEC_TC]]>
+; CHECK-NEXT: No successors
+; CHECK-NEXT: }
+; CHECK-NEXT: No successors
+; CHECK-NEXT: }
+;
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %iv = phi i64 [ %iv.next, %for.body ], [ 0, %entry ]
+  %red = phi float [ %red.next, %for.body ], [ 0.0, %entry ]
+  %arrayidx = getelementptr inbounds float, float* %y, i64 %iv
+  %lv = load float, float* %arrayidx, align 4
+  %red.next = fadd fast float %lv, %red
+  store float %red.next, float* %dst, align 4
+  %iv.next = add i64 %iv, 1
+  %exitcond = icmp eq i64 %iv.next, %n
+  br i1 %exitcond, label %for.end, label %for.body
+
+for.end:                                          ; preds = %for.body, %entry
+  ret void
+}
+
 define void @print_replicate_predicated_phi(i64 %n, i64* %x) {
 ; CHECK-LABEL: Checking a loop in 'print_replicate_predicated_phi'
 ; CHECK:      VPlan 'Initial VPlan for VF={4},UF>=1' {
 ; CHECK-NEXT: Live-in vp<[[VEC_TC:%.+]]> = vector-trip-count
 ; CHECK-EMPTY:
 ; CHECK-NEXT: <x1> vector loop: {
 ; CHECK-NEXT: for.body:
 ; CHECK-NEXT:   EMIT vp<[[CAN_IV:%.+]]> = CANONICAL-INDUCTION