Update MergedLoadStoreMotion to use MemorySSA

[• dberlin]

Summary

This updates MergedLoadStoreMotion to use and preserve MemorySSA.
It depends on D7864

Prior to this, the algorithm for loads was N^2, and for stores, it was N^2R, where N is the number of instructions in
the block, and R is the number of removed stores
(It restarted the reverse walk every time it removed a store due to iterator invalidation).

It is now O(M) (where M is the number of memory instructions in the two blocks) for loads,
and O(max(M,S^2)) for stores (because we have no downwards clobbering API yet, the hash table does not help
us determine memory dependence for our uses).

I have deliberately not changed behavior in terms of what loads/stores it will remove or the compile time
controls, in order to make the change minimal.

(Hopefully phabricator will not screwup a revision that is a branch of a branch,
arc diff --preview showed the right stuff)

[• dberlin]

Delete dead functions
Trigger an update so phabricator sends out this review, since it got eaten during sendgrid issues

[• dberlin]

Update for MemorySSA API Update

[reames]

FYI, I'm basically reviewing this as a new algorithm. Trying to match it up with what's there before doesn't seem particularly worth doing given the algorithmic problems you commented on.

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

106

Can you use lambdas for these? I'm not entirely sure that works, but if it does, the code would be cleaner.

Also, I'm surprised that a generic instruction hash function doesn't already exist. What does GVN use for this?

117

Is the extra A == B check needed? I would expect that to be inside isSameOperationAs.

324

Not quite clear what you're trying to say with this comment? The code appears to just be updating MemorySSA for the transform performed.

426

I'm a bit confused here. If we found the set of load instructions which share the generation number at the start of a block and we're merging two blocks with a single shared successor, why do we need any further legality checks? Aren't all the loads guaranteed to have the generation number at the end of the common block?

Are you possibly trying to do two generations at once here? Or something else more complex?

532

Are you guaranteed there's not another phi here? Not entirely sure what this is doing. It may be completely correct, but it looks suspicious when compared to instruction insertion into a basic block.

571

Given this is a linear traversal of Pred0, does the reserve really gain you anything here?

581

Neat!

611–614

Ok, I'm again a bit confused by the algorithm. :)

What I expected to see was something along the following:

• Find store in each block which is a candidate for merge (using memory-def edges from MemoryPhi)
• Merge any loads in those basic blocks which use the generation if possible.
• Search downward from each store to find a conflicting memory access or end of block. (This could either by by walking instructions or mem-use edges.)
• If both reach end of block, merge.
• Repeat previous steps with newly created MemoryPhi for earlier stores if any.

[• dberlin]

During building, MemorySSA currently optimizes uses to the nearest
dominating clobbering access for you, done at the suggestion of 2
people. We could make it stop doing this (it's just a flag), but the
tradeoff is most things end up walking more in practice.
Additionally, as things optimize, and you update memssa, you will
often end up with the *same result* anyway. The only way to stop that
is to make it so you are required to use "nearest dominating
memorydef" instead of "dominating memorydef" when updating. We do
this in the API's where we do insertion (addNewMemoryUse), but we
don't check/verify it in replaceMemoryAccess style API's, we only
check standard domination.

The other tradeoff, btw, is that while doing MemoryUse optimization
while building makes this particular algorithm slightly harder, it
makes an (IMHO) much more useful thing a lot easier.

A lot of passes (GVN, memcpyopt, etc) want to know whether they can
replace a load with a load to eliminate something:

1= MemoryDef(liveonentry)
store a
2 = MemoryDef(1)
store b
3 = MemoryDef(2)
store c
4 = MemoryDef(3)
store d
MemoryUse(1)
load a, 4 bytes
load a, 8 bytes
(or different types, or whatever)

By optimizing MemoryUse's during building, we are guaranteed that if
we do getMemoryAccess(load A)->definition(which is store a)->uses, or
we now have all loads that actually use *that* store's value. This
means we can look at the other loads of that store value and see if we
can reuse them.

Doing this without use optimization of uses would be much harder.
you'd need to use a downwards API from the store to get all possible
real uses, which may encompass walking the entire program.

On the other hand, the tradeoff of doing MemoryUse optimization is
that it means to get all the loads in a block, you really have to ask
for all the loads in the block, instead of trying to hope you have
chains that give it to you :)

[• dberlin]

• Update for new insertion API. Clean up comments a bit
• Update to walk stores

This now uses an immediate use walker rather than the alias ABI to determine downward exposure of stores

[reames]

Waiting for update or possibly submission. Have lost track.

[• dberlin]

I have a large update coming to this, was just trying to make a sane update
API (the rest is easy).

[• dberlin]

Update to work in progress
This will work for loads, stores coming up.

This also includes the memoryssa update API, which will be split out and committed separately

[• dberlin]

• Add support for stores. Code generated is correct but updating of memoryssa is wrong and will be fixed

[• dberlin]

• Add verification for completeness of PHI nodes
• Add basic PHI creation API
• Update MemorySSA for store sinking

[• dberlin]

This version should correctly update memoryssa for store sinking :)
(I checked that it does on the llvm testcases, i'll test it further as i
split out the memoryssa part)

[• dberlin]

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

106

To answer the original question, i'm checking to see about lambdas.

No generic Instruction hasher exists that covers precisely the set of stuff we care about. GVN does what newgvn does, which is it defines its own expression class, shoves operands in those classes, and calls hash_combine on the whole shebang:

friend hash_code hash_value(const Expression &Value) {
    return hash_combine(
        Value.opcode, Value.type,
        hash_combine_range(Value.varargs.begin(), Value.varargs.end()));
  }

The default hasher will end up hashing the pointer values themselves, which is what gvn wanted in this case.

It is not useful, however, for hashing "parts of operations you care about", or "do two operations look pretty similar".

Past that, i did not change the things included in the hash to ensure consistency of what is optimized with the existing code.

117

