[Mlir-commits] [mlir] [mlir][Affine] Fix loadCSE replacing a load with one already scheduled for erasure (PR #217511)

[Alessandro Potenza]

https://github.com/alepot55 commented:

Built `11e915f2b75d` and ran your reproducer. Confirmed both ways:

```
$ mlir-opt --affine-scalrep loadcse.mlir          # without the patch
error: 'affine.load' op operation destroyed but still has uses
note: see current operation: %0 = "affine.load"(<<UNKNOWN SSA VALUE>>) ...
```

With the patch it completes, the three loads become two, and `mlir/test/Dialect/Affine` is 72 of 72.

**One thing worth adding to the description**, because it is the first question a reviewer will have: is `loadCSE` the only phase with this problem? It is, and for a reason that is easy to check. `affineScalarReplace` runs three walks and clears `opsToErase` between them, at the lines right after the `forwardStoreToLoad` and `findUnusedStore` loops. So a queued operation can only be picked up as a replacement inside the same phase that queued it, and of the three only `loadCSE` chooses a replacement from among other operations on the memref: the other two replace with a stored value or delete a dead store. Saying that turns "here is a fix" into "here is the fix, and there is no second one."

**Two things on the check itself.**

It is the cheapest test in the loop and it runs last, after `MemRefAccess` comparison, `domInfo.dominates`, and `hasNoInterveningEffect`, which walks. It only looks at `loadB`, so it can go at the top, next to `if (!loadB || loadB == loadA) continue;`. That skips the expensive work for candidates that are already dead, and it reads better there: an operation about to be erased is not a candidate at all, rather than a candidate that is rejected at the end.

`loadOpsToErase` is a `SmallVectorImpl<Operation *>`, so `llvm::is_contained` is a linear scan, and this check is now inside the candidate loop, which is inside the walk over every load. On a function with many loads on one memref that is a new quadratic term in a pass that had none. A `SmallPtrSet` kept alongside the vector would make it constant. Not blocking, and possibly not worth it, but it is the kind of thing that is easier to decide now than after it lands.

https://github.com/llvm/llvm-project/pull/217511