It is not, it goes straight into comparing all sorts of fields/data.

[gberry]

A couple of minor comments I noticed when reading through this for my own edification:

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

650

It seems like you forgot to finish this sentence.

979

typo: "MemoryDefwith"

991

typo: "replaceALlUsesWith"

[• dberlin]

• Update MergedLoadStoreMotion to optionally use MemorySSA
• Add support for stores. Code generated is correct but updating of memoryssa is wrong
• Add verification for completeness of PHI nodes
• Add basic PHI creation API
• Update MemorySSA for store sinking
• Fix typos

[• dberlin]

rebased

[mcrosier]

A few comments in passing.

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

99

Should we try and unique the name a bit more (e.g., "use-memssa-mlsm") under the assumption that other passes will have a command-line option to opt into using MemorySSA? Alternatively, we could include the option in the PassManagerBuilder and use this single flag to control the use of MemorySSA for all passes during the transition period. Just my 2c.

158

Perhaps I missed something, but why is TargetLibraryInfo required?

158–164

Should the DominatorTree be added as preserved?

526

No need for extra curly brackets.

710

but -> put

725

Pre-increment here?

732

Any reason we don't just increment in the for statement itself?

737

Please add a period.

821

Please add a period.

[• dberlin]

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

99

I'll rename it to use-memoryssa-mldst (or whatever the shortname for this pass is).

A single flag would be nice, but would make it much harder to turn memoryssa on and off for each pass, and the point of the flag is to be able to test/debug regressions until we are satisified and get rid of the non-memoryssa version.

(ala what happened with SROA, etc and the ssaupdater)

158

It is not, this is a merge error.

158–164

This is not necessary, setPreservesCFG will preserve it, and all other passes marked CFG-only

725

Not sure what line this is supposed to go with, i can't find a post-increment around :)

732

Fixed. Originally, the updater was modifying the accesslist we were walking on and invalidating the iterator.
It turns out to be non-trivial to change this, but i never moved this loop back to account for this.

[• dberlin]

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

732

Actually, it can't be done, we still invalidate the current iterator for loads.
I'll document this.

[• dberlin]

Update for requested changes

[mcrosier]

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

99

I agree we should have a flag per pass to simplify debugging.

725

Sorry if I wasn't clear. I was suggesting we remove the pre-increment above and do it in the if statement.

++NLoads;
if (NLoads * Size1 ...)

vs

if (++NLoads * Size1 ...)

[• dberlin]

Rebase for passmanager conversion

[gberry]

It looks like this change still contains all of the MemorySSA changes that were split out into D21463

[• dberlin]

Properly rebase on top of other patch

[• dberlin]

Should be fixed now.

[gberry]

Would it make sense to split out the DFS numbering code motion barrier code into a separate change?

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

896

Typo: "explosed".
Also, there is no 'End' parameter as referenced in the comment.

[• dberlin]

Rebase after dependent revision commited

[• dberlin]

Also run verification pass with new pass manager for tests.

[• dberlin]

Any other comments?
If not, because it's off by default (and it looked like there are no real objections), my plan was to commit it.

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

725

Oh, sorry. I copied this code from the non-memoryssa version.
(In truth, the code can stand a lot of cleaning up, improvement, and reducing of restrictions)

[gberry]

I'm still digesting this, but here are a couple more nits I noticed.
I don't want to hold up approval from anyone else though.

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

680

Isn't MemorySSA only preserved if 'UseMemorySSA' is true?

901

getDFSNumIn/getDFSNumOut are documented as: These are an internal implementation detail, do not call them.
which seems to conflict with this usage.

[• dberlin]

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

680

Yes, fixed.

901

I will change the docs for them in a separate patch.

The change to make updateDFSNumbers public, etc, was done specifically to enable this usage.
See the thread titlte "[PATCH] Make updateDFSNumbers API public" from last year, and there was a review where i computed them separately, and it was suggested to just make it public.

[• dberlin]

(and happy to wait if you want to review it)

[george.burgess.iv]

Drive-by comments about comments (insert "we need to go deeper" meme here)

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

710

Potentially unfinished thought?

716

Did you mean "only move simple (non-atomic, non-volatile)"?

[• dberlin]

• Update for review comments on comments

[• dberlin]

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

901

This is now fixed in the docs for those functions

[Gerolf]

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

292–293

Loc0,1 should be moved below the conditional. Then they are definitely needed.

323

This is a MSSA update problem. How about
if (MSSA && MSSA->update(HoistCan, HoistedInst))

MSSA->removeMemoryAccess(ElseInst);

526

You could remove the conditional here and simply return from updateMemorySSA when MSSA is null.

650

How about: Visit blocks in DFS order and number instructions. For each block track first hoist barrier and last sink barrier.

681

I think this is the only place where you need UseMemorySSA. All other instances can be handled by if (MSSA).

951

Shouldn't updates be owned by MSSA rather than each client?

[• dberlin]

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

323

As mentioned, if you look early enough in the thread on this, you can see at one point it was, in fact that API, and i was explicitly asked to make simpler apis instead of these kinds of composed ones :)

The way it is now is also consistent with how we do updates for everything else both SSA related (where you do exactly what you see below), and analysis related, in LLVM.
(IE memdep has you call invalidateCachedPointerInfo and such, dominators requires you set and redirect dominators on your own)

In fact, if memoryssa was part of the IR, there would be no difference between the way this code looks and the rest of the function (which is updating the IR) looks . Because it's a side datastructure the naming is a little different and we haven't built all the same utilities yet.

Additonally, the kind of API you suggest above would have wildly varying complexity and is overkill vs building utilities that handle the common cases, which is what we do elsewhere for ssa/dom/etc updates.

I'm more than happy to commit to building those utilities as we come across common cases.

951

(replying here since phab is not good at processing email like this):

That would be inconsistent with what we do in the rest of LLVM :)
MemorySSA is an SSA form, and like SSA in LLVM, passes are responsible for keeping SSA up to date. This is also true of most analysis preservation.
You can see that splitCriticalEdge takes a memdep argument, for example, and will update memdep instead of memdep updating itself.
Same with dominators. Passes that want to preserve Dom are responsible for redirecting the dominators where they should go. They can screw it up, too!

Doing general updates (IE "I have no idea what happened, you go find out and fix it") is also expensive, and at least for every pass i've converted so far, completely unnecessary.

I also had more functional updaters (IE replaceAccessWithHoistedAccess) in memoryssa at earlier points, but if you look earlier in the review, you'll see others wanted it to be simple functionality instead.
In retrospect, i agree.
We can build a general updater if we need it, and we can build utilities to handle common updates if we need it. So far, it hasn't been necessary.

At some point, the better question is really "why is memoryssa a side data structure instead of just part of the normal IR". But that's a question for another day.

[• dberlin]

• Update for review comments

[• dberlin]

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

951

(and, just to also point out, the update algorithm we use is specific to diamonds, and would not work for more general control flow. It also would not work without the set of legality testing, etc, we perform elsewhere in the code )

[Gerolf]

I'll take a more in depth look into the load and store merge routines also. At a first glance it seems one could just add a few MSSA hooks rather than copy-paste-modify the code base.
I also take a (late :-() look at the core MSSA design. Could you outline the high-level MSSA design in more detail than http://reviews.llvm.org/D7864?Eg. how does it compare/relate to http://www.airs.com/dnovillo/Papers/mem-ssa.pdf? Thanks!

-Gerolf

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

323

Ok, it would probably not be possible foresee all use cases. I certainly agree to the iterative approach you propose. I'm surprised though that a simple replacement of a load has not come up elsewhere.

951

What got me here is that I expected the update to be simple. Why is here more involved than replacing two MD's with one, possibly removing a MP? From the clients perspective updates like this look too complicated to get right. Maybe there is verifier in place, but in the code I don't even see a single assertion but could help debugging issues should something go wrong. Perhaps we can come up with some utility routines that make simpler to grasp: when two memory or more memory operation get moved, this is the blueprint for updating MSSA and guarantee it stays consistent.

1018

There could be a lambda or function that takes care of the two loops.

[• dberlin]

Also, ust to give you some idea of general update complexity: If you don't
bound the problem with some common cases, it's simply not possible to make
it faster than current creation (Ie it would be simpler/easier to not
preserve memoryssa in those cases).

For the common cases, assuming you tell us what variables to remove and
what variables were added in what blocks:

1. If you only removed and replaced variables (Ie literally no code motion

or insertion), it requires the renaming pass only. It could be further
optimized to not redo use chain optimization for memoryuses if you can
guarantee you only touched memoryuses. For memorydefs, if you can guarantee
(and you probably can't) that none of the aa relationships have changed,
you would not have to redo use chain opt there either. You can't guarantee
this because basicaa, et al have walkers with random hard limits in them,
and so you may get better results simply by reducing the length of a chain
of things it had to walk, even if there was in fact, no aliasing either way.

2. if the cfg has not changed, and you have added/removed variables: if you

added memoryuses, if we stay with pruned ssa, we have to recompute the IDF.
We could minimize the size we calculate by tracking some things related to
where the last uses in the program are relative to the iterated dominance
frontier. It's complicated however.
After recomputing IDF, we also have to redo renaming. For inserting
memoryuses, if we tracked the first/last version of a block (and updated
it through the creation/removal apis), we only have to rename parts of the
dom tree where the first/last version changes.

(IE if we changed first version in a block due to phi node insertion, but
last version is the same, we do not have to rename children of that block.
if we changed last version but not first version, we only have to rename
children but not the current block)
For inserting stores, it's something close to this but a bit more
complicated.

3. if the cfg has changed, without tracking a lot more under the covers

related to merge/join sets, it's not possible to make it faster than doing
it from scratch.

The TL;DR of all this is that a general update mechanism is very hard to
make faster than for scratch without giving it all kinds of specific info,
and usually figuring out that info is as hard as figuring out the update
algorithms, in my experience.

[• dberlin]

Ping!

[gberry]

Sorry if this isn't the right place, but I have a general question about what it means to preserve MemorySSA, specifically regarding the defining accesses of MemoryUse nodes. Is the idea here that we make a best effort to keep the MemoryUse defining access links optimized (i.e. never pointing to a no-alias def)? Because of the limits of basic-aa, it isn't possible to guarantee this property after any code transformation, even in the limited case here, since the alias results for completely unrelated load/stores may have been affected, right?

[Gerolf]

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

728

The code from here could be a function. I prefer a function to fit on my screen.

751

Why do you need this? When the first condition is false the while loop does not execute. But then SafeToLoadUnconditionally is not needed anyway

767

Within a block I thought the DFS numbers are
0
1
2
...
So I would expect loop up(load) < BBHoistBarrier. I'm probably wrong about the DFSs.

769

Load1 is loop invariant and the continue condition can be computed in the header.

781

That comment does not look right. Load0 and load1 just got hoisted and LookupIter->first could still be Load0.

Also, not hoisting above loads that had not been hoisted is too restrictive. With better/cheaper DA this optimization can be more aggressive. Perhaps this could be a FIXME.

[• dberlin]

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

728

Fixed

751

It's moved here because it's loop-invariant. That's also why the check matches the loop check, to avoid calculating it if the loop won't execute :)
Otherwise, we would calculate it on every loop iteration, when it doesn't change.

767

You are thinking about it backwards i think.
The dfs numbers are indeed as you list them

If the block looks like this:

0
1 - load barrier
2 - load

We can't hoist.

If it's
0
1 - load
2 - load barrier

We can hoist.

So if DFS(load) > DFS(load barrier), that means the load barrier is *above* us in the block (and we come *after* the load barrier), and we can't hoist past it.

769

Fixed

781

1. Yes, i'll fix.
2. This code was copied from mergeLoads originally. It is indeed too restrictive (there are a lot of random restrictions we could relax). I'll add a FIXME.

[• dberlin]

Update for review comments

[gberry]

A few more comments/questions below. I still have a few parts I haven't fully grasped yet, but don't let me hold up approval.

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

96

vector doesn't appear to be used?

108

Can't you just compute these instruction indices for the diamond then/else blocks lazily? The same goes for LastSinkBarrier, which wouldn't even need a map in that case.

692

Maybe check Load1->isSimple() here? Maybe factor out all these checks into a function and use it for the Load0 hash table insertion as well?

720

Isn't this still O(M^2) if e.g. all of the loads have the same clobbering access and the same types? I guess the limit above controls how big M can get though.

918

Is there a reason you're using a queue for this work-list instead of the more common SmallVector?

[• dberlin]

I'll be honest. I think we are going a bit overboard for something that is off by default and is meant to be an identical algorithm to an existing pass that we know can stand improvements.

While I am all for trying to get code to be as good as possible, i don't think we need to make it completely perfect on the first pass when it's easily subject to incremental improvements :)

IE i think there is some point that is "good enough" to start, particularly for things that we can improve before we flip them to be the default.

For all we know, changes we have to make to account for converting the other passes are going to also make us change code here anyway.

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

96

fixed

108

Yes, you can compute them once on-demand for the blocks, but it takes nearly zero time right now. It did not seem worth it for a first pass to build something to do it :)
Remember this is supposed to be an NFC conversion to start, and building the barrier maps already converts what was N^2 before to N.

692

Let us please not do that in this patch.
Doing that would change this algorithm to catch a different set of loads than the non-memoryssa version.
I would like to keep the algorithms identical in what they catch for the moment, so debugging is easy.

Eventually, when the mssa version is the default and the non-mssa version killed, we can improve it to be whatever we want.

720

If that was the case, all the loads in that block are equal (because they also passed isSameOperation, etc) and GVN should have removed them already :)

Note that neither the mssa version or the non-mssa version will handle the case of two identical loads in one side of the diamond, and one load in the other. It will only hoist/merge one of identical ones.

918

We push back and pull front because that exploration order guarantees we terminate at the earliest possible point.

If we popped off the back, it would be an order that would explore lower things before higher things.

[Gerolf]

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

751

Ok, I assumed the while loop is on the hot path usually. Then the code checks the (part of the) while condition twice.

767

Hm, it looks alright today :-). Thanks.

810

Footer -> Tail in error message

828

There is similar code in mergeLoads. That could be a utility function that returns false when there is no memory access in at least one of the blocks, e.g.. static bool hasMemoryAccess(BasicBlock *B1, BasicBlock *B2).

850

Should be if (++NStores) to be consistent with your code in mergeLoads. Or change the code for NLoads there.

854

Could be a separate function:

871

Please add a comment about what the equal stores are.

875

Since it is sinking stores would it speed up the code reversing the loop traversal (from second to first)?

900

A comment would be nice why the accesses get cleared here. Or would that fit better in sinkStores?

909

Typo: explosed

924

What is a good invariant for this loop? Like there shouldn't be more checks than memory references in the code that contains Start? A checked invariant would make me more comfortable with this code.

940

That looks very wrong. At the minimum this is a convoluted equality test.

973

I continue struggling with this function. For now I just need some help understanding the comment.
When you say "above" S0, S1 you refer to DFS numbering like Sx is above S0 in this example:

0 Sx
1
2 S0

Correct? This would be consistent with your DFS explanation in a previous explanation.

975

What is the "top of the diamond"? The head?

978

What is the "sink location block"? The Tail?

981

What is "the block"? Perhaps "uses in the tail and/or below" describes it better.

982

Would it make sense to clearly distinguish phi and memory phi (mphi) nodes?

[gberry]

I'm personally fine with committing this as is since it's is disabled, though I don't feel I have the authority to approve it.

I do have a couple more comments/questions below, mainly to help me gain confidence that my understanding of MemorySSA is correct.

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

720

Is that right (the bit about all the loads being equal)?

I'm thinking of a case like the below where all of the loads pass isSameOperation and all have the call as their clobbering access. In this case you'll end up comparing all pairs of loads (assuming none of the loads must alias).

call foo
br %c, %then, %else

then:
load i32, %p1
load i32, %p2
load i32, %p3
...
br %end

else:
load i32, %p4
load i32, %p5
load i32, %p6
...
br %end

end:

932

Do you even need the domtree info here to check for barriers in a diamond? There are references to hammocks, but as far as I can tell we never actually attempt to optimize hammocks. If that is the case, couldn't this just check for barriers in the same block as Start?

973

I think there is a bug in the case 1 scenario if there is a pre-existing phi in the bottom of the diamond that references just one of the sunk stores.
For example:

; 1 = MemDef ...
call foo
br %c %then, %else

then:
; 2 = MemDef(1)
store @A
; 3 = MemDef(2)
store @B
br %end

else:
; 4 = MemDef(1)
store @A
br %end

end:
5 = MemPhi(3, 4)

I believe in this case, after updating the end block will look like:

end:
5 = MemPhi(3, 6)
; 6 = MemDef(1)
store @A

Which seems wrong, since the phi is before it's use, but also it seems like having 6's defining access skip over the phi 5 and def 3 could cause trouble, though I'm less sure about the latter.

[• dberlin]

Let me clarify my last comment, since I hit send too early. This was a bug
i already fixed. Case 1 should also be checking that there are no stores
below. If there are, it should be doing case 2. In essence, the check for
case two should happen before the check for case 1 (since it's not possible
to have stores below without a phi node existing)

This will cause your phi node to get changed to memphi(3, 1), and do the
right thing.

I have a testcase for this.

[davide]

hrrrm, are there any real showstoppers to get this one in? Maybe Dan has no more time time to work on it but this patch already constitutes a strong basis on which we can iterate and it's also disabled by default.
@Gerolf , any opinions?

lib/Transforms/Scalar/MergedLoadStoreMotion.cpp

 //
 // 1) Generalize to regions other than diamonds
 // 2) Be more aggressive merging memory operations
 // Note that both changes require register pressure control
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Scalar.h"
+#include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/MemoryDependenceAnalysis.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
-#include "llvm/Transforms/Utils/SSAUpdater.h"
+#include "llvm/Transforms/Utils/MemorySSA.h"
 #include <vector>
+#include <unordered_set>
 using namespace llvm;
 
 #define DEBUG_TYPE "mldst-motion"
 
 //===----------------------------------------------------------------------===//
 //                         MergedLoadStoreMotion Pass
 //===----------------------------------------------------------------------===//
 
 namespace {
+
+struct StoreInstHash {
+  size_t operator()(const StoreInst *A) const {
+    return hash_combine(A->getType(), A->getPointerOperand()->getType(),
+                        A->getValueOperand()->getType(), A->isVolatile(),
+                        A->getAlignment(), A->getOrdering(),
+                        A->getSynchScope());
+  }
+};
+
+struct StoreInstEq {
+  bool operator()(const StoreInst *A, const StoreInst *B) const {
+    return (A == B || A->isSameOperationAs(B));
+  }
+};
+
+typedef std::pair<LoadInst *, const MemoryAccess *> LoadPair;
+
+struct LoadPairHash {
+  size_t operator()(const LoadPair &LP) const {
+    const LoadInst *A = LP.first;
+    return hash_combine(A->getType(), A->getPointerOperand()->getType(),
+                        A->isVolatile(), A->getAlignment(), A->getOrdering(),
+                        A->getSynchScope(), LP.second);
+  }
+};
+
+struct LoadPairEq {
+  bool operator()(const LoadPair &A, const LoadPair &B) const {
+    return (A.second == B.second &&
+            (A.first == B.first || A.first->isSameOperationAs(B.first)));
+  }
+};
+
 class MergedLoadStoreMotion : public FunctionPass {
   AliasAnalysis *AA;
   MemoryDependenceAnalysis *MD;
+  MemorySSA *MSSA;
+  MemorySSAWalker *CachingWalker;
+  DominatorTree *DT;
 
 public:
   static char ID; // Pass identification, replacement for typeid
   explicit MergedLoadStoreMotion(void)
       : FunctionPass(ID), MD(nullptr), MagicCompileTimeControl(250) {
     initializeMergedLoadStoreMotionPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnFunction(Function &F) override;
 
 private:
   // This transformation requires dominator postdominator info
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<TargetLibraryInfoWrapperPass>();
-    AU.addRequired<MemoryDependenceAnalysis>();
+    AU.addRequired<DominatorTreeWrapperPass>();
+    AU.addRequired<MemorySSALazy>();
+    AU.addPreserved<MemoryDependenceAnalysis>();
     AU.addRequired<AliasAnalysis>();
     AU.addPreserved<AliasAnalysis>();
+    AU.addPreserved<MemorySSALazy>();
   }
 
   // Helper routines
 
   ///
   /// \brief Remove instruction from parent and update memory dependence
   /// analysis.
   ///
   void removeInstruction(Instruction *Inst);
   BasicBlock *getDiamondTail(BasicBlock *BB);
   bool isDiamondHead(BasicBlock *BB);
   // Routines for hoisting loads
-  bool isLoadHoistBarrierInRange(const Instruction& Start,
-                                 const Instruction& End,
-                                 LoadInst* LI);
-  LoadInst *canHoistFromBlock(BasicBlock *BB, LoadInst *LI);
+  bool isLoadHoistBarrierInRange(const Instruction &Start,
+                                 const Instruction &End, LoadInst *LI);
   void hoistInstruction(BasicBlock *BB, Instruction *HoistCand,
                         Instruction *ElseInst);
   bool isSafeToHoist(Instruction *I) const;
   bool hoistLoad(BasicBlock *BB, LoadInst *HoistCand, LoadInst *ElseInst);
   bool mergeLoads(BasicBlock *BB);
   // Routines for sinking stores
-  StoreInst *canSinkFromBlock(BasicBlock *BB, StoreInst *SI);
   PHINode *getPHIOperand(BasicBlock *BB, StoreInst *S0, StoreInst *S1);
-  bool isStoreSinkBarrierInRange(const Instruction& Start,
-                                 const Instruction& End,
+  bool isStoreSinkBarrierInRange(const Instruction &Start,
+                                 const Instruction &End,
                                  AliasAnalysis::Location Loc);
   bool sinkStore(BasicBlock *BB, StoreInst *SinkCand, StoreInst *ElseInst);
   bool mergeStores(BasicBlock *BB);
   // The mergeLoad/Store algorithms could have Size0 * Size1 complexity,
   // where Size0 and Size1 are the #instructions on the two sides of
   // the diamond. The constant chosen here is arbitrary. Compiler Time
   // Control is enforced by the check Size0 * Size1 < MagicCompileTimeControl.
   const int MagicCompileTimeControl;
+  std::unordered_multiset<LoadPair, LoadPairHash, LoadPairEq> LoadInfo;
+  std::unordered_multiset<StoreInst *, StoreInstHash, StoreInstEq> StoreInfo;
+  SmallVector<MemoryAccess *, 8> AccessesToDelete;
 };
 
 char MergedLoadStoreMotion::ID = 0;
 }
 
 ///
 /// \brief createMergedLoadStoreMotionPass - The public interface to this file.
 ///
 FunctionPass *llvm::createMergedLoadStoreMotionPass() {
   return new MergedLoadStoreMotion();
 }
 
 INITIALIZE_PASS_BEGIN(MergedLoadStoreMotion, "mldst-motion",
                       "MergedLoadStoreMotion", false, false)
+INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(MemorySSALazy)
 INITIALIZE_PASS_DEPENDENCY(MemoryDependenceAnalysis)
 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
 INITIALIZE_AG_DEPENDENCY(AliasAnalysis)
 INITIALIZE_PASS_END(MergedLoadStoreMotion, "mldst-motion",
                     "MergedLoadStoreMotion", false, false)
 
 ///
 /// \brief Remove instruction from parent and update memory dependence analysis.

 ///
 /// \brief True when instruction is a hoist barrier for a load
 ///
 /// Whenever an instruction could possibly modify the value
 /// being loaded or protect against the load from happening
 /// it is considered a hoist barrier.
 ///
 
-bool MergedLoadStoreMotion::isLoadHoistBarrierInRange(const Instruction& Start, 
-                                                      const Instruction& End,
-                                                      LoadInst* LI) {
+bool MergedLoadStoreMotion::isLoadHoistBarrierInRange(const Instruction &Start,
+                                                      const Instruction &End,
+                                                      LoadInst *LI) {
   AliasAnalysis::Location Loc = AA->getLocation(LI);
   return AA->canInstructionRangeModRef(Start, End, Loc, AliasAnalysis::Mod);
 }
 
 ///
-/// \brief Decide if a load can be hoisted
-///
-/// When there is a load in \p BB to the same address as \p LI
-/// and it can be hoisted from \p BB, return that load.
-/// Otherwise return Null.
-///
-LoadInst *MergedLoadStoreMotion::canHoistFromBlock(BasicBlock *BB1,
-                                                   LoadInst *Load0) {
-
-  for (BasicBlock::iterator BBI = BB1->begin(), BBE = BB1->end(); BBI != BBE;
-       ++BBI) {
-    Instruction *Inst = BBI;
-
-    // Only merge and hoist loads when their result in used only in BB
-    if (!isa<LoadInst>(Inst) || Inst->isUsedOutsideOfBlock(BB1))
-      continue;
-
-    LoadInst *Load1 = dyn_cast<LoadInst>(Inst);
-    BasicBlock *BB0 = Load0->getParent();
-
-    AliasAnalysis::Location Loc0 = AA->getLocation(Load0);
-    AliasAnalysis::Location Loc1 = AA->getLocation(Load1);
-    if (AA->isMustAlias(Loc0, Loc1) && Load0->isSameOperationAs(Load1) &&
-        !isLoadHoistBarrierInRange(BB1->front(), *Load1, Load1) &&
-        !isLoadHoistBarrierInRange(BB0->front(), *Load0, Load0)) {
-      return Load1;
-    }
-  }
-  return nullptr;
-}
-
-///
 /// \brief Merge two equivalent instructions \p HoistCand and \p ElseInst into
 /// \p BB
 ///
 /// BB is the head of a diamond
 ///
 void MergedLoadStoreMotion::hoistInstruction(BasicBlock *BB,
                                              Instruction *HoistCand,
                                              Instruction *ElseInst) {
   DEBUG(dbgs() << " Hoist Instruction into BB \n"; BB->dump();
void MergedLoadStoreMotion::hoistInstruction(BasicBlock *BB,
   Instruction *HoistedInst = HoistCand->clone();
 
   // Notify AA of the new value.
   if (isa<LoadInst>(HoistCand))
     AA->copyValue(HoistCand, HoistedInst);
 
   // Hoist instruction.
   HoistedInst->insertBefore(HoistPt);
 
+  // We also hoist operands, so see if this is even a memory access we want to
+  // fix up
+  MemoryAccess *HoistCandAccess = MSSA->getMemoryAccess(HoistCand);
+  if (HoistCandAccess) {
+    MemoryAccess *NewAccess =
+        MSSA->replaceMemoryAccessWithNewAccess(HoistCandAccess, HoistedInst);
+    MSSA->replaceMemoryAccess(MSSA->getMemoryAccess(ElseInst), NewAccess);
+  }
+
   HoistCand->replaceAllUsesWith(HoistedInst);
   removeInstruction(HoistCand);
   // Replace the else block instruction.
   ElseInst->replaceAllUsesWith(HoistedInst);
   removeInstruction(ElseInst);
 }
 
 ///
bool MergedLoadStoreMotion::mergeLoads(BasicBlock *BB) {
   bool MergedLoads = false;
   assert(isDiamondHead(BB));
   BranchInst *BI = dyn_cast<BranchInst>(BB->getTerminator());
   BasicBlock *Succ0 = BI->getSuccessor(0);
   BasicBlock *Succ1 = BI->getSuccessor(1);
   // #Instructions in Succ1 for Compile Time Control
   int Size1 = Succ1->size();
   int NLoads = 0;
-  for (BasicBlock::iterator BBI = Succ0->begin(), BBE = Succ0->end();
-       BBI != BBE;) {
 
-    Instruction *I = BBI;
-    ++BBI;
+  // Skip if we don't have memory accesses in both blocks
+  auto *Succ0Accesses = MSSA->getBlockAccesses(Succ0);
+  if (!Succ0Accesses)
+    return false;
+  auto *Succ1Accesses = MSSA->getBlockAccesses(Succ1);
+  if (!Succ1Accesses)
+    return false;
 
+  for (auto &Access : *Succ0Accesses) {
+    const MemoryUse *MU = dyn_cast<MemoryUse>(&Access);
+    if (!MU)
+      continue;
+    Instruction *I = MU->getMemoryInst();
     // Only move non-simple (atomic, volatile) loads.
-    LoadInst *L0 = dyn_cast<LoadInst>(I);
-    if (!L0 || !L0->isSimple() || L0->isUsedOutsideOfBlock(Succ0))
+    LoadInst *Load0 = dyn_cast<LoadInst>(I);
+    if (!Load0 || !Load0->isSimple() || Load0->isUsedOutsideOfBlock(Succ0))
       continue;
 
+    LoadInfo.insert({Load0, CachingWalker->getClobberingMemoryAccess(Load0)});
     ++NLoads;
     if (NLoads * Size1 >= MagicCompileTimeControl)
       break;
-    if (LoadInst *L1 = canHoistFromBlock(Succ1, L0)) {
-      bool Res = hoistLoad(BB, L0, L1);
+  }
+
+  for (auto AI = Succ1Accesses->begin(), AE = Succ1Accesses->end(); AI != AE;) {
+    auto CurrIter = AI;
+    ++AI;
+    const MemoryUse *MU = dyn_cast<MemoryUse>(&*CurrIter);
+    if (!MU)
+      continue;
+    LoadInst *Load1 = dyn_cast<LoadInst>(MU->getMemoryInst());
+    if (!Load1 || Load1->isUsedOutsideOfBlock(Succ1))
+      continue;
+
+    // We know that if the load has the same clobbering access as this one,,
+    // they must not be killed until the same point
+
+    auto LookupResult = LoadInfo.equal_range(
+        {Load1, CachingWalker->getClobberingMemoryAccess(Load1)});
+    bool Res = false;
+    auto LookupIter = LookupResult.first;
+    while (!Res && LookupIter != LookupResult.second) {
+      LoadInst *Load0 = LookupIter->first;
+      auto OldIter = LookupIter;
+      ++LookupIter;
+      AliasAnalysis::Location Loc0 = AA->getLocation(Load0);
+      AliasAnalysis::Location Loc1 = AA->getLocation(Load1);
+      if (AA->isMustAlias(Loc0, Loc1))
+        Res = hoistLoad(BB, Load0, Load1);
       MergedLoads |= Res;
       // Don't attempt to hoist above loads that had not been hoisted.
-      if (!Res)
-        break;
+      if (Res)
+        LoadInfo.erase(OldIter);
     }
   }
+  LoadInfo.clear();
   return MergedLoads;
 }
 
 ///
 /// \brief True when instruction is a sink barrier for a store
 /// located in Loc
 ///
 /// Whenever an instruction could possibly read or modify the
 /// value being stored or protect against the store from
 /// happening it is considered a sink barrier.
 ///
 
-bool MergedLoadStoreMotion::isStoreSinkBarrierInRange(const Instruction& Start,
-                                                      const Instruction& End,
-                                                      AliasAnalysis::Location
-                                                      Loc) {
+bool MergedLoadStoreMotion::isStoreSinkBarrierInRange(
+    const Instruction &Start, const Instruction &End,
+    AliasAnalysis::Location Loc) {
   return AA->canInstructionRangeModRef(Start, End, Loc, AliasAnalysis::ModRef);
 }
 
 ///
-/// \brief Check if \p BB contains a store to the same address as \p SI
-///
-/// \return The store in \p  when it is safe to sink. Otherwise return Null.
-///
-StoreInst *MergedLoadStoreMotion::canSinkFromBlock(BasicBlock *BB1,
-                                                   StoreInst *Store0) {
-  DEBUG(dbgs() << "can Sink? : "; Store0->dump(); dbgs() << "\n");
-  BasicBlock *BB0 = Store0->getParent();
-  for (BasicBlock::reverse_iterator RBI = BB1->rbegin(), RBE = BB1->rend();
-       RBI != RBE; ++RBI) {
-    Instruction *Inst = &*RBI;
-
-    if (!isa<StoreInst>(Inst))
-       continue;
-
-    StoreInst *Store1 = cast<StoreInst>(Inst);
-
-    AliasAnalysis::Location Loc0 = AA->getLocation(Store0);
-    AliasAnalysis::Location Loc1 = AA->getLocation(Store1);
-    if (AA->isMustAlias(Loc0, Loc1) && Store0->isSameOperationAs(Store1) &&
-      !isStoreSinkBarrierInRange(*(std::next(BasicBlock::iterator(Store1))),
-                                 BB1->back(), Loc1) &&
-      !isStoreSinkBarrierInRange(*(std::next(BasicBlock::iterator(Store0))),
-                                 BB0->back(), Loc0)) {
-      return Store1;
-    }
-  }
-  return nullptr;
-}
-
-///
 /// \brief Create a PHI node in BB for the operands of S0 and S1
 ///
 PHINode *MergedLoadStoreMotion::getPHIOperand(BasicBlock *BB, StoreInst *S0,
                                               StoreInst *S1) {
   // Create a phi if the values mismatch.
   PHINode *NewPN = 0;
   Value *Opd1 = S0->getValueOperand();
   Value *Opd2 = S1->getValueOperand();
if (A0 && A1 && A0->isIdenticalTo(A1) && A0->hasOneUse() &&
     StoreInst *SNew = (StoreInst *)(S0->clone());
     Instruction *ANew = A0->clone();
     AA->copyValue(S0, SNew);
     SNew->insertBefore(InsertPt);
     ANew->insertBefore(SNew);
 
     assert(S0->getParent() == A0->getParent());
     assert(S1->getParent() == A1->getParent());
 
     PHINode *NewPN = getPHIOperand(BB, S0, S1);
     // New PHI operand? Use it.
     if (NewPN)
       SNew->setOperand(0, NewPN);
+    MSSA->replaceMemoryAccessWithNewAccess(MSSA->getMemoryAccess(BB), SNew);
+    MSSA->removeMemoryAccess(MSSA->getMemoryAccess(S0));
     removeInstruction(S0);
+    // ilist's reverse iterator invalidation semantics basically mean we need to
+    // wait until the loop in our caller is dead before we kill these
+    AccessesToDelete.push_back(MSSA->getMemoryAccess(S1));
     removeInstruction(S1);
     A0->replaceAllUsesWith(ANew);
     removeInstruction(A0);
     A1->replaceAllUsesWith(ANew);
     removeInstruction(A1);
     return true;
   }
   return false;
 }
 
 ///
 /// \brief True when two stores are equivalent and can sink into the footer
 ///
-/// Starting from a diamond tail block, iterate over the instructions in one
-/// predecessor block and try to match a store in the second predecessor.
+/// Starting from a diamond tail block, place all the stores in one predecessor
+/// in a hash table, and try to match them against stores in the second
+/// predecessor
 ///
 bool MergedLoadStoreMotion::mergeStores(BasicBlock *T) {
 
   bool MergedStores = false;
   assert(T && "Footer of a diamond cannot be empty");
 
   pred_iterator PI = pred_begin(T), E = pred_end(T);
   assert(PI != E);
   BasicBlock *Pred0 = *PI;
   ++PI;
   BasicBlock *Pred1 = *PI;
   ++PI;
   // tail block  of a diamond/hammock?
   if (Pred0 == Pred1)
     return false; // No.
   if (PI != E)
     return false; // No. More than 2 predecessors.
 
+  StoreInfo.reserve(Pred0->size());
+
   // #Instructions in Succ1 for Compile Time Control
   int Size1 = Pred1->size();
   int NStores = 0;
 
-  for (BasicBlock::reverse_iterator RBI = Pred0->rbegin(), RBE = Pred0->rend();
-       RBI != RBE;) {
+  // Short circuit all of this if there is no memory phi in T
+  // It's not possible for both sides of the branch to have stores unless there
+  // is a memory phi in T merging the results
+  auto *TAccesses = MSSA->getBlockAccesses(T);
+  if (!TAccesses || !isa<MemoryPhi>(&TAccesses->front()))
+    return false;
+
+  // Skip all this if we don't have any memory accesses to look at
+  auto *Pred0Accesses = MSSA->getBlockAccesses(Pred0);
+  if (!Pred0Accesses)
+    return false;
+  auto *Pred1Accesses = MSSA->getBlockAccesses(Pred1);
+  if (!Pred1Accesses)
+    return false;
 
-    Instruction *I = &*RBI;
-    ++RBI;
+  for (auto AI = Pred0Accesses->rbegin(), AE = Pred0Accesses->rend(); AI != AE;
+       ++AI) {
+    if (const MemoryDef *MD = dyn_cast<MemoryDef>(&*AI)) {
+      Instruction *I = MD->getMemoryInst();
 
       // Sink move non-simple (atomic, volatile) stores
       if (!isa<StoreInst>(I))
         continue;
-    StoreInst *S0 = (StoreInst *)I;
+      StoreInst *S0 = cast<StoreInst>(I);
       if (!S0->isSimple())
         continue;
+      StoreInfo.insert(S0);
 
       ++NStores;
       if (NStores * Size1 >= MagicCompileTimeControl)
         break;
-    if (StoreInst *S1 = canSinkFromBlock(Pred1, S0)) {
-      bool Res = sinkStore(T, S0, S1);
+    }
+  }
+
+  for (auto AI = Pred1Accesses->rbegin(), AE = Pred1Accesses->rend(); AI != AE;
+       ++AI) {
+    const MemoryDef *MD = dyn_cast<MemoryDef>(&*AI);
+    if (!MD)
+      continue;
+    Instruction *Inst = MD->getMemoryInst();
+    if (!isa<StoreInst>(Inst))
+      continue;
+
+    StoreInst *Store1 = cast<StoreInst>(Inst);
+    auto LookupResult = StoreInfo.equal_range(Store1);
+    bool Res = false;
+    auto LookupIter = LookupResult.first;
+    while (!Res && LookupIter != LookupResult.second) {
+      StoreInst *Store0 = *LookupIter;
+      auto OldIter = LookupIter;
+      ++LookupIter;
+
+      AliasAnalysis::Location Loc0 = AA->getLocation(Store0);
+      AliasAnalysis::Location Loc1 = AA->getLocation(Store1);
+      if (!AA->isMustAlias(Loc0, Loc1) ||
+          isStoreSinkBarrierInRange(*(std::next(BasicBlock::iterator(Store1))),
+                                    Pred1->back(), Loc1) ||
+          isStoreSinkBarrierInRange(*(std::next(BasicBlock::iterator(Store0))),
+                                    Pred0->back(), Loc0))
+        continue;
+      Res = sinkStore(T, Store0, Store1);
+      if (Res)
+        StoreInfo.erase(OldIter);
       MergedStores |= Res;
-      // Don't attempt to sink below stores that had to stick around
-      // But after removal of a store and some of its feeding
-      // instruction search again from the beginning since the iterator
-      // is likely stale at this point.
-      if (!Res)
-        break;
-      else {
-        RBI = Pred0->rbegin();
-        RBE = Pred0->rend();
-        DEBUG(dbgs() << "Search again\n"; Instruction *I = &*RBI; I->dump());
-      }
-    }
+  for (auto &V : AccessesToDelete) {
+    MSSA->removeMemoryAccess(V);
   }
+  AccessesToDelete.clear();
+  StoreInfo.clear();
   return MergedStores;
 }
+
 ///
 /// \brief Run the transformation for each function
 ///
 bool MergedLoadStoreMotion::runOnFunction(Function &F) {
-  MD = &getAnalysis<MemoryDependenceAnalysis>();
+  DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+  MD = getAnalysisIfAvailable<MemoryDependenceAnalysis>();
   AA = &getAnalysis<AliasAnalysis>();
+  MSSA = &getAnalysis<MemorySSALazy>().getMSSA();
+  CachingWalker = MSSA->buildMemorySSA(AA, DT);
 
   bool Changed = false;
   DEBUG(dbgs() << "Instruction Merger\n");
 
   // Merge unconditional branches, allowing PRE to catch more
   // optimization opportunities.
   for (Function::iterator FI = F.begin(), FE = F.end(); FI != FE;) {
     BasicBlock *BB = FI++;
 
     // Hoist equivalent loads and sink stores
     // outside diamonds when possible
     if (isDiamondHead(BB)) {
       Changed |= mergeLoads(BB);
       Changed |= mergeStores(getDiamondTail(BB));
     }
   }
+  StoreInfo.clear();
   return Changed;
 